src/phase2d_mla.c

/* Phase 2d: MLA attention in C, validated vs Python (attn_out.f32).
 * Weights pre-dumped as f32 (weight-loading proven in Phase 1). Isolates the attention MATH.
 * Build: gcc -O2 -o phase2d_mla phase2d_mla.c -lm
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#define T 5
#define HID 6144
#define H 64
#define NOPE 128
#define ROPE 64
#define HD 192          /* qk_head_dim */
#define VD 128
#define QLORA 1536
#define KVLORA 512
#define QB_OUT 12288    /* H*HD */
#define KVB_OUT 16384   /* H*(NOPE+VD) */
#define O_IN 8192       /* H*VD */
#define SQ 2.0f
#define SKV 3.4641016151377544f
#define SCALE 0.07216878364870322f
#define EPS 1e-5f

static float* L(const char* n, long cnt) {
    char p[512]; snprintf(p, sizeof p, "ref/%s.f32", n);
    FILE* f = fopen(p, "rb"); if (!f) { perror(p); exit(1); }
    float* b = malloc(cnt * 4);
    if (fread(b, 4, cnt, f) != (size_t)cnt) { fprintf(stderr, "short %s\n", n); exit(1); }
    fclose(f); return b;
}
/* out[M] = W[M,K] @ x[K] */
static void matvec(const float* W, const float* x, float* out, int M, int K) {
    for (int i = 0; i < M; i++) { double a = 0; const float* w = W + (long)i * K;
        for (int k = 0; k < K; k++) a += (double)w[k] * x[k]; out[i] = (float)a; }
}
static void rmsnorm(const float* x, const float* w, float* o, int d) {
    double ss = 0; for (int i = 0; i < d; i++) ss += (double)x[i] * x[i];
    float inv = 1.0f / sqrtf((float)(ss / d) + EPS);
    for (int i = 0; i < d; i++) o[i] = x[i] * inv * w[i];
}
/* interleave (use_mla): [d] viewed as [d/2,2] transposed to [2,d/2] flattened */
static void interleave(const float* in, float* out, int d) {
    int half = d / 2;
    for (int i = 0; i < half; i++) { out[i] = in[2*i]; out[half + i] = in[2*i + 1]; }
}
/* rope in place on [ROPE] with cos/sin[ROPE]; rotate_half = [-x[half:], x[:half]] */
static void rope(float* t, const float* cos, const float* sin) {
    int half = ROPE / 2; float tmp[ROPE];
    for (int i = 0; i < ROPE; i++) {
        float rh = (i < half) ? -t[half + i] : t[i - half];
        tmp[i] = t[i] * cos[i] + rh * sin[i];
    }
    memcpy(t, tmp, ROPE * 4);
}

int main(void) {
    float* qa=L("attnw_q_a_proj",(long)QLORA*HID);   float* qaln=L("attnw_q_a_layernorm",QLORA);
    float* qb=L("attnw_q_b_proj",(long)QB_OUT*QLORA);
    float* kva=L("attnw_kv_a_proj_with_mqa",(long)(KVLORA+ROPE)*HID); float* kvaln=L("attnw_kv_a_layernorm",KVLORA);
    float* kvb=L("attnw_kv_b_proj",(long)KVB_OUT*KVLORA); float* op=L("attnw_o_proj",(long)HID*O_IN);
    float* cosT=L("mla_cos",(long)T*ROPE); float* sinT=L("mla_sin",(long)T*ROPE);
    float* hn=L("attn_in",(long)T*HID);   float* ref=L("attn_out",(long)T*HID);

    /* per-token compute Q,K,V, then attention */
    static float Q[H][T][HD], K[H][T][HD], V[H][T][VD];
    for (int t = 0; t < T; t++) {
        const float* h = hn + (long)t * HID;
        float ql[QLORA], qln[QLORA], q[QB_OUT];
        matvec(qa, h, ql, QLORA, HID); rmsnorm(ql, qaln, qln, QLORA); matvec(qb, qln, q, QB_OUT, QLORA);
        float c[KVLORA + ROPE]; matvec(kva, h, c, KVLORA + ROPE, HID);
        float kln[KVLORA]; rmsnorm(c, kvaln, kln, KVLORA);
        for (int i = 0; i < KVLORA; i++) kln[i] *= SKV;
        float kpass[KVB_OUT]; matvec(kvb, kln, kpass, KVB_OUT, KVLORA);
        /* k_rot shared: interleave + rope once per token */
        float krot[ROPE], krot_i[ROPE];
        interleave(c + KVLORA, krot_i, ROPE); memcpy(krot, krot_i, ROPE*4);
        rope(krot, cosT + t*ROPE, sinT + t*ROPE);
        for (int hd = 0; hd < H; hd++) {
            /* q for this head: q[hd*HD ..], split nope/rope, scale, interleave+rope the rope part */
            const float* qh = q + hd * HD;
            for (int i = 0; i < NOPE; i++) Q[hd][t][i] = qh[i] * SQ;
            float qrot[ROPE], qrot_i[ROPE];
            for (int i = 0; i < ROPE; i++) qrot[i] = qh[NOPE + i] * SQ;
            interleave(qrot, qrot_i, ROPE); memcpy(qrot, qrot_i, ROPE*4);
            rope(qrot, cosT + t*ROPE, sinT + t*ROPE);
            for (int i = 0; i < ROPE; i++) Q[hd][t][NOPE + i] = qrot[i];
            /* k = [k_nope | krot(shared)] , v */
            const float* kp = kpass + hd * (NOPE + VD);
            for (int i = 0; i < NOPE; i++) K[hd][t][i] = kp[i];
            for (int i = 0; i < ROPE; i++) K[hd][t][NOPE + i] = krot[i];
            for (int i = 0; i < VD; i++) V[hd][t][i] = kp[NOPE + i];
        }
    }
    /* attention per head, causal */
    static float ctx[T][O_IN];   /* [T, H*VD] */
    for (int hd = 0; hd < H; hd++)
        for (int qi = 0; qi < T; qi++) {
            float sc[T]; float mx = -1e30f;
            for (int ki = 0; ki <= qi; ki++) { double a = 0;
                for (int d = 0; d < HD; d++) a += (double)Q[hd][qi][d] * K[hd][ki][d];
                sc[ki] = (float)a * SCALE; if (sc[ki] > mx) mx = sc[ki]; }
            double sum = 0; for (int ki = 0; ki <= qi; ki++) { sc[ki] = expf(sc[ki] - mx); sum += sc[ki]; }
            for (int d = 0; d < VD; d++) { double a = 0;
                for (int ki = 0; ki <= qi; ki++) a += sc[ki] * V[hd][ki][d];
                ctx[qi][hd * VD + d] = (float)(a / sum); }
        }
    /* o_proj */
    double dmax = 0;
    for (int t = 0; t < T; t++) { float o[HID]; matvec(op, ctx[t], o, HID, O_IN);
        for (int i = 0; i < HID; i++) { double d = fabs((double)o[i] - ref[(long)t*HID + i]); if (d > dmax) dmax = d; }
        if (t == 0) { printf("out[0..8]: "); for (int i = 0; i < 8; i++) printf("%.4f ", o[i]); printf("\n"); }
    }
    printf("MLA attention max abs diff vs Python: %.3e\n", dmax);
    printf("VERDICT: %s\n", dmax < 5e-2 ? "PASS — MLA attention matches Python" : "FAIL");
    return dmax < 5e-2 ? 0 : 1;
}