mlp_kernels_bf16.c File Reference

Optimized BF16 MLP Kernels. More...

#include <stddef.h>
#include <stdint.h>
#include <math.h>
#include "bf16_utils.h"
#include "ckernel_engine.h"

Go to the source code of this file.

Functions
static float	gelu_derivative_scalar (float x)
static float	gelu_scalar (float x)
void	gemm_bf16_fp32out (const uint16_t *A, const uint16_t *B, const float *bias, float *C, int M, int N, int K)
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)
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)
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)

Detailed Description

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).

Definition in file mlp_kernels_bf16.c.

Function Documentation

gelu_derivative_scalar()

static float gelu_derivative_scalar ( float  x )

inlinestatic

Definition at line 53 of file mlp_kernels_bf16.c.

{
    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;
}

Referenced by mlp_token_parallel_bf16_backward_mixed().

gelu_scalar()

static float gelu_scalar ( float  x )

inlinestatic

Definition at line 45 of file mlp_kernels_bf16.c.

{
    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)));
}

Referenced by mlp_token_parallel_bf16(), mlp_token_parallel_bf16_backward_mixed(), and mlp_token_parallel_bf16_fp32act().

gemm_bf16_fp32out()

void gemm_bf16_fp32out ( const uint16_t *  A, const uint16_t *  B, const float *  bias, float *  C, int  M, int  N, int  K  )

extern

Definition at line 453 of file gemm_kernels_bf16.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
#if HAVE_NATIVE_BF16
#if HAVE_AMX_BF16
    const char *amx_env = getenv("CK_BF16_AMX");
    if (amx_env && amx_env[0] == '1' &&
        (M % 16) == 0 && (N % 16) == 0 && (K % 32) == 0 &&
        M >= 16 && N >= 16 && K >= 32 && ck_amx_request_xtile_data()) {
        gemm_bf16_fp32out_amx(A, B, bias, C, M, N, K);
        return;
    }
#endif
 
    #pragma omp parallel for schedule(dynamic)
    for (int i = 0; i < M; ++i) {
        const uint16_t *a_row = A + (size_t)i * K;
        int j = 0;
 
        for (; j + 4 <= N; j += 4) {
            const uint16_t *b0 = B + (size_t)(j + 0) * K;
            const uint16_t *b1 = B + (size_t)(j + 1) * K;
            const uint16_t *b2 = B + (size_t)(j + 2) * K;
            const uint16_t *b3 = B + (size_t)(j + 3) * K;
            __m512 acc0 = _mm512_setzero_ps();
            __m512 acc1 = _mm512_setzero_ps();
            __m512 acc2 = _mm512_setzero_ps();
            __m512 acc3 = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 32; k += 32) {
                const __m512bh a_vec = load_bf16x32(a_row + k);
                acc0 = _mm512_dpbf16_ps(acc0, a_vec, load_bf16x32(b0 + k));
                acc1 = _mm512_dpbf16_ps(acc1, a_vec, load_bf16x32(b1 + k));
                acc2 = _mm512_dpbf16_ps(acc2, a_vec, load_bf16x32(b2 + k));
                acc3 = _mm512_dpbf16_ps(acc3, a_vec, load_bf16x32(b3 + k));
            }
 
            float s0 = _mm512_reduce_add_ps(acc0);
            float s1 = _mm512_reduce_add_ps(acc1);
            float s2 = _mm512_reduce_add_ps(acc2);
            float s3 = _mm512_reduce_add_ps(acc3);
            for (; k < K; ++k) {
                const float a = bf16_to_float(a_row[k]);
                s0 += a * bf16_to_float(b0[k]);
                s1 += a * bf16_to_float(b1[k]);
                s2 += a * bf16_to_float(b2[k]);
                s3 += a * bf16_to_float(b3[k]);
            }
            if (bias) {
                s0 += bias[j + 0];
                s1 += bias[j + 1];
                s2 += bias[j + 2];
                s3 += bias[j + 3];
            }
            C[(size_t)i * N + (j + 0)] = s0;
            C[(size_t)i * N + (j + 1)] = s1;
            C[(size_t)i * N + (j + 2)] = s2;
            C[(size_t)i * N + (j + 3)] = s3;
        }
 
        for (; j < N; ++j) {
            const uint16_t *b_row = B + (size_t)j * K;
            __m512 sum_vec = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 32; k += 32) {
                const __m512bh a_vec = load_bf16x32(a_row + k);
                const __m512bh b_vec = load_bf16x32(b_row + k);
                sum_vec = _mm512_dpbf16_ps(sum_vec, a_vec, b_vec);
            }
 
            float sum = _mm512_reduce_add_ps(sum_vec);
            for (; k < K; ++k) {
                sum += bf16_to_float(a_row[k]) * bf16_to_float(b_row[k]);
            }
            if (bias) {
                sum += bias[j];
            }
            C[(size_t)i * N + j] = sum;
        }
    }
#elif defined(__AVX512F__)
    #pragma omp parallel for schedule(dynamic)
    for (int i = 0; i < M; ++i) {
        const uint16_t *a_row = A + (size_t)i * K;
 
        for (int j = 0; j < N; ++j) {
            const uint16_t *b_row = B + (size_t)j * K;
 
            __m512 sum_vec = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= K - 16; k += 16) {
                __m256i a_bf16 = _mm256_loadu_si256((const __m256i *)(a_row + k));
                __m256i b_bf16 = _mm256_loadu_si256((const __m256i *)(b_row + k));
                sum_vec = bf16_dot16(a_bf16, b_bf16, sum_vec);
            }
 
            float sum = _mm512_reduce_add_ps(sum_vec);
 
            for (; k < K; ++k) {
                sum += bf16_to_float(a_row[k]) * bf16_to_float(b_row[k]);
            }
 
            if (bias) {
                sum += bias[j];
            }
 
            C[(size_t)i * N + j] = sum;
        }
    }
#else
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += bf16_to_float(A[(size_t)i * K + k]) *
                       bf16_to_float(B[(size_t)j * K + k]);
            }
            C[(size_t)i * N + j] = sum;
        }
    }
#endif
}

References bf16_to_float(), and C.

Referenced by mlp_token_parallel_bf16(), and mlp_token_parallel_bf16_fp32act().

mlp_token_parallel_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
	)

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)

Definition at line 104 of file mlp_kernels_bf16.c.

{
    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);
}

References bf16_to_float(), float_to_bf16(), gelu_scalar(), and gemm_bf16_fp32out().

mlp_token_parallel_bf16_backward_mixed()

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
	)

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.

Definition at line 268 of file mlp_kernels_bf16.c.

{
    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;
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), gelu_derivative_scalar(), and gelu_scalar().

mlp_token_parallel_bf16_fp32act()

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
	)

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)

Definition at line 199 of file mlp_kernels_bf16.c.

{
    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);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), gelu_scalar(), and gemm_bf16_fp32out().