82 lines
3.4 KiB
C
82 lines
3.4 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#ifndef FOUNDRY_FT_LAYER_OPS_H
|
|
#define FOUNDRY_FT_LAYER_OPS_H
|
|
|
|
#include "ternarize.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* Layer forward/backward dispatch for ternarizer.
|
|
*
|
|
* Each layer applies ternary weights via {add, nop, sub} integer ops
|
|
* scaled by per-group float scales. Forward computes activations,
|
|
* backward accumulates gradient statistics.
|
|
*
|
|
* LLaMA-family layer structure:
|
|
* RMSNorm -> Q/K/V proj -> RoPE -> GQA attention -> O proj -> residual
|
|
* -> RMSNorm -> gate/up proj -> SiLU -> down proj -> residual
|
|
*/
|
|
|
|
/* Forward pass through one transformer layer with ternary weights.
|
|
* input: [batch, n_embd] (pre-residual activations)
|
|
* output: [batch, n_embd] (post-residual activations)
|
|
* Adds residual internally. */
|
|
void ft_layer_forward(const ft_layer_state *state, const ft_config *cfg,
|
|
const float *input, float *output, int batch_size);
|
|
|
|
/* Backward pass through one transformer layer.
|
|
* Accumulates gradient statistics (grad_sum, grad_sq_sum, hessian_diag)
|
|
* into the provided ft_grad_stats structure.
|
|
* d_output: [batch, n_embd] gradient w.r.t. layer output
|
|
* d_input: [batch, n_embd] gradient w.r.t. layer input (accumulated) */
|
|
void ft_layer_backward(const ft_layer_state *state, const ft_config *cfg,
|
|
const float *input, const float *d_output,
|
|
float *d_input, ft_grad_stats *stats, int batch_size);
|
|
|
|
/* ---- Helper ops ---- */
|
|
|
|
/* RMSNorm: out[i] = weight[i] * x[i] / rms(x), where rms = sqrt(mean(x^2) + eps) */
|
|
void ft_rmsnorm(const float *x, const float *weight, float *out,
|
|
int n, float eps);
|
|
|
|
/* RMSNorm backward: d_x given d_out, x, weight, eps */
|
|
void ft_rmsnorm_backward(const float *x, const float *weight,
|
|
const float *d_out, float *d_x,
|
|
int n, float eps);
|
|
|
|
/* Ternary matvec: out[row] = scale * sum_k(ternary[row,k] * x[k])
|
|
* Uses sign+mask bit arrays + per-group scales.
|
|
* state_offset: offset into the layer state's sign/mask arrays. */
|
|
void ft_ternary_matvec(const ft_layer_state *state, int64_t state_offset,
|
|
int rows, int cols,
|
|
const float *x, float *out);
|
|
|
|
/* Ternary matvec backward (d_input): d_x[k] += scale * sum_row(ternary[row,k] * d_out[row]) */
|
|
void ft_ternary_matvec_backward_input(const ft_layer_state *state, int64_t state_offset,
|
|
int rows, int cols,
|
|
const float *d_out, float *d_x);
|
|
|
|
/* Accumulate weight gradient stats for a ternary matmul.
|
|
* For each weight [row,col]: grad += x[col] * d_out[row]
|
|
* Accumulates into stats->grad_sum, grad_sq_sum, hessian_diag. */
|
|
void ft_ternary_matvec_backward_weight(const ft_layer_state *state, int64_t state_offset,
|
|
int rows, int cols,
|
|
const float *x, const float *d_out,
|
|
ft_grad_stats *stats);
|
|
|
|
/* SiLU activation: out[i] = x[i] * sigmoid(x[i]) */
|
|
void ft_silu(const float *x, float *out, int n);
|
|
|
|
/* SiLU backward: d_x[i] = d_out[i] * (sigmoid(x[i]) + x[i] * sigmoid(x[i]) * (1 - sigmoid(x[i]))) */
|
|
void ft_silu_backward(const float *x, const float *d_out, float *d_x, int n);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_FT_LAYER_OPS_H */
|