swiglu_kernels.c File Reference

SwiGLU activation kernels with SIMD (SSE/AVX/AVX512) More...

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

Go to the source code of this file.

Functions
void	swiglu_backward (const float *input, const float *d_output, float *d_input, int tokens, int dim)
void	swiglu_backward_exact (const float *input, const float *d_output, float *d_input, int tokens, int dim)
void	swiglu_forward (const float *input, float *output, int tokens, int dim)
void	swiglu_forward_exact (const float *input, float *output, int tokens, int dim)

Detailed Description

SwiGLU activation kernels with SIMD (SSE/AVX/AVX512)

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

Function Documentation

swiglu_backward()

void swiglu_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		int	tokens,
		int	dim
	)

SwiGLU backward pass

Test:

test_swiglu.py::TestSwiGLUBackward::test_backward_tokens

test_swiglu.py::TestSwiGLUBackward::test_backward_single

test_parity.py::test_swiglu_backward_parity

Computes dGate and dUp given dY. dGate = dy * b * silu'(a), dUp = dy * silu(a)

After changes: make test && make llamacpp-parity-full

Definition at line 215 of file swiglu_kernels.c.

{
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        const float *dy_row = d_output + (size_t)t * D;
        float *dx_row = d_input + (size_t)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 = _mm512_loadu_ps(&row[d]);         // gate
            __m512 b = _mm512_loadu_ps(&row[D + d]);     // value
            __m512 dy = _mm512_loadu_ps(&dy_row[d]);
 
            __m512 s = sigmoid512_fast(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);
 
            _mm512_storeu_ps(&dx_row[d], dA);
            _mm512_storeu_ps(&dx_row[D + d], dB);
        }
#elif defined(__AVX2__)
        // AVX2: Process 8 floats at a time
        __m256 one = _mm256_set1_ps(1.0f);
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
            __m256 dy = _mm256_loadu_ps(&dy_row[d]);
 
            __m256 s = sigmoid256_fast(a);              // sigmoid(a)
            __m256 silu = _mm256_mul_ps(a, s);          // silu(a) = a * s
            __m256 s_prime = _mm256_mul_ps(s, _mm256_sub_ps(one, s)); // s * (1 - s)
            __m256 silu_prime = _mm256_fmadd_ps(a, s_prime, s);       // s + a * s_prime
 
            // dA = dy * b * silu_prime
            __m256 dA = _mm256_mul_ps(dy, _mm256_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m256 dB = _mm256_mul_ps(dy, silu);
 
            _mm256_storeu_ps(&dx_row[d], dA);
            _mm256_storeu_ps(&dx_row[D + d], dB);
        }
#elif defined(__AVX__)
        // AVX1: Vectorize arithmetic, use scalar sigmoid
        __m256 one = _mm256_set1_ps(1.0f);
        float a_arr[8] __attribute__((aligned(32)));
        float s_arr[8] __attribute__((aligned(32)));
 
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
            __m256 dy = _mm256_loadu_ps(&dy_row[d]);
 
            // Compute sigmoid scalarly
            _mm256_store_ps(a_arr, a);
            for (int j = 0; j < 8; ++j) {
                s_arr[j] = sigmoid_scalar(a_arr[j]);
            }
            __m256 s = _mm256_load_ps(s_arr);
 
            __m256 silu = _mm256_mul_ps(a, s);                        // silu(a) = a * s
            __m256 s_prime = _mm256_mul_ps(s, _mm256_sub_ps(one, s)); // s * (1 - s)
            // silu_prime = s + a * s_prime (no FMA in AVX1)
            __m256 a_s_prime = _mm256_mul_ps(a, s_prime);
            __m256 silu_prime = _mm256_add_ps(s, a_s_prime);
 
            // dA = dy * b * silu_prime
            __m256 dA = _mm256_mul_ps(dy, _mm256_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m256 dB = _mm256_mul_ps(dy, silu);
 
            _mm256_storeu_ps(&dx_row[d], dA);
            _mm256_storeu_ps(&dx_row[D + d], dB);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
            float dy = dy_row[d];
 
            float s = sigmoid_scalar(a);               // sigmoid(a)
            float silu = a * s;                       // silu(a)
            float s_prime = s * (1.0f - s);           // sigmoid'(a)
            float silu_prime = s + a * s_prime;       // silu'(a)
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = dA;
            dx_row[D + d] = dB;
        }
    }
}

References __attribute__(), sigmoid_scalar(), and silu().

Referenced by ck_layer_backward_rmsnorm_swiglu().

swiglu_backward_exact()

void swiglu_backward_exact	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		int	tokens,
		int	dim
	)

SwiGLU backward pass (exact version using stdlib sigmoid)

Test:

test_swiglu.py::TestSwiGLUBackward::test_exact_vs_fast

test_swiglu.py::TestSwiGLUBackward::test_exact_single

Uses standard library expf for numerical accuracy reference.

After changes: make test

Definition at line 373 of file swiglu_kernels.c.

{
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        const float *dy_row = d_output + (size_t)t * D;
        float *dx_row = d_input + (size_t)t * (2 * D);
 
        for (int d = 0; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
            float dy = dy_row[d];
 
            // Use standard library expf via sigmoid_scalar
            float s = sigmoid_scalar(a);               // sigmoid(a)
            float silu = a * s;                       // silu(a)
            float s_prime = s * (1.0f - s);           // sigmoid'(a)
            float silu_prime = s + a * s_prime;       // silu'(a)
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = dA;
            dx_row[D + d] = dB;
        }
    }
}

References sigmoid_scalar(), and silu().

swiglu_forward()

void swiglu_forward	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim
	)

SwiGLU forward pass

Test:

test_swiglu.py::TestSwiGLUForward::test_forward_tokens

test_swiglu.py::TestSwiGLUForward::test_forward_single

test_mlp.py::TestMLPForward::test_swiglu_mlp

test_fused_swiglu_decode.py::TestFusedSwiGLUDecode::test_fused_swiglu_decode

test_parity.py::test_swiglu_parity

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

After changes: make test && make llamacpp-parity-full

Definition at line 131 of file swiglu_kernels.c.

{
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        float *out_row = output + (size_t)t * D;
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        for (; d + 16 <= D; d += 16) {
            __m512 a = _mm512_loadu_ps(&row[d]);         // gate
            __m512 b = _mm512_loadu_ps(&row[D + d]);     // value
 
            __m512 s = sigmoid512_fast(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
 
            _mm512_storeu_ps(&out_row[d], y);
        }
#elif defined(__AVX2__)
        // AVX2: Process 8 floats at a time
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
 
            __m256 s = sigmoid256_fast(a);              // sigmoid(a)
            __m256 silu = _mm256_mul_ps(a, s);          // silu(a) = a * sigmoid(a)
            __m256 y = _mm256_mul_ps(silu, b);          // y = silu(a) * b
 
            _mm256_storeu_ps(&out_row[d], y);
        }
#elif defined(__AVX__)
        // AVX1: Vectorize arithmetic, use scalar sigmoid
        float a_arr[8] __attribute__((aligned(32)));
        float s_arr[8] __attribute__((aligned(32)));
 
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
 
            // Compute sigmoid scalarly
            _mm256_store_ps(a_arr, a);
            for (int j = 0; j < 8; ++j) {
                s_arr[j] = sigmoid_scalar(a_arr[j]);
            }
            __m256 s = _mm256_load_ps(s_arr);
 
            __m256 silu = _mm256_mul_ps(a, s);          // silu(a) = a * sigmoid(a)
            __m256 y = _mm256_mul_ps(silu, b);          // y = silu(a) * b
 
            _mm256_storeu_ps(&out_row[d], y);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
 
            float s = sigmoid_scalar(a);         // sigmoid(a)
            float silu = a * s;                  // silu(a) = a * sigmoid(a)
 
            out_row[d] = silu * b;
        }
    }
}

References __attribute__(), sigmoid_scalar(), and silu().

Referenced by ck_mlp_swiglu_forward(), ck_mlp_swiglu_forward_q4_k(), ck_mlp_swiglu_forward_q4_k_q8_k(), ck_mlp_swiglu_forward_q4_k_q8_k_prefill(), ck_mlp_swiglu_forward_quant(), ck_mlp_swiglu_forward_ref(), ck_test_swiglu(), model_layer_0_decode(), model_layer_0_prefill(), model_layer_10_decode(), model_layer_10_prefill(), model_layer_11_decode(), model_layer_11_prefill(), model_layer_12_decode(), model_layer_12_prefill(), model_layer_13_decode(), model_layer_13_prefill(), model_layer_14_decode(), model_layer_14_prefill(), model_layer_15_decode(), model_layer_15_prefill(), model_layer_16_decode(), model_layer_16_prefill(), model_layer_17_decode(), model_layer_17_prefill(), model_layer_18_decode(), model_layer_18_prefill(), model_layer_19_decode(), model_layer_19_prefill(), model_layer_1_decode(), model_layer_1_prefill(), model_layer_20_decode(), model_layer_20_prefill(), model_layer_21_decode(), model_layer_21_prefill(), model_layer_22_decode(), model_layer_22_prefill(), model_layer_23_decode(), model_layer_23_prefill(), model_layer_2_decode(), model_layer_2_prefill(), model_layer_3_decode(), model_layer_3_prefill(), model_layer_4_decode(), model_layer_4_prefill(), model_layer_5_decode(), model_layer_5_prefill(), model_layer_6_decode(), model_layer_6_prefill(), model_layer_7_decode(), model_layer_7_prefill(), model_layer_8_decode(), model_layer_8_prefill(), model_layer_9_decode(), model_layer_9_prefill(), qwen2_0_5b_decode_layer_0_decode(), qwen2_0_5b_decode_layer_0_prefill(), qwen2_0_5b_decode_layer_10_decode(), qwen2_0_5b_decode_layer_10_prefill(), qwen2_0_5b_decode_layer_11_decode(), qwen2_0_5b_decode_layer_11_prefill(), qwen2_0_5b_decode_layer_12_decode(), qwen2_0_5b_decode_layer_12_prefill(), qwen2_0_5b_decode_layer_13_decode(), qwen2_0_5b_decode_layer_13_prefill(), qwen2_0_5b_decode_layer_14_decode(), qwen2_0_5b_decode_layer_14_prefill(), qwen2_0_5b_decode_layer_15_decode(), qwen2_0_5b_decode_layer_15_prefill(), qwen2_0_5b_decode_layer_16_decode(), qwen2_0_5b_decode_layer_16_prefill(), qwen2_0_5b_decode_layer_17_decode(), qwen2_0_5b_decode_layer_17_prefill(), qwen2_0_5b_decode_layer_18_decode(), qwen2_0_5b_decode_layer_18_prefill(), qwen2_0_5b_decode_layer_19_decode(), qwen2_0_5b_decode_layer_19_prefill(), qwen2_0_5b_decode_layer_1_decode(), qwen2_0_5b_decode_layer_1_prefill(), qwen2_0_5b_decode_layer_20_decode(), qwen2_0_5b_decode_layer_20_prefill(), qwen2_0_5b_decode_layer_21_decode(), qwen2_0_5b_decode_layer_21_prefill(), qwen2_0_5b_decode_layer_22_decode(), qwen2_0_5b_decode_layer_22_prefill(), qwen2_0_5b_decode_layer_23_decode(), qwen2_0_5b_decode_layer_23_prefill(), qwen2_0_5b_decode_layer_2_decode(), qwen2_0_5b_decode_layer_2_prefill(), qwen2_0_5b_decode_layer_3_decode(), qwen2_0_5b_decode_layer_3_prefill(), qwen2_0_5b_decode_layer_4_decode(), qwen2_0_5b_decode_layer_4_prefill(), qwen2_0_5b_decode_layer_5_decode(), qwen2_0_5b_decode_layer_5_prefill(), qwen2_0_5b_decode_layer_6_decode(), qwen2_0_5b_decode_layer_6_prefill(), qwen2_0_5b_decode_layer_7_decode(), qwen2_0_5b_decode_layer_7_prefill(), qwen2_0_5b_decode_layer_8_decode(), qwen2_0_5b_decode_layer_8_prefill(), qwen2_0_5b_decode_layer_9_decode(), and qwen2_0_5b_decode_layer_9_prefill().

swiglu_forward_exact()

void swiglu_forward_exact	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim
	)

SwiGLU forward pass (exact version using stdlib sigmoid)

Test:

test_swiglu.py::TestSwiGLUForward::test_exact_vs_fast

test_swiglu.py::TestSwiGLUForward::test_exact_single

Uses standard library expf for numerical accuracy reference.

After changes: make test

Definition at line 339 of file swiglu_kernels.c.

{
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        float *out_row = output + (size_t)t * D;
 
        for (int d = 0; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
 
            // Use standard library expf via sigmoid_scalar
            float s = sigmoid_scalar(a);         // sigmoid(a) = 1/(1+expf(-a))
            float silu = a * s;                  // silu(a) = a * sigmoid(a)
 
            out_row[d] = silu * b;
        }
    }
}

References sigmoid_scalar(), and silu().