103 lines
4.9 KiB
Plaintext
103 lines
4.9 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* fm_arena.cuh — device scratch bump-arena for the resident GPU trainer.
|
|
*
|
|
* Optimization #2 (resident scratch): the per-step forward/backward used to
|
|
* cudaMalloc/cudaFree ~1500 transient device buffers per training step — each
|
|
* cudaFree synchronizes the K80, which dominated step time at the 1M config.
|
|
* This arena pre-allocates ONE device slab and hands out bump-pointer slices;
|
|
* it is reset (rewound to empty) at the top of every fm_forward / fm_backward,
|
|
* so no transient buffer is ever cudaMalloc'd/cudaFree'd in steady state.
|
|
*
|
|
* Parity: this changes ONLY memory management. Every kernel/GEMM sees the same
|
|
* values in the same order, so the loss curve is bit-identical to the pre-arena
|
|
* trainer (both SR4 and AdamW arms).
|
|
*
|
|
* Dual-mode (so the standalone parity/test builds are byte-identical):
|
|
* - active (fmA_init called, train_gpu path): fmA_alloc bumps the slab,
|
|
* fmA_free is a no-op, fmA_reset rewinds. Buffers freed en-masse at reset.
|
|
* - inactive(default, test_*_cuda builds that include the merged .cu files
|
|
* without ever calling fmA_init): fmA_alloc == cudaMalloc, fmA_free ==
|
|
* cudaFree, fmA_reset is a no-op. Identical to the original code.
|
|
*
|
|
* Robustness: if the slab is exhausted while active, fmA_alloc falls back to a
|
|
* tracked cudaMalloc (freed at the next fmA_reset), so correctness never depends
|
|
* on the slab being large enough — only performance does.
|
|
*/
|
|
#ifndef FM_ARENA_CUH
|
|
#define FM_ARENA_CUH
|
|
#include <cuda_runtime.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <stdint.h>
|
|
|
|
#define FMA_OVF_MAX 512
|
|
typedef struct {
|
|
char *base; /* device slab */
|
|
size_t cap; /* slab bytes */
|
|
size_t off; /* bump offset */
|
|
int active; /* 0 => fall back to cudaMalloc/cudaFree */
|
|
void *ovf[FMA_OVF_MAX]; /* slab-exhaustion fallbacks, freed at reset */
|
|
int novf;
|
|
size_t hiwater; /* peak off (diagnostic) */
|
|
} fm_arena_t;
|
|
|
|
static fm_arena_t g_fmA = {0,0,0,0,{0},0,0};
|
|
|
|
/* ---- CUDA Graph capture support (opt #3) --------------------------------
|
|
* The resident train step is captured once as a CUDA graph and replayed each
|
|
* step to escape the launch-bound floor. Two globals steer capture-safety:
|
|
* g_fm_stream : the work stream every per-step memcpy/memset is issued on.
|
|
* 0 in the standalone test builds (legacy default stream,
|
|
* matches the kernels' <<<...,0>>>); set to cudaStreamPerThread
|
|
* by the train_gpu driver (which compiles with
|
|
* --default-stream per-thread, so its <<<...,0>>> kernels and
|
|
* cuBLAS-on-per-thread all share one capturable stream).
|
|
* g_fm_capturing : nonzero WHILE a train step is being stream-captured. Host
|
|
* syncs and the per-step H2D input copies / D2H loss readback
|
|
* are skipped during capture (a sync would abort the capture;
|
|
* inputs/loss are marshalled by the driver around replay).
|
|
* In every non-capture path (all test builds, the eval forwards, the non-graph
|
|
* fallback) both are their defaults and behaviour is byte-identical to before. */
|
|
static cudaStream_t g_fm_stream = 0;
|
|
static int g_fm_capturing = 0;
|
|
/* skip the terminal device sync while capturing (replay is synced by the driver) */
|
|
static inline void fm_sync(void){ if(!g_fm_capturing) cudaDeviceSynchronize(); }
|
|
|
|
static inline void fmA_init(size_t cap_bytes){
|
|
if(g_fmA.base){ cudaFree(g_fmA.base); g_fmA.base=NULL; }
|
|
if(cudaMalloc((void**)&g_fmA.base, cap_bytes)!=cudaSuccess){
|
|
fprintf(stderr,"fmA_init: cudaMalloc(%zu) failed\n",cap_bytes); exit(1);
|
|
}
|
|
g_fmA.cap=cap_bytes; g_fmA.off=0; g_fmA.active=1; g_fmA.novf=0; g_fmA.hiwater=0;
|
|
}
|
|
static inline void fmA_reset(void){
|
|
if(!g_fmA.active) return;
|
|
for(int i=0;i<g_fmA.novf;i++) cudaFree(g_fmA.ovf[i]);
|
|
g_fmA.novf=0; g_fmA.off=0;
|
|
}
|
|
/* allocate n bytes of device scratch (256B-aligned). */
|
|
static inline void* fmA_alloc(size_t n){
|
|
if(!g_fmA.active){
|
|
void*p=NULL; if(cudaMalloc(&p,n)!=cudaSuccess){
|
|
fprintf(stderr,"fmA_alloc(fallback) cudaMalloc(%zu) failed\n",n); exit(1);} return p;
|
|
}
|
|
size_t a=(n+255)&~(size_t)255;
|
|
if(g_fmA.off+a<=g_fmA.cap){
|
|
void*p=g_fmA.base+g_fmA.off; g_fmA.off+=a;
|
|
if(g_fmA.off>g_fmA.hiwater) g_fmA.hiwater=g_fmA.off;
|
|
return p;
|
|
}
|
|
/* slab exhausted: tracked cudaMalloc fallback (freed next reset) */
|
|
void*p=NULL; if(cudaMalloc(&p,n)!=cudaSuccess){
|
|
fprintf(stderr,"fmA_alloc(ovf) cudaMalloc(%zu) failed\n",n); exit(1);}
|
|
if(g_fmA.novf<FMA_OVF_MAX) g_fmA.ovf[g_fmA.novf++]=p;
|
|
return p;
|
|
}
|
|
/* free: no-op while active (arena owns it); real cudaFree when inactive. */
|
|
static inline void fmA_free(void*p){
|
|
if(!g_fmA.active){ if(p) cudaFree(p); }
|
|
}
|
|
#endif /* FM_ARENA_CUH */
|