/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /* * optim.cu — CUDA GPU optimizer-step kernels for Foundry (Phase 3, Chunk A). * * Target: sm_37 (K80, Kepler). Compile: nvcc -ptx -arch=sm_37 ... (or included * directly by the parity harness test_optim_cuda.cu and linked with the CPU * oracle objects). * * Two GPU optimizer steps, each with an extern "C" host launcher, parity-checked * against the CPU optimizer: * * 1. sr4_step_cuda — the SR4 4-bit stochastic-rounded latent-accumulator * update (ternary_optim_step). BIT-IDENTICAL to the * CPU: same SplitMix64 counter-hash keyed on * (seed, stratum, idx_base+i), same stratified Weyl * low-discrepancy dither, same 8.8 fixed-point grid * snap, same saturating add / clamp order. * 2. adamw_step_cuda — the continuous-param AdamW update (foundry optim.c * adam_step, is_adamw path). Matches beta1/beta2/eps/ * weight-decay/bias-correction; elementwise so it is * order-independent -> parity to fp rounding. * * The updates are O(params) elementwise (no matmul, no cuBLAS). * * Bit-exactness note for SR4: the ONLY float->discrete boundary is the per-matrix * step scale s = target_step/(mean|g|+1e-12) and the 8.8 grid snap lrintf(prop). * mean|g| is a sequential DOUBLE reduction in the CPU code. To be bit-identical * we (a) run the reduction on ONE GPU thread in the SAME index order in double, * and (b) build the CPU oracle without -ffast-math so gcc cannot reassociate / * vectorize that double reduction (-fassociative-math off). Everything after the * grid snap is pure integer (the counter-hash, dither, saturating add) and is * flag-immune, so it is bit-identical by construction. */ #include #include #include /* Stratified low-discrepancy dither constants (ternary_optim.c verbatim). */ #define TO_LD_STRATUM_BITS 8 #define TO_LD_STRIDE 159u /* --------------------------------------------------------------------------- * SplitMix64-style counter-based hash — byte-identical to ternary_optim.c. * ------------------------------------------------------------------------- */ __device__ static inline uint64_t d_sm64_mix(uint64_t x) { x += 0x9E3779B97F4A7C15ULL; x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL; x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL; return x ^ (x >> 31); } __device__ static inline uint64_t d_sr_hash(uint64_t seed, uint64_t batch_counter, uint64_t global_idx) { uint64_t x = d_sm64_mix(seed ^ 0xD1B54A32D192ED03ULL); x = d_sm64_mix(x ^ batch_counter); x = d_sm64_mix(x ^ global_idx); return x; } /* --------------------------------------------------------------------------- * SR4 kernel 1: per-matrix step scale s, computed by ONE thread in double in the * SAME sequential index order as the CPU (ternary_optim.c lines 163-167). This * is the only reduction; the O(params) work is the parallel update below. * ------------------------------------------------------------------------- */ extern "C" __global__ void fc_sr4_scale(const float *grad_sum, int numel, float inv_bs, float target_step, float *s_out) { if (blockIdx.x != 0 || threadIdx.x != 0) return; double abs_sum = 0.0; for (int i = 0; i < numel; i++) abs_sum += fabsf(grad_sum[i] * inv_bs); float mean_abs = (float)(abs_sum / numel); *s_out = target_step / (mean_abs + 1e-12f); } /* --------------------------------------------------------------------------- * SR4 kernel 2: the parallel elementwise stochastic-rounded integer update into * the int4 accumulator q. Byte-identical to ternary_optim.c lines 169-203. * stratum / j / weyl are pure functions of batch_counter (computed here exactly * as the CPU does). * ------------------------------------------------------------------------- */ extern "C" __global__ void fc_sr4_update(int8_t *q, const float *grad_sum, int numel, float inv_bs, const float *s_in, uint64_t seed, int64_t idx_base, uint64_t batch_counter, int q_clamp, int use_dither) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= numel) return; float s = *s_in; uint64_t stratum = batch_counter >> TO_LD_STRATUM_BITS; uint32_t j = (uint32_t)(batch_counter & ((1u << TO_LD_STRATUM_BITS) - 1u)); uint32_t weyl = (TO_LD_STRIDE * j) & 255u; float g = grad_sum[i] * inv_bs; float prop = -g * s * 256.0f; if (prop > 4096.0f) prop = 4096.0f; if (prop < -4096.0f) prop = -4096.0f; int32_t step_fx = (int32_t)lrintf(prop); int32_t k = step_fx >> 8; int32_t f = step_fx & 255; uint64_t h = d_sr_hash(seed, stratum, (uint64_t)(idx_base + i)); uint32_t u; if (use_dither) u = ((uint32_t)(h & 255) + weyl) & 255u; /* stratified Weyl */ else u = (uint32_t)(h & 255); /* plain i.i.d. */ int32_t sr_bit = ((int32_t)u < f) ? 1 : 0; int32_t nq = (int32_t)q[i] + k + sr_bit; if (nq > q_clamp) nq = q_clamp; if (nq < -q_clamp) nq = -q_clamp; q[i] = (int8_t)nq; } /* --------------------------------------------------------------------------- * opt#3 graph variant of fc_sr4_update: byte-identical math, but the per-step * batch_counter (which MUST advance every replay or the stochastic-rounding * dither repeats and every step reuses the same round) is read from a device * buffer instead of a by-value arg. The driver refreshes *batch_counter_dev * each step before cudaGraphLaunch, so the captured node sees a fresh counter on * every replay while its pointers stay fixed. Reading the counter from memory * rather than a register changes no arithmetic -> parity is bit-for-bit. * ------------------------------------------------------------------------- */ extern "C" __global__ void fc_sr4_update_g(int8_t *q, const float *grad_sum, int numel, float inv_bs, const float *s_in, uint64_t seed, int64_t idx_base, const unsigned long long *batch_counter_dev, int q_clamp, int use_dither) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= numel) return; uint64_t batch_counter = (uint64_t)(*batch_counter_dev); float s = *s_in; uint64_t stratum = batch_counter >> TO_LD_STRATUM_BITS; uint32_t j = (uint32_t)(batch_counter & ((1u << TO_LD_STRATUM_BITS) - 1u)); uint32_t weyl = (TO_LD_STRIDE * j) & 255u; float g = grad_sum[i] * inv_bs; float prop = -g * s * 256.0f; if (prop > 4096.0f) prop = 4096.0f; if (prop < -4096.0f) prop = -4096.0f; int32_t step_fx = (int32_t)lrintf(prop); int32_t k = step_fx >> 8; int32_t f = step_fx & 255; uint64_t h = d_sr_hash(seed, stratum, (uint64_t)(idx_base + i)); uint32_t u; if (use_dither) u = ((uint32_t)(h & 255) + weyl) & 255u; else u = (uint32_t)(h & 255); int32_t sr_bit = ((int32_t)u < f) ? 1 : 0; int32_t nq = (int32_t)q[i] + k + sr_bit; if (nq > q_clamp) nq = q_clamp; if (nq < -q_clamp) nq = -q_clamp; q[i] = (int8_t)nq; } /* --------------------------------------------------------------------------- * AdamW kernel — elementwise, byte-for-byte the per-element math of optim.c * adam_step() on the FO_ADAMW path (decoupled weight decay). bc1/bc2 are passed * in (host computes powf, identical to the CPU) so t-dependent bias correction * matches exactly. * ------------------------------------------------------------------------- */ extern "C" __global__ void fc_adamw_update(float *data, const float *grad, float *m, float *v, int64_t n, float lr, float beta1, float beta2, float eps, float weight_decay, float bc1, float bc2, int is_adamw) { int64_t j = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; if (j >= n) return; float g = grad[j]; /* For classic Adam (not AdamW), fold weight decay into the gradient. */ if (!is_adamw && weight_decay != 0.0f) g += weight_decay * data[j]; float mj = beta1 * m[j] + (1.0f - beta1) * g; float vj = beta2 * v[j] + (1.0f - beta2) * g * g; m[j] = mj; v[j] = vj; float m_hat = mj / bc1; float v_hat = vj / bc2; float update = lr * m_hat / (sqrtf(v_hat) + eps); float d = data[j]; if (is_adamw && weight_decay != 0.0f) /* decoupled weight decay */ d -= lr * weight_decay * d; d -= update; data[j] = d; } /* --------------------------------------------------------------------------- * opt#3 graph variant of fc_adamw_update: identical math, but the per-step * scalars lr / bc1 / bc2 (cosine LR + t-dependent bias correction, all change * every step) are read from a device buffer opt[3] = {lr, bc1, bc2}. The host * computes them EXACTLY as before (powf on the CPU) and stages them into that * buffer each step before replay, so the float bits are identical — delivering * them via memory instead of by-value changes no arithmetic (bit-for-bit parity). * ------------------------------------------------------------------------- */ extern "C" __global__ void fc_adamw_update_g(float *data, const float *grad, float *m, float *v, int64_t n, const float *opt, float beta1, float beta2, float eps, float weight_decay, int is_adamw) { int64_t j = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; if (j >= n) return; float lr = opt[0], bc1 = opt[1], bc2 = opt[2]; float g = grad[j]; if (!is_adamw && weight_decay != 0.0f) g += weight_decay * data[j]; float mj = beta1 * m[j] + (1.0f - beta1) * g; float vj = beta2 * v[j] + (1.0f - beta2) * g * g; m[j] = mj; v[j] = vj; float m_hat = mj / bc1; float v_hat = vj / bc2; float update = lr * m_hat / (sqrtf(v_hat) + eps); float d = data[j]; if (is_adamw && weight_decay != 0.0f) d -= lr * weight_decay * d; d -= update; data[j] = d; } /* =========================================================================== * Host launchers * ========================================================================= */ #define CUDA_OK(call) do { cudaError_t _e = (call); if (_e != cudaSuccess) { \ fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(_e), \ __FILE__, __LINE__); return -1; } } while (0) #include /* * One SR4 step on GPU 0. q_io[0..numel) is updated in place (int8, clamped to * [-q_clamp,+q_clamp]). Mirrors ternary_optim_step(opt, grad_sum, batch_size, * batch_counter, target_step) exactly. Returns 0 on success. */ extern "C" int sr4_step_cuda(int8_t *q_io, const float *grad_sum, int numel, int batch_size, uint64_t batch_counter, float target_step, uint64_t seed, int64_t idx_base, int q_clamp, int use_dither) { if (numel <= 0 || batch_size <= 0) return -1; float inv_bs = 1.0f / (float)batch_size; int8_t *d_q = NULL; float *d_g = NULL, *d_s = NULL; CUDA_OK(cudaMalloc(&d_q, (size_t)numel * sizeof(int8_t))); CUDA_OK(cudaMalloc(&d_g, (size_t)numel * sizeof(float))); CUDA_OK(cudaMalloc(&d_s, sizeof(float))); CUDA_OK(cudaMemcpy(d_q, q_io, (size_t)numel * sizeof(int8_t), cudaMemcpyHostToDevice)); CUDA_OK(cudaMemcpy(d_g, grad_sum, (size_t)numel * sizeof(float), cudaMemcpyHostToDevice)); fc_sr4_scale<<<1, 1>>>(d_g, numel, inv_bs, target_step, d_s); int threads = 256; int blocks = (numel + threads - 1) / threads; fc_sr4_update<<>>(d_q, d_g, numel, inv_bs, d_s, seed, idx_base, batch_counter, q_clamp, use_dither); CUDA_OK(cudaGetLastError()); CUDA_OK(cudaDeviceSynchronize()); CUDA_OK(cudaMemcpy(q_io, d_q, (size_t)numel * sizeof(int8_t), cudaMemcpyDeviceToHost)); cudaFree(d_q); cudaFree(d_g); cudaFree(d_s); return 0; } /* * One AdamW step on GPU 0. data[0..n), m[0..n), v[0..n) updated in place. * `t` is the 1-based step index for bias correction (CPU: step_count+1). * Mirrors optim.c adam_step() on the FO_ADAMW path. Returns 0 on success. */ extern "C" int adamw_step_cuda(float *data, const float *grad, float *m, float *v, int64_t n, int t, float lr, float beta1, float beta2, float eps, float weight_decay, int is_adamw) { if (n <= 0 || t < 1) return -1; float bc1 = 1.0f - powf(beta1, (float)t); float bc2 = 1.0f - powf(beta2, (float)t); float *d_data = NULL, *d_grad = NULL, *d_m = NULL, *d_v = NULL; size_t bytes = (size_t)n * sizeof(float); CUDA_OK(cudaMalloc(&d_data, bytes)); CUDA_OK(cudaMalloc(&d_grad, bytes)); CUDA_OK(cudaMalloc(&d_m, bytes)); CUDA_OK(cudaMalloc(&d_v, bytes)); CUDA_OK(cudaMemcpy(d_data, data, bytes, cudaMemcpyHostToDevice)); CUDA_OK(cudaMemcpy(d_grad, grad, bytes, cudaMemcpyHostToDevice)); CUDA_OK(cudaMemcpy(d_m, m, bytes, cudaMemcpyHostToDevice)); CUDA_OK(cudaMemcpy(d_v, v, bytes, cudaMemcpyHostToDevice)); int threads = 256; int64_t blocks = (n + threads - 1) / threads; fc_adamw_update<<<(unsigned)blocks, threads>>>(d_data, d_grad, d_m, d_v, n, lr, beta1, beta2, eps, weight_decay, bc1, bc2, is_adamw); CUDA_OK(cudaGetLastError()); CUDA_OK(cudaDeviceSynchronize()); CUDA_OK(cudaMemcpy(data, d_data, bytes, cudaMemcpyDeviceToHost)); CUDA_OK(cudaMemcpy(m, d_m, bytes, cudaMemcpyDeviceToHost)); CUDA_OK(cudaMemcpy(v, d_v, bytes, cudaMemcpyDeviceToHost)); cudaFree(d_data); cudaFree(d_grad); cudaFree(d_m); cudaFree(d_v); return 0; }