src/phase2f_layer.c

/* Phase 2f: full ScMoE decoder layer in C, validated vs layer0_target.f32 (clean f32 ref).
 * Composes the 5 proven ops: RMSNorm, matmul, MLA attention, dense MLP (SwiGLU), MoE.
 * Weights pre-dumped f32 (native shard streaming = Phase 3). Build: gcc -O2 -o phase2f_layer phase2f_layer.c -lm
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>

#define T 5
#define HID 6144
#define H 64
#define NOPE 128
#define ROPE 64
#define HD 192
#define VD 128
#define QLORA 1536
#define KVLORA 512
#define FFN 12288      /* dense mlp intermediate */
#define EFFN 2048      /* expert intermediate */
#define K 12
#define NR 512
#define EPS 1e-5f
#define SQ 2.0f
#define SKV 3.4641016151377544f
#define ASC 0.07216878364870322f
#define RSCALE 6.0f

static const char* DIR = "ref";
static float* Lf(const char* n, long c){char p[512];snprintf(p,sizeof p,"%s/%s.f32",DIR,n);FILE*f=fopen(p,"rb");if(!f){perror(p);exit(1);}float*b=malloc(c*4);if(fread(b,4,c,f)!=(size_t)c){fprintf(stderr,"short %s\n",n);exit(1);}fclose(f);return b;}
static int32_t* Li(const char* n, long c){char p[512];snprintf(p,sizeof p,"%s/%s.i32",DIR,n);FILE*f=fopen(p,"rb");if(!f){perror(p);exit(1);}int32_t*b=malloc(c*4);fread(b,4,c,f);fclose(f);return b;}
static long Lcount(const char* n){char p[512];snprintf(p,sizeof p,"%s/%s.i32",DIR,n);FILE*f=fopen(p,"rb");if(!f)return 0;fseek(f,0,SEEK_END);long b=ftell(f);fclose(f);return b/4;}

static void mv(const float*W,const float*x,float*o,int M,int Kd){for(int i=0;i<M;i++){double a=0;const float*w=W+(long)i*Kd;for(int k=0;k<Kd;k++)a+=(double)w[k]*x[k];o[i]=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];}
static inline float silu(float z){return z/(1.0f+expf(-z));}
static void il(const float*in,float*out,int d){int h=d/2;for(int i=0;i<h;i++){out[i]=in[2*i];out[h+i]=in[2*i+1];}}
static void rope(float*t,const float*co,const float*si){int h=ROPE/2;float tmp[ROPE];for(int i=0;i<ROPE;i++){float rh=(i<h)?-t[h+i]:t[i-h];tmp[i]=t[i]*co[i]+rh*si[i];}memcpy(t,tmp,ROPE*4);}

/* MLA attention for sub-block li: writes attn output [T,HID] into out */
static void mla(int li, const float* hn /*[T,HID]*/, const float* cosT, const float* sinT, float* out){
    char b[64];
    #define W(nm) ({snprintf(b,64,"L0_attn%d_%s",li,nm); b;})
    float*qa=Lf(W("q_a_proj"),(long)QLORA*HID),*qaln=Lf(W("q_a_layernorm"),QLORA),*qb=Lf(W("q_b_proj"),(long)H*HD*QLORA);
    float*kva=Lf(W("kv_a_proj_with_mqa"),(long)(KVLORA+ROPE)*HID),*kvaln=Lf(W("kv_a_layernorm"),KVLORA),*kvb=Lf(W("kv_b_proj"),(long)H*(NOPE+VD)*KVLORA),*op=Lf(W("o_proj"),(long)HID*H*VD);
    static float Q[H][T][HD],Kk[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[H*HD]; mv(qa,h,ql,QLORA,HID); rmsnorm(ql,qaln,qln,QLORA); mv(qb,qln,q,H*HD,QLORA);
        float c[KVLORA+ROPE]; mv(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 kp[H*(NOPE+VD)]; mv(kvb,kln,kp,H*(NOPE+VD),KVLORA);
        float krot[ROPE],ki[ROPE]; il(c+KVLORA,ki,ROPE); memcpy(krot,ki,ROPE*4); rope(krot,cosT+t*ROPE,sinT+t*ROPE);
        for(int hd=0;hd<H;hd++){
            const float* qh=q+hd*HD;
            for(int i=0;i<NOPE;i++) Q[hd][t][i]=qh[i]*SQ;
            float qr[ROPE],qi[ROPE]; for(int i=0;i<ROPE;i++)qr[i]=qh[NOPE+i]*SQ; il(qr,qi,ROPE);memcpy(qr,qi,ROPE*4);rope(qr,cosT+t*ROPE,sinT+t*ROPE);
            for(int i=0;i<ROPE;i++) Q[hd][t][NOPE+i]=qr[i];
            const float* kk=kp+hd*(NOPE+VD);
            for(int i=0;i<NOPE;i++) Kk[hd][t][i]=kk[i];
            for(int i=0;i<ROPE;i++) Kk[hd][t][NOPE+i]=krot[i];
            for(int i=0;i<VD;i++) V[hd][t][i]=kk[NOPE+i];
        }
    }
    static float ctx[T][H*VD];
    for(int hd=0;hd<H;hd++)for(int qi=0;qi<T;qi++){
        float sc[T],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]*Kk[hd][ki][d];sc[ki]=(float)a*ASC;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]=a/sum;}
    }
    for(int t=0;t<T;t++) mv(op,ctx[t],out+(long)t*HID,HID,H*VD);
    free(qa);free(qaln);free(qb);free(kva);free(kvaln);free(kvb);free(op);
}
static void dense(int li,const float*hn,float*out){
    char b[64];float*g,*u,*d; snprintf(b,64,"L0_mlp%d_gate",li);g=Lf(b,(long)FFN*HID);snprintf(b,64,"L0_mlp%d_up",li);u=Lf(b,(long)FFN*HID);snprintf(b,64,"L0_mlp%d_down",li);d=Lf(b,(long)HID*FFN);
    for(int t=0;t<T;t++){float gp[FFN],up[FFN],ac[FFN];mv(g,hn+(long)t*HID,gp,FFN,HID);mv(u,hn+(long)t*HID,up,FFN,HID);for(int i=0;i<FFN;i++)ac[i]=silu(gp[i])*up[i];mv(d,ac,out+(long)t*HID,HID,FFN);}
    free(g);free(u);free(d);
}
static void moe(const float*hn,float*out){
    float*cw=Lf("L0_router_w",(long)768*HID),*bias=Lf("L0_router_b",768);
    memset(out,0,(long)T*HID*4);
    for(int t=0;t<T;t++){
        float lg[768]; mv(cw,hn+(long)t*HID,lg,768,HID);
        float mx=-1e30f;for(int i=0;i<768;i++)if(lg[i]>mx)mx=lg[i];
        double sum=0;float sf[768];for(int i=0;i<768;i++){sf[i]=expf(lg[i]-mx);sum+=sf[i];} for(int i=0;i<768;i++)sf[i]/=sum;
        /* topk-12 by (softmax + bias) */
        int idx[K];float chosen[768];for(int i=0;i<768;i++)chosen[i]=sf[i]+bias[i];
        for(int j=0;j<K;j++){int bi=-1;float bv=-1e30f;for(int i=0;i<768;i++){int taken=0;for(int q=0;q<j;q++)if(idx[q]==i)taken=1;if(!taken&&chosen[i]>bv){bv=chosen[i];bi=i;}}idx[j]=bi;}
        for(int j=0;j<K;j++){int e=idx[j];float w=sf[e]*RSCALE;
            if(e>=NR){for(int i=0;i<HID;i++)out[(long)t*HID+i]+=hn[(long)t*HID+i]*w;continue;}
            char b[64];float*g,*u,*d;snprintf(b,64,"L0exp_%d_gate",e);g=Lf(b,(long)EFFN*HID);snprintf(b,64,"L0exp_%d_up",e);u=Lf(b,(long)EFFN*HID);snprintf(b,64,"L0exp_%d_down",e);d=Lf(b,(long)HID*EFFN);
            float gp[EFFN],up[EFFN],ac[EFFN],dp[HID];mv(g,hn+(long)t*HID,gp,EFFN,HID);mv(u,hn+(long)t*HID,up,EFFN,HID);for(int i=0;i<EFFN;i++)ac[i]=silu(gp[i])*up[i];mv(d,ac,dp,HID,EFFN);
            for(int i=0;i<HID;i++)out[(long)t*HID+i]+=dp[i]*w; free(g);free(u);free(d);
        }
    }
    free(cw);free(bias);
}

int main(void){
    float*x=Lf("embed_out",(long)T*HID);
    float*cosT=Lf("mla_cos",(long)T*ROPE),*sinT=Lf("mla_sin",(long)T*ROPE);
    float*ref=Lf("layer0_target",(long)T*HID);
    float *tmp=malloc((long)T*HID*4),*xn=malloc((long)T*HID*4),*shortcut=malloc((long)T*HID*4);
    char b[64];
    for(int i=0;i<2;i++){
        /* attn: r + attn(rms(x, inln)) */
        snprintf(b,64,"L0_inln%d",i); float* inln=Lf(b,HID);
        for(int t=0;t<T;t++) rmsnorm(x+(long)t*HID,inln,xn+(long)t*HID,HID);
        mla(i,xn,cosT,sinT,tmp);
        for(long j=0;j<(long)T*HID;j++) x[j]+=tmp[j];
        free(inln);
        /* post-norm */
        snprintf(b,64,"L0_paln%d",i); float* paln=Lf(b,HID);
        for(int t=0;t<T;t++) rmsnorm(x+(long)t*HID,paln,xn+(long)t*HID,HID);
        free(paln);
        if(i==0) moe(xn,shortcut);
        dense(i,xn,tmp);
        for(long j=0;j<(long)T*HID;j++) x[j]+=tmp[j];
        if(i==1) for(long j=0;j<(long)T*HID;j++) x[j]+=shortcut[j];
    }
    double dmax=0;for(long j=0;j<(long)T*HID;j++){double dd=fabs((double)x[j]-ref[j]);if(dd>dmax)dmax=dd;}
    printf("layer0[0,:6]: ");for(int i=0;i<6;i++)printf("%.4f ",x[i]);printf("\n");
    printf("FULL C LAYER max abs diff vs clean Python: %.3e\n",dmax);
    printf("VERDICT: %s\n",dmax<0.1?"PASS — full ScMoE layer works in C":"FAIL");
    return dmax<0.1?0:1;
}