swiglu_kernels_bf16.c File Reference

SwiGLU activation kernels for BF16 tensors. More...

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

Go to the source code of this file.

Functions
void	swiglu_backward_bf16 (const uint16_t *input, const uint16_t *d_output, uint16_t *d_input, int tokens, int dim)
void	swiglu_forward_bf16 (const uint16_t *input, uint16_t *output, int tokens, int dim)

Detailed Description

SwiGLU activation kernels for BF16 tensors.

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

SwiGLU: y = silu(gate) * up = (gate * sigmoid(gate)) * up

Definition in file swiglu_kernels_bf16.c.

Function Documentation

swiglu_backward_bf16()

void swiglu_backward_bf16	(	const uint16_t *	input,
		const uint16_t *	d_output,
		uint16_t *	d_input,
		int	tokens,
		int	dim
	)

Definition at line 108 of file swiglu_kernels_bf16.c.

{
    if (!input || !d_output || !d_input || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const int T = tokens;
    const int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const uint16_t *row = input + (size_t)t * (size_t)(2 * D);
        const uint16_t *dy_row = d_output + (size_t)t * (size_t)D;
        uint16_t *dx_row = d_input + (size_t)t * (size_t)(2 * D);
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        __m512 one = _mm512_set1_ps(1.0f);
        for (; d + 16 <= D; d += 16) {
            __m512 a = bf16_loadu_cvt_fp32(&row[d]);         // gate
            __m512 b = bf16_loadu_cvt_fp32(&row[D + d]);     // value
            __m512 dy = bf16_loadu_cvt_fp32(&dy_row[d]);
 
            __m512 s = sigmoid512_fast_bf16(a);             // sigmoid(a)
            __m512 silu = _mm512_mul_ps(a, s);              // silu(a) = a * s
            __m512 s_prime = _mm512_mul_ps(s, _mm512_sub_ps(one, s)); // s * (1 - s)
            __m512 silu_prime = _mm512_fmadd_ps(a, s_prime, s);       // s + a * s_prime
 
            // dA = dy * b * silu_prime
            __m512 dA = _mm512_mul_ps(dy, _mm512_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m512 dB = _mm512_mul_ps(dy, silu);
 
            fp32_cvt_storeu_bf16(&dx_row[d], dA);
            fp32_cvt_storeu_bf16(&dx_row[D + d], dB);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = bf16_to_float(row[d]);
            float b = bf16_to_float(row[D + d]);
            float dy = bf16_to_float(dy_row[d]);
 
            float s = sigmoid_scalar(a);
            float silu = a * s;
            float s_prime = s * (1.0f - s);
            float silu_prime = s + a * s_prime;
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = float_to_bf16(dA);
            dx_row[D + d] = float_to_bf16(dB);
        }
    }
}

References bf16_to_float(), float_to_bf16(), sigmoid_scalar(), and silu().

swiglu_forward_bf16()

void swiglu_forward_bf16	(	const uint16_t *	input,
		uint16_t *	output,
		int	tokens,
		int	dim
	)

Definition at line 66 of file swiglu_kernels_bf16.c.

{
    if (!input || !output || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const int T = tokens;
    const int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const uint16_t *row = input + (size_t)t * (size_t)(2 * D);
        uint16_t *out_row = output + (size_t)t * (size_t)D;
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        for (; d + 16 <= D; d += 16) {
            __m512 a = bf16_loadu_cvt_fp32(&row[d]);         // gate
            __m512 b = bf16_loadu_cvt_fp32(&row[D + d]);     // value
 
            __m512 s = sigmoid512_fast_bf16(a);             // sigmoid(a)
            __m512 silu = _mm512_mul_ps(a, s);              // silu(a) = a * sigmoid(a)
            __m512 y = _mm512_mul_ps(silu, b);              // y = silu(a) * b
 
            fp32_cvt_storeu_bf16(&out_row[d], y);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = bf16_to_float(row[d]);
            float b = bf16_to_float(row[D + d]);
            float s = sigmoid_scalar(a);
            float silu = a * s;
            out_row[d] = float_to_bf16(silu * b);
        }
    }
}

References bf16_to_float(), float_to_bf16(), sigmoid_scalar(), and silu().