mlp_kernels_bf16.c

Go to the documentation of this file.

/**
 * @file mlp_kernels_bf16.c
 * @brief Optimized BF16 MLP Kernels
 *
 * CK-ENGINE KERNEL RULES:
 * =======================
 * 1. NO malloc/free - memory via bump allocator, pointers passed in
 * 2. NO OpenMP - parallelization at orchestrator/codegen layer
 * 3. API must define: inputs, outputs, workspace, and memory layouts
 * 4. Pure computation - deterministic, no side effects
 *
 * After changes: make test && make llamacpp-parity-full
 *
 * Uses direct BF16 GEMM instead of converting to FP32.
 * Layout: input[T,D] -> fc1[T,4D] -> GELU -> fc2[T,D]
 *
 * All functions use caller-provided scratch buffers (no internal malloc).
 */
 
#include <stddef.h>
#include <stdint.h>
#include <math.h>
 
#if defined(__AVX512F__)
#include <immintrin.h>
#endif
 
#ifdef _OPENMP
#include <omp.h>
#endif
 
#include "bf16_utils.h"
#include "ckernel_engine.h"
 
// Suppress false positive warnings about uninitialized variables
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
 
/* Forward declaration of optimized BF16 GEMM */
extern void gemm_bf16_fp32out(const uint16_t *A, const uint16_t *B,
                               const float *bias, float *C,
                               int M, int N, int K);
 
/* GELU activation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) */
static inline float gelu_scalar(float x)
{
    const float c = 0.7978845608f;  /* sqrt(2/pi) */
    const float k = 0.044715f;
    float x3 = x * x * x;
    return 0.5f * x * (1.0f + tanhf(c * (x + k * x3)));
}
 
static inline float gelu_derivative_scalar(float x)
{
    const float c = 0.7978845608f;  /* sqrt(2/pi) */
    const float k = 0.044715f;
    const float x2 = x * x;
    const float x3 = x2 * x;
    const float g = c * (x + k * x3);
    const float tanh_g = tanhf(g);
    const float sech2_g = 1.0f - tanh_g * tanh_g;
    const float g_prime = c * (1.0f + 3.0f * k * x2);
    return 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
}
 
#if defined(__AVX512F__)
/* Vectorized GELU using polynomial approximation of tanh */
static inline __m512 gelu_avx512(__m512 x)
{
    const __m512 c = _mm512_set1_ps(0.7978845608f);
    const __m512 k = _mm512_set1_ps(0.044715f);
    const __m512 half = _mm512_set1_ps(0.5f);
    const __m512 one = _mm512_set1_ps(1.0f);
 
    __m512 x2 = _mm512_mul_ps(x, x);
    __m512 x3 = _mm512_mul_ps(x2, x);
 
    __m512 inner = _mm512_fmadd_ps(k, x3, x);
    inner = _mm512_mul_ps(c, inner);
 
    __m512 inner2 = _mm512_mul_ps(inner, inner);
    __m512 num = _mm512_add_ps(_mm512_set1_ps(27.0f), inner2);
    __m512 den = _mm512_fmadd_ps(_mm512_set1_ps(9.0f), inner2, _mm512_set1_ps(27.0f));
    __m512 tanh_approx = _mm512_mul_ps(inner, _mm512_div_ps(num, den));
 
    tanh_approx = _mm512_min_ps(tanh_approx, one);
    tanh_approx = _mm512_max_ps(tanh_approx, _mm512_set1_ps(-1.0f));
 
    __m512 result = _mm512_add_ps(one, tanh_approx);
    result = _mm512_mul_ps(half, _mm512_mul_ps(x, result));
 
    return result;
}
#endif
 
/**
 * Optimized MLP Forward (BF16 weights, FP32 activations)
 *
 * Caller-provided scratch buffers:
 *   scratch_bias1_f: [4*D] floats
 *   scratch_bias2_f: [D] floats
 *   scratch_fc1_bf16: [T * 4*D] uint16_t (BF16)
 */
void mlp_token_parallel_bf16(const uint16_t *input,
                             const uint16_t *W_fc1,
                             const uint16_t *b_fc1,
                             const uint16_t *W_fc2,
                             const uint16_t *b_fc2,
                             float *fc1_output,
                             float *output,
                             int T,
                             int aligned_dim,
                             int num_threads,
                             float *scratch_bias1_f,
                             float *scratch_bias2_f,
                             uint16_t *scratch_fc1_bf16)
{
    if (!input || !W_fc1 || !b_fc1 || !W_fc2 || !b_fc2 || !fc1_output || !output) return;
    if (!scratch_bias1_f || !scratch_bias2_f || !scratch_fc1_bf16) return;
 
    (void)num_threads;
    const int D = aligned_dim;
    const int fourD = 4 * D;
 
    /* Convert biases to FP32 */
    for (int i = 0; i < fourD; ++i) {
        scratch_bias1_f[i] = bf16_to_float(b_fc1[i]);
    }
    for (int i = 0; i < D; ++i) {
        scratch_bias2_f[i] = bf16_to_float(b_fc2[i]);
    }
 
    /* FC1: [T, D] x [4D, D].T -> [T, 4D] */
    gemm_bf16_fp32out(input, W_fc1, scratch_bias1_f, fc1_output, T, fourD, D);
 
    /* GELU activation */
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int t = 0; t < T; ++t) {
        float *row = fc1_output + (size_t)t * fourD;
        int j = 0;
        for (; j <= fourD - 16; j += 16) {
            __m512 x = _mm512_loadu_ps(row + j);
            _mm512_storeu_ps(row + j, gelu_avx512(x));
        }
        for (; j < fourD; ++j) {
            row[j] = gelu_scalar(row[j]);
        }
    }
#else
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            fc1_output[t * fourD + j] = gelu_scalar(fc1_output[t * fourD + j]);
        }
    }
#endif
 
    /* Convert FP32 activations to BF16 */
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int t = 0; t < T; ++t) {
        float *src = fc1_output + (size_t)t * fourD;
        uint16_t *dst = scratch_fc1_bf16 + (size_t)t * fourD;
        int j = 0;
        for (; j <= fourD - 16; j += 16) {
            __m512 fp32 = _mm512_loadu_ps(src + j);
            __m512i as_int = _mm512_castps_si512(fp32);
            __m512i lsb = _mm512_srli_epi32(as_int, 16);
            lsb = _mm512_and_si512(lsb, _mm512_set1_epi32(1));
            __m512i rounding = _mm512_add_epi32(_mm512_set1_epi32(0x7FFF), lsb);
            __m512i rounded = _mm512_add_epi32(as_int, rounding);
            __m512i shifted = _mm512_srli_epi32(rounded, 16);
            __m256i bf16 = _mm512_cvtepi32_epi16(shifted);
            _mm256_storeu_si256((__m256i *)(dst + j), bf16);
        }
        for (; j < fourD; ++j) {
            dst[j] = float_to_bf16(src[j]);
        }
    }
#else
    for (size_t i = 0; i < (size_t)T * fourD; ++i) {
        scratch_fc1_bf16[i] = float_to_bf16(fc1_output[i]);
    }
#endif
 
    /* FC2: BF16 GEMM with FP32 output */
    gemm_bf16_fp32out(scratch_fc1_bf16, W_fc2, scratch_bias2_f, output, T, D, fourD);
}
 
/**
 * Alternative: Fully FP32 activations throughout
 *
 * Caller-provided scratch buffers:
 *   scratch_input_f: [T * D] floats
 *   scratch_bias1_f: [4*D] floats
 *   scratch_bias2_f: [D] floats
 *   scratch_fc1_bf16: [T * 4*D] uint16_t (BF16)
 */
void mlp_token_parallel_bf16_fp32act(const uint16_t *input,
                                      const uint16_t *W_fc1,
                                      const uint16_t *b_fc1,
                                      const uint16_t *W_fc2,
                                      const uint16_t *b_fc2,
                                      float *fc1_output,
                                      float *output,
                                      int T,
                                      int aligned_dim,
                                      int num_threads,
                                      float *scratch_input_f,
                                      float *scratch_bias1_f,
                                      float *scratch_bias2_f,
                                      uint16_t *scratch_fc1_bf16)
{
    if (!input || !W_fc1 || !b_fc1 || !W_fc2 || !b_fc2 || !fc1_output || !output) return;
    if (!scratch_input_f || !scratch_bias1_f || !scratch_bias2_f || !scratch_fc1_bf16) return;
 
    (void)num_threads;
    const int D = aligned_dim;
    const int fourD = 4 * D;
 
    /* Convert input and biases to FP32 */
    bf16_tensor_to_float(input, scratch_input_f, (size_t)T * D);
    bf16_tensor_to_float(b_fc1, scratch_bias1_f, fourD);
    bf16_tensor_to_float(b_fc2, scratch_bias2_f, D);
 
    /* FC1 */
    gemm_bf16_fp32out(input, W_fc1, scratch_bias1_f, fc1_output, T, fourD, D);
 
    /* GELU */
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int t = 0; t < T; ++t) {
        float *row = fc1_output + (size_t)t * fourD;
        int j = 0;
        for (; j <= fourD - 16; j += 16) {
            __m512 x = _mm512_loadu_ps(row + j);
            _mm512_storeu_ps(row + j, gelu_avx512(x));
        }
        for (; j < fourD; ++j) {
            row[j] = gelu_scalar(row[j]);
        }
    }
#else
    for (size_t i = 0; i < (size_t)T * fourD; ++i) {
        fc1_output[i] = gelu_scalar(fc1_output[i]);
    }
#endif
 
    /* Convert fc1_output to BF16 for FC2 */
    float_tensor_to_bf16(fc1_output, scratch_fc1_bf16, (size_t)T * fourD);
    gemm_bf16_fp32out(scratch_fc1_bf16, W_fc2, scratch_bias2_f, output, T, D, fourD);
}
 
 
/**
 * BF16 MLP backward with FP32 gradient accumulation.
 *
 * Forward contract matched here:
 *   z1 = input_bf16 @ W_fc1_bf16.T + b_fc1_bf16
 *   h  = GELU(z1)
 *   hq = round_to_bf16(h)
 *   y  = hq @ W_fc2_bf16.T + b_fc2_bf16
 *
 * Gradients are accumulated and written as FP32. The BF16 activation cast is
 * treated like PyTorch's mixed-precision cast: gradient flows through to h,
 * while d_W_fc2 uses the rounded hq values that FC2 actually consumed.
 */
void mlp_token_parallel_bf16_backward_mixed(const uint16_t *input,
                                            const uint16_t *W_fc1,
                                            const uint16_t *b_fc1,
                                            const uint16_t *W_fc2,
                                            const uint16_t *d_output,
                                            float *d_input,
                                            float *d_W_fc1,
                                            float *d_b_fc1,
                                            float *d_W_fc2,
                                            float *d_b_fc2,
                                            int T,
                                            int aligned_dim,
                                            int num_threads,
                                            float *scratch_fc1_pre,
                                            uint16_t *scratch_fc1_act_bf16,
                                            float *scratch_d_fc1)
{
    if (!input || !W_fc1 || !b_fc1 || !W_fc2 || !d_output) return;
    if (!scratch_fc1_pre || !scratch_fc1_act_bf16 || !scratch_d_fc1) return;
    if (T <= 0 || aligned_dim <= 0) return;
 
    (void)num_threads;
    const int D = aligned_dim;
    const int fourD = 4 * D;
 
    /* Recompute FC1 pre-activation and the rounded activation consumed by FC2. */
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            float sum = bf16_to_float(b_fc1[j]);
            for (int i = 0; i < D; ++i) {
                const float x = bf16_to_float(input[(size_t)t * (size_t)D + (size_t)i]);
                const float w = bf16_to_float(W_fc1[(size_t)j * (size_t)D + (size_t)i]);
                sum += x * w;
            }
            scratch_fc1_pre[(size_t)t * (size_t)fourD + (size_t)j] = sum;
            scratch_fc1_act_bf16[(size_t)t * (size_t)fourD + (size_t)j] = float_to_bf16(gelu_scalar(sum));
        }
    }
 
    if (d_input) {
        for (int t = 0; t < T; ++t) {
            for (int i = 0; i < D; ++i) {
                d_input[(size_t)t * (size_t)D + (size_t)i] = 0.0f;
            }
        }
    }
    if (d_W_fc1) {
        for (size_t i = 0; i < (size_t)fourD * (size_t)D; ++i) d_W_fc1[i] = 0.0f;
    }
    if (d_b_fc1) {
        for (int j = 0; j < fourD; ++j) d_b_fc1[j] = 0.0f;
    }
    if (d_W_fc2) {
        for (size_t i = 0; i < (size_t)D * (size_t)fourD; ++i) d_W_fc2[i] = 0.0f;
    }
    if (d_b_fc2) {
        for (int o = 0; o < D; ++o) d_b_fc2[o] = 0.0f;
    }
 
    /* d_W_fc2, d_b_fc2, and d_h = d_output @ W_fc2. */
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            float dh = 0.0f;
            const float hq = bf16_to_float(scratch_fc1_act_bf16[(size_t)t * (size_t)fourD + (size_t)j]);
            for (int o = 0; o < D; ++o) {
                const float dy = bf16_to_float(d_output[(size_t)t * (size_t)D + (size_t)o]);
                const float w2 = bf16_to_float(W_fc2[(size_t)o * (size_t)fourD + (size_t)j]);
                dh += dy * w2;
                if (d_W_fc2) {
                    d_W_fc2[(size_t)o * (size_t)fourD + (size_t)j] += dy * hq;
                }
            }
            const float z = scratch_fc1_pre[(size_t)t * (size_t)fourD + (size_t)j];
            scratch_d_fc1[(size_t)t * (size_t)fourD + (size_t)j] = dh * gelu_derivative_scalar(z);
        }
        if (d_b_fc2) {
            for (int o = 0; o < D; ++o) {
                d_b_fc2[o] += bf16_to_float(d_output[(size_t)t * (size_t)D + (size_t)o]);
            }
        }
    }
 
    /* Backprop through FC1. */
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            const float dz = scratch_d_fc1[(size_t)t * (size_t)fourD + (size_t)j];
            if (d_b_fc1) d_b_fc1[j] += dz;
            for (int i = 0; i < D; ++i) {
                const float x = bf16_to_float(input[(size_t)t * (size_t)D + (size_t)i]);
                const float w1 = bf16_to_float(W_fc1[(size_t)j * (size_t)D + (size_t)i]);
                if (d_W_fc1) d_W_fc1[(size_t)j * (size_t)D + (size_t)i] += dz * x;
                if (d_input) d_input[(size_t)t * (size_t)D + (size_t)i] += dz * w1;
            }
        }
    }
}
 
#pragma GCC diagnostic pop