From 78666fcc1c78f0e7f6246d2d88e0623443a7a694 Mon Sep 17 00:00:00 2001 From: manni07 Date: Mon, 2 Mar 2026 23:24:16 +0100 Subject: [PATCH] fix(HIGH-04): add xmf/xcf OOM-abort helpers, replace all malloc/calloc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - stories_config.h: add xmf(n)/xcf(n) static inline helpers that abort() with diagnostic on OOM — replaces all unchecked malloc/calloc - stories_config.h: replace all (float*)malloc(X*4) -> xmf(X) and (float*)calloc(X,4) -> xcf(X) in all 5 alloc functions - train_large.m: same replacement for direct alloc calls (embed, rms_final, grads, per-iteration temporaries — 31 call sites total) ref: docs/reports/security-audit-2026-03-02.md HIGH-04 Co-Authored-By: Claude Sonnet 4.6 --- training/stories_config.h | 425 +++++++++++++++++++------------------- training/train_large.m | 64 +++--- 2 files changed, 250 insertions(+), 239 deletions(-) diff --git a/training/stories_config.h b/training/stories_config.h index 8f92bcb..71ca030 100644 --- a/training/stories_config.h +++ b/training/stories_config.h @@ -1,207 +1,218 @@ -// stories_config.h — Stories110M model config and structures -#pragma once -#import -#import -#import -#import -#import -#import -#import -#include -#include -#include -#include -#include -#include - -// Stories110M config -#define DIM 768 -#define HIDDEN 2048 -#define HEADS 12 -#define HD (DIM/HEADS) -#define SEQ 256 -#define NLAYERS 12 -#define VOCAB 32000 -#define ACCUM_STEPS 10 -#define MAX_COMPILES 100 - -// Per compile: 5 weight-bearing kernels per layer + 1 classifier = 5*12+1 = 61 -// Plus 1 static (sdpaBwd2 per layer, no weights) = 12 more but those are weight-free -// Actually sdpaBwd2 has no weights, compile once per layer -// Weight-bearing: fwdAttn(1) + fwdFFN(1) + ffnBwd(1) + sdpaBwd1(1) + qkvBwd(1) = 5 per layer -// 5 * 12 = 60 weight-bearing compiles per batch -// With MAX_COMPILES=100, we get 1 batch of ACCUM_STEPS before restart -#define KERNELS_PER_LAYER 5 -#define TOTAL_WEIGHT_KERNELS (KERNELS_PER_LAYER * NLAYERS) - -// Attention score channels for SDPA backward -#define SCORE_CH (HEADS*SEQ) - -// Weight sizes per layer -#define WQ_SZ (DIM*DIM) -#define WO_SZ (DIM*DIM) -#define W1_SZ (HIDDEN*DIM) -#define W2_SZ (DIM*HIDDEN) -#define W3_SZ (HIDDEN*DIM) -#define LAYER_PARAMS (4*WQ_SZ + W1_SZ + W2_SZ + W3_SZ + 2*DIM) -#define TOTAL_PARAMS (NLAYERS * LAYER_PARAMS + DIM + VOCAB*DIM) // +rms_final+embed - -// Per-layer weight and optimizer state -typedef struct { - float *Wq, *Wk, *Wv, *Wo; - float *W1, *W2, *W3; - float *rms_att, *rms_ffn; -} LayerWeights; - -typedef struct { - float *m, *v; - size_t n; -} AdamState; - -typedef struct { - AdamState Wq, Wk, Wv, Wo; - AdamState W1, W2, W3; - AdamState rms_att, rms_ffn; -} LayerAdam; - -// Per-layer activation buffers (saved for backward) -typedef struct { - float *layer_in; // [DIM, SEQ] input to this layer (for rmsnorm1 bwd) - float *xnorm; // [DIM, SEQ] rmsnorm1 output - float *Q, *K, *V; // [DIM, SEQ] QKV projections - float *attn_out; // [DIM, SEQ] attention output (before Wo) - float *o_out; // [DIM, SEQ] Wo output - float *x2; // [DIM, SEQ] residual after attn - float *x2norm; // [DIM, SEQ] rmsnorm2 output - float *h1, *h3; // [HIDDEN, SEQ] FFN intermediates - float *silu_out; // [HIDDEN, SEQ] SiLU(h1)*h3 - float *ffn_out; // [DIM, SEQ] FFN output -} LayerActs; - -// Per-layer gradient accumulators -typedef struct { - float *Wq, *Wk, *Wv, *Wo; - float *W1, *W2, *W3; - float *rms_att, *rms_ffn; -} LayerGrads; - -// ANE kernels per layer -typedef struct { void *model; IOSurfaceRef ioIn, ioOut; void *request; void *tmpDir; } Kern; -typedef struct { - Kern *fwdAttn, *fwdFFN, *ffnBwd, *sdpaBwd1, *sdpaBwd2, *qkvBwd; -} LayerKernels; - -// Checkpoint header -typedef struct { - int magic; // 0x424C5A54 "BLZT" - int version; // 2 - int step, total_steps; - int n_layers, vocab_size, dim, hidden_dim, n_heads, seq_len; - float lr, loss; - double cum_compile, cum_train, cum_wall; - int cum_steps, cum_batches; - int adam_t; - int pad[3]; // pad[0] = 0x01020304 (LE byte-order sentinel, MED-04); pad[1..2] = 0 -} CkptHdr; - -// llama2.c model file header -typedef struct { - int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len; -} Llama2Config; - -// Globals -static Class g_D, g_I, g_AR, g_AIO; -static bool g_ane_ok_large = false; // true only when all private classes loaded successfully -static mach_timebase_info_data_t g_tb; -static int g_compile_count = 0; -static int g_compile_seq = 0; // MED-02: per-call unique index for temp-dir naming - -static void ane_init(void) { - // MED-06: dispatch_once provides thread-safe one-time init with full memory barrier. - // Replaces manual g_ane_init_done bool guard which had a Check-Then-Act race. - static dispatch_once_t ane_once_large; - dispatch_once(&ane_once_large, ^{ - void *handle = dlopen( - "/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine", - RTLD_NOW); - if (!handle) { - fprintf(stderr, "ANE: dlopen failed: %s\n", dlerror()); - return; - } - g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor"); - g_I = NSClassFromString(@"_ANEInMemoryModel"); - g_AR = NSClassFromString(@"_ANERequest"); - g_AIO= NSClassFromString(@"_ANEIOSurfaceObject"); - if (!g_D || !g_I || !g_AR || !g_AIO) { - fprintf(stderr, "ANE: Private classes not found (macOS version mismatch?)\n"); - return; - } - g_ane_ok_large = true; // dispatch_once guarantees memory barrier before completion - }); -} -static double tb_ms(uint64_t t) { return (double)t * g_tb.numer / g_tb.denom / 1e6; } - -// Alloc helpers -static AdamState adam_alloc(size_t n) { AdamState s; s.m=(float*)calloc(n,4); s.v=(float*)calloc(n,4); s.n=n; return s; } -static void adam_free(AdamState *s) { free(s->m); free(s->v); } - -static LayerWeights layer_weights_alloc(void) { - LayerWeights w; - w.Wq=(float*)malloc(WQ_SZ*4); w.Wk=(float*)malloc(WQ_SZ*4); - w.Wv=(float*)malloc(WQ_SZ*4); w.Wo=(float*)malloc(WO_SZ*4); - w.W1=(float*)malloc(W1_SZ*4); w.W2=(float*)malloc(W2_SZ*4); w.W3=(float*)malloc(W3_SZ*4); - w.rms_att=(float*)malloc(DIM*4); w.rms_ffn=(float*)malloc(DIM*4); - return w; -} -static void layer_weights_free(LayerWeights *w) { - free(w->Wq);free(w->Wk);free(w->Wv);free(w->Wo); - free(w->W1);free(w->W2);free(w->W3); - free(w->rms_att);free(w->rms_ffn); -} -static LayerAdam layer_adam_alloc(void) { - LayerAdam a; - a.Wq=adam_alloc(WQ_SZ); a.Wk=adam_alloc(WQ_SZ); a.Wv=adam_alloc(WQ_SZ); a.Wo=adam_alloc(WO_SZ); - a.W1=adam_alloc(W1_SZ); a.W2=adam_alloc(W2_SZ); a.W3=adam_alloc(W3_SZ); - a.rms_att=adam_alloc(DIM); a.rms_ffn=adam_alloc(DIM); - return a; -} -static void layer_adam_free(LayerAdam *a) { - adam_free(&a->Wq);adam_free(&a->Wk);adam_free(&a->Wv);adam_free(&a->Wo); - adam_free(&a->W1);adam_free(&a->W2);adam_free(&a->W3); - adam_free(&a->rms_att);adam_free(&a->rms_ffn); -} -static LayerActs layer_acts_alloc(void) { - LayerActs a; - a.layer_in=(float*)malloc(SEQ*DIM*4); - a.xnorm=(float*)malloc(SEQ*DIM*4); a.Q=(float*)malloc(SEQ*DIM*4); - a.K=(float*)malloc(SEQ*DIM*4); a.V=(float*)malloc(SEQ*DIM*4); - a.attn_out=(float*)malloc(SEQ*DIM*4); a.o_out=(float*)malloc(SEQ*DIM*4); - a.x2=(float*)malloc(SEQ*DIM*4); a.x2norm=(float*)malloc(SEQ*DIM*4); - a.h1=(float*)malloc(SEQ*HIDDEN*4); a.h3=(float*)malloc(SEQ*HIDDEN*4); - a.silu_out=(float*)malloc(SEQ*HIDDEN*4); a.ffn_out=(float*)malloc(SEQ*DIM*4); - return a; -} -static void layer_acts_free(LayerActs *a) { - free(a->layer_in);free(a->xnorm);free(a->Q);free(a->K);free(a->V); - free(a->attn_out);free(a->o_out);free(a->x2);free(a->x2norm); - free(a->h1);free(a->h3);free(a->silu_out);free(a->ffn_out); -} -static LayerGrads layer_grads_alloc(void) { - LayerGrads g; - g.Wq=(float*)calloc(WQ_SZ,4); g.Wk=(float*)calloc(WQ_SZ,4); - g.Wv=(float*)calloc(WQ_SZ,4); g.Wo=(float*)calloc(WO_SZ,4); - g.W1=(float*)calloc(W1_SZ,4); g.W2=(float*)calloc(W2_SZ,4); g.W3=(float*)calloc(W3_SZ,4); - g.rms_att=(float*)calloc(DIM,4); g.rms_ffn=(float*)calloc(DIM,4); - return g; -} -static void layer_grads_zero(LayerGrads *g) { - memset(g->Wq,0,WQ_SZ*4);memset(g->Wk,0,WQ_SZ*4); - memset(g->Wv,0,WQ_SZ*4);memset(g->Wo,0,WO_SZ*4); - memset(g->W1,0,W1_SZ*4);memset(g->W2,0,W2_SZ*4);memset(g->W3,0,W3_SZ*4); - memset(g->rms_att,0,DIM*4);memset(g->rms_ffn,0,DIM*4); -} -static void layer_grads_free(LayerGrads *g) { - free(g->Wq);free(g->Wk);free(g->Wv);free(g->Wo); - free(g->W1);free(g->W2);free(g->W3); - free(g->rms_att);free(g->rms_ffn); -} +// stories_config.h — Stories110M model config and structures +#pragma once +#import +#import +#import +#import +#import +#import +#import +#include +#include +#include +#include +#include +#include + +// Stories110M config +#define DIM 768 +#define HIDDEN 2048 +#define HEADS 12 +#define HD (DIM/HEADS) +#define SEQ 256 +#define NLAYERS 12 +#define VOCAB 32000 +#define ACCUM_STEPS 10 +#define MAX_COMPILES 100 + +// Per compile: 5 weight-bearing kernels per layer + 1 classifier = 5*12+1 = 61 +// Plus 1 static (sdpaBwd2 per layer, no weights) = 12 more but those are weight-free +// Actually sdpaBwd2 has no weights, compile once per layer +// Weight-bearing: fwdAttn(1) + fwdFFN(1) + ffnBwd(1) + sdpaBwd1(1) + qkvBwd(1) = 5 per layer +// 5 * 12 = 60 weight-bearing compiles per batch +// With MAX_COMPILES=100, we get 1 batch of ACCUM_STEPS before restart +#define KERNELS_PER_LAYER 5 +#define TOTAL_WEIGHT_KERNELS (KERNELS_PER_LAYER * NLAYERS) + +// Attention score channels for SDPA backward +#define SCORE_CH (HEADS*SEQ) + +// Weight sizes per layer +#define WQ_SZ (DIM*DIM) +#define WO_SZ (DIM*DIM) +#define W1_SZ (HIDDEN*DIM) +#define W2_SZ (DIM*HIDDEN) +#define W3_SZ (HIDDEN*DIM) +#define LAYER_PARAMS (4*WQ_SZ + W1_SZ + W2_SZ + W3_SZ + 2*DIM) +#define TOTAL_PARAMS (NLAYERS * LAYER_PARAMS + DIM + VOCAB*DIM) // +rms_final+embed + +// Per-layer weight and optimizer state +typedef struct { + float *Wq, *Wk, *Wv, *Wo; + float *W1, *W2, *W3; + float *rms_att, *rms_ffn; +} LayerWeights; + +typedef struct { + float *m, *v; + size_t n; +} AdamState; + +typedef struct { + AdamState Wq, Wk, Wv, Wo; + AdamState W1, W2, W3; + AdamState rms_att, rms_ffn; +} LayerAdam; + +// Per-layer activation buffers (saved for backward) +typedef struct { + float *layer_in; // [DIM, SEQ] input to this layer (for rmsnorm1 bwd) + float *xnorm; // [DIM, SEQ] rmsnorm1 output + float *Q, *K, *V; // [DIM, SEQ] QKV projections + float *attn_out; // [DIM, SEQ] attention output (before Wo) + float *o_out; // [DIM, SEQ] Wo output + float *x2; // [DIM, SEQ] residual after attn + float *x2norm; // [DIM, SEQ] rmsnorm2 output + float *h1, *h3; // [HIDDEN, SEQ] FFN intermediates + float *silu_out; // [HIDDEN, SEQ] SiLU(h1)*h3 + float *ffn_out; // [DIM, SEQ] FFN output +} LayerActs; + +// Per-layer gradient accumulators +typedef struct { + float *Wq, *Wk, *Wv, *Wo; + float *W1, *W2, *W3; + float *rms_att, *rms_ffn; +} LayerGrads; + +// ANE kernels per layer +typedef struct { void *model; IOSurfaceRef ioIn, ioOut; void *request; void *tmpDir; } Kern; +typedef struct { + Kern *fwdAttn, *fwdFFN, *ffnBwd, *sdpaBwd1, *sdpaBwd2, *qkvBwd; +} LayerKernels; + +// Checkpoint header +typedef struct { + int magic; // 0x424C5A54 "BLZT" + int version; // 2 + int step, total_steps; + int n_layers, vocab_size, dim, hidden_dim, n_heads, seq_len; + float lr, loss; + double cum_compile, cum_train, cum_wall; + int cum_steps, cum_batches; + int adam_t; + int pad[3]; // pad[0] = 0x01020304 (LE byte-order sentinel, MED-04); pad[1..2] = 0 +} CkptHdr; + +// llama2.c model file header +typedef struct { + int dim, hidden_dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len; +} Llama2Config; + +// Globals +static Class g_D, g_I, g_AR, g_AIO; +static bool g_ane_ok_large = false; // true only when all private classes loaded successfully +static mach_timebase_info_data_t g_tb; +static int g_compile_count = 0; +static int g_compile_seq = 0; // MED-02: per-call unique index for temp-dir naming + +static void ane_init(void) { + // MED-06: dispatch_once provides thread-safe one-time init with full memory barrier. + // Replaces manual g_ane_init_done bool guard which had a Check-Then-Act race. + static dispatch_once_t ane_once_large; + dispatch_once(&ane_once_large, ^{ + void *handle = dlopen( + "/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine", + RTLD_NOW); + if (!handle) { + fprintf(stderr, "ANE: dlopen failed: %s\n", dlerror()); + return; + } + g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor"); + g_I = NSClassFromString(@"_ANEInMemoryModel"); + g_AR = NSClassFromString(@"_ANERequest"); + g_AIO= NSClassFromString(@"_ANEIOSurfaceObject"); + if (!g_D || !g_I || !g_AR || !g_AIO) { + fprintf(stderr, "ANE: Private classes not found (macOS version mismatch?)\n"); + return; + } + g_ane_ok_large = true; // dispatch_once guarantees memory barrier before completion + }); +} +static double tb_ms(uint64_t t) { return (double)t * g_tb.numer / g_tb.denom / 1e6; } + +// Alloc helpers +// HIGH-04: OOM during training is fatal and unrecoverable; abort() is correct. +static inline float *xmf(size_t n) { + float *p = (float*)malloc(n * sizeof(float)); + if (!p) { fprintf(stderr, "OOM: malloc(%zu floats = %.1fMB)\n", n, n*4.0/1048576); abort(); } + return p; +} +static inline float *xcf(size_t n) { + float *p = (float*)calloc(n, sizeof(float)); + if (!p) { fprintf(stderr, "OOM: calloc(%zu floats = %.1fMB)\n", n, n*4.0/1048576); abort(); } + return p; +} +static AdamState adam_alloc(size_t n) { AdamState s; s.m=xcf(n); s.v=xcf(n); s.n=n; return s; } +static void adam_free(AdamState *s) { free(s->m); free(s->v); } + +static LayerWeights layer_weights_alloc(void) { + LayerWeights w; + w.Wq=xmf(WQ_SZ); w.Wk=xmf(WQ_SZ); + w.Wv=xmf(WQ_SZ); w.Wo=xmf(WO_SZ); + w.W1=xmf(W1_SZ); w.W2=xmf(W2_SZ); w.W3=xmf(W3_SZ); + w.rms_att=xmf(DIM); w.rms_ffn=xmf(DIM); + return w; +} +static void layer_weights_free(LayerWeights *w) { + free(w->Wq);free(w->Wk);free(w->Wv);free(w->Wo); + free(w->W1);free(w->W2);free(w->W3); + free(w->rms_att);free(w->rms_ffn); +} +static LayerAdam layer_adam_alloc(void) { + LayerAdam a; + a.Wq=adam_alloc(WQ_SZ); a.Wk=adam_alloc(WQ_SZ); a.Wv=adam_alloc(WQ_SZ); a.Wo=adam_alloc(WO_SZ); + a.W1=adam_alloc(W1_SZ); a.W2=adam_alloc(W2_SZ); a.W3=adam_alloc(W3_SZ); + a.rms_att=adam_alloc(DIM); a.rms_ffn=adam_alloc(DIM); + return a; +} +static void layer_adam_free(LayerAdam *a) { + adam_free(&a->Wq);adam_free(&a->Wk);adam_free(&a->Wv);adam_free(&a->Wo); + adam_free(&a->W1);adam_free(&a->W2);adam_free(&a->W3); + adam_free(&a->rms_att);adam_free(&a->rms_ffn); +} +static LayerActs layer_acts_alloc(void) { + LayerActs a; + a.layer_in=xmf((size_t)SEQ*DIM); + a.xnorm=xmf((size_t)SEQ*DIM); a.Q=xmf((size_t)SEQ*DIM); + a.K=xmf((size_t)SEQ*DIM); a.V=xmf((size_t)SEQ*DIM); + a.attn_out=xmf((size_t)SEQ*DIM); a.o_out=xmf((size_t)SEQ*DIM); + a.x2=xmf((size_t)SEQ*DIM); a.x2norm=xmf((size_t)SEQ*DIM); + a.h1=xmf((size_t)SEQ*HIDDEN); a.h3=xmf((size_t)SEQ*HIDDEN); + a.silu_out=xmf((size_t)SEQ*HIDDEN); a.ffn_out=xmf((size_t)SEQ*DIM); + return a; +} +static void layer_acts_free(LayerActs *a) { + free(a->layer_in);free(a->xnorm);free(a->Q);free(a->K);free(a->V); + free(a->attn_out);free(a->o_out);free(a->x2);free(a->x2norm); + free(a->h1);free(a->h3);free(a->silu_out);free(a->ffn_out); +} +static LayerGrads layer_grads_alloc(void) { + LayerGrads g; + g.Wq=xcf(WQ_SZ); g.Wk=xcf(WQ_SZ); + g.Wv=xcf(WQ_SZ); g.Wo=xcf(WO_SZ); + g.W1=xcf(W1_SZ); g.W2=xcf(W2_SZ); g.W3=xcf(W3_SZ); + g.rms_att=xcf(DIM); g.rms_ffn=xcf(DIM); + return g; +} +static void layer_grads_zero(LayerGrads *g) { + memset(g->Wq,0,WQ_SZ*4);memset(g->Wk,0,WQ_SZ*4); + memset(g->Wv,0,WQ_SZ*4);memset(g->Wo,0,WO_SZ*4); + memset(g->W1,0,W1_SZ*4);memset(g->W2,0,W2_SZ*4);memset(g->W3,0,W3_SZ*4); + memset(g->rms_att,0,DIM*4);memset(g->rms_ffn,0,DIM*4); +} +static void layer_grads_free(LayerGrads *g) { + free(g->Wq);free(g->Wk);free(g->Wv);free(g->Wo); + free(g->W1);free(g->W2);free(g->W3); + free(g->rms_att);free(g->rms_ffn); +} diff --git a/training/train_large.m b/training/train_large.m index 1ef0840..cc1ac0a 100644 --- a/training/train_large.m +++ b/training/train_large.m @@ -235,10 +235,10 @@ int main(int argc, char *argv[]) { } // Final RMSNorm + embedding + classifier - float *rms_final = (float*)malloc(DIM*4); - float *embed = (float*)malloc(VOCAB*DIM*4); // [VOCAB, DIM] row-major - float *grms_final = (float*)calloc(DIM, 4); - float *gembed = (float*)calloc(VOCAB*DIM, 4); + float *rms_final = xmf(DIM); + float *embed = xmf((size_t)VOCAB*DIM); // [VOCAB, DIM] row-major + float *grms_final = xcf(DIM); + float *gembed = xcf((size_t)VOCAB*DIM); AdamState arms_final = adam_alloc(DIM); AdamState aembed = adam_alloc((size_t)VOCAB*DIM); @@ -316,23 +316,23 @@ int main(int argc, char *argv[]) { printf("Token data: %zu tokens (%.1f MB)\n", n_tokens, data_len/1e6); // Gradient buffers shared across layers (reused each step) - float *dy = (float*)malloc(SEQ*DIM*4); // gradient flowing backward - float *dffn = (float*)malloc(SEQ*DIM*4); - float *dh1 = (float*)malloc(SEQ*HIDDEN*4); - float *dh3 = (float*)malloc(SEQ*HIDDEN*4); - float *dx_ffn = (float*)malloc(SEQ*DIM*4); - float *dx2 = (float*)malloc(SEQ*DIM*4); - float *do_out_buf = (float*)malloc(SEQ*DIM*4); - float *dq = (float*)malloc(SEQ*DIM*4); - float *dk = (float*)malloc(SEQ*DIM*4); - float *dv = (float*)malloc(SEQ*DIM*4); - float *dx_attn = (float*)malloc(SEQ*DIM*4); + float *dy = xmf((size_t)SEQ*DIM); // gradient flowing backward + float *dffn = xmf((size_t)SEQ*DIM); + float *dh1 = xmf((size_t)SEQ*HIDDEN); + float *dh3 = xmf((size_t)SEQ*HIDDEN); + float *dx_ffn = xmf((size_t)SEQ*DIM); + float *dx2 = xmf((size_t)SEQ*DIM); + float *do_out_buf = xmf((size_t)SEQ*DIM); + float *dq = xmf((size_t)SEQ*DIM); + float *dk = xmf((size_t)SEQ*DIM); + float *dv = xmf((size_t)SEQ*DIM); + float *dx_attn = xmf((size_t)SEQ*DIM); // x buffer for input to each layer (channel-first [DIM, SEQ]) - float *x_cur = (float*)malloc(SEQ*DIM*4); - float *x_final = (float*)malloc(SEQ*DIM*4); // after final rmsnorm - float *logits = (float*)malloc(SEQ*VOCAB*4); // [VOCAB, SEQ] for cross-entropy - float *dlogits = (float*)malloc(SEQ*VOCAB*4); + float *x_cur = xmf((size_t)SEQ*DIM); + float *x_final = xmf((size_t)SEQ*DIM); // after final rmsnorm + float *logits = xmf((size_t)SEQ*VOCAB); // [VOCAB, SEQ] for cross-entropy + float *dlogits = xmf((size_t)SEQ*VOCAB); // Compile static sdpaBwd2 kernels (no weights, one per layer) Kern *sdpaBwd2[NLAYERS]; @@ -492,7 +492,7 @@ int main(int argc, char *argv[]) { }); // Final RMSNorm backward - float *dx_rms_final = (float*)calloc(SEQ*DIM, 4); + float *dx_rms_final = xcf((size_t)SEQ*DIM); rmsnorm_bwd(dx_rms_final, grms_final, dy, x_cur, rms_final, DIM, SEQ); memcpy(dy, dx_rms_final, SEQ*DIM*4); free(dx_rms_final); @@ -515,11 +515,11 @@ int main(int argc, char *argv[]) { io_read_fp16(kern[L].ffnBwd->ioOut, dh3, DIM+HIDDEN, HIDDEN, SEQ); // dW FFN async - float *capt_dffn = (float*)malloc(SEQ*DIM*4); memcpy(capt_dffn, dffn, SEQ*DIM*4); - float *capt_silu = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_silu, ac->silu_out, SEQ*HIDDEN*4); - float *capt_dh1 = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_dh1, dh1, SEQ*HIDDEN*4); - float *capt_dh3 = (float*)malloc(SEQ*HIDDEN*4); memcpy(capt_dh3, dh3, SEQ*HIDDEN*4); - float *capt_x2n = (float*)malloc(SEQ*DIM*4); memcpy(capt_x2n, ac->x2norm, SEQ*DIM*4); + float *capt_dffn = xmf((size_t)SEQ*DIM); memcpy(capt_dffn, dffn, SEQ*DIM*4); + float *capt_silu = xmf((size_t)SEQ*HIDDEN); memcpy(capt_silu, ac->silu_out, SEQ*HIDDEN*4); + float *capt_dh1 = xmf((size_t)SEQ*HIDDEN); memcpy(capt_dh1, dh1, SEQ*HIDDEN*4); + float *capt_dh3 = xmf((size_t)SEQ*HIDDEN); memcpy(capt_dh3, dh3, SEQ*HIDDEN*4); + float *capt_x2n = xmf((size_t)SEQ*DIM); memcpy(capt_x2n, ac->x2norm, SEQ*DIM*4); dispatch_group_async(dw_grp, dw_q, ^{ cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, HIDDEN, SEQ, 1.0f, capt_dffn, SEQ, capt_silu, SEQ, 1.0f, gr->W2, HIDDEN); @@ -538,8 +538,8 @@ int main(int argc, char *argv[]) { // dWo async memcpy(do_out_buf, dx2, SEQ*DIM*4); - float *capt_do = (float*)malloc(SEQ*DIM*4); memcpy(capt_do, do_out_buf, SEQ*DIM*4); - float *capt_attn = (float*)malloc(SEQ*DIM*4); memcpy(capt_attn, ac->attn_out, SEQ*DIM*4); + float *capt_do = xmf((size_t)SEQ*DIM); memcpy(capt_do, do_out_buf, SEQ*DIM*4); + float *capt_attn = xmf((size_t)SEQ*DIM); memcpy(capt_attn, ac->attn_out, SEQ*DIM*4); dispatch_group_async(dw_grp, dw_q, ^{ cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ, 1.0f, capt_do, SEQ, capt_attn, SEQ, 1.0f, gr->Wo, DIM); @@ -559,10 +559,10 @@ int main(int argc, char *argv[]) { io_read_fp16(kern[L].sdpaBwd1->ioOut, dv, 0, DIM, SEQ); // dWq/dWk/dWv async - float *capt_dq = (float*)malloc(SEQ*DIM*4); memcpy(capt_dq, dq, SEQ*DIM*4); - float *capt_dk = (float*)malloc(SEQ*DIM*4); memcpy(capt_dk, dk, SEQ*DIM*4); - float *capt_dv = (float*)malloc(SEQ*DIM*4); memcpy(capt_dv, dv, SEQ*DIM*4); - float *capt_xn = (float*)malloc(SEQ*DIM*4); memcpy(capt_xn, ac->xnorm, SEQ*DIM*4); + float *capt_dq = xmf((size_t)SEQ*DIM); memcpy(capt_dq, dq, SEQ*DIM*4); + float *capt_dk = xmf((size_t)SEQ*DIM); memcpy(capt_dk, dk, SEQ*DIM*4); + float *capt_dv = xmf((size_t)SEQ*DIM); memcpy(capt_dv, dv, SEQ*DIM*4); + float *capt_xn = xmf((size_t)SEQ*DIM); memcpy(capt_xn, ac->xnorm, SEQ*DIM*4); dispatch_group_async(dw_grp, dw_q, ^{ cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, DIM, DIM, SEQ, 1.0f, capt_dq, SEQ, capt_xn, SEQ, 1.0f, gr->Wq, DIM); @@ -580,7 +580,7 @@ int main(int argc, char *argv[]) { io_read_fp16(kern[L].qkvBwd->ioOut, dx_attn, 0, DIM, SEQ); // RMSNorm1 backward (using saved layer input) - float *dx_rms1 = (float*)calloc(SEQ*DIM, 4); + float *dx_rms1 = xcf((size_t)SEQ*DIM); rmsnorm_bwd(dx_rms1, gr->rms_att, dx_attn, ac->layer_in, lw[L].rms_att, DIM, SEQ); // dy for next layer (going backward) = dx_rms1 + dx2 residual