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

#include "bf16_utils.h"
#include "ckernel_engine.h"
#include "ckernel_quant.h"
#include <dlfcn.h>
#include <math.h>
#include <pthread.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>

Macros
#define	_GNU_SOURCE

Functions
static float	sigmoid_scalar_parity (float x)

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)

void	swiglu_forward_ggml (const float input, float output, int tokens, int dim)

void	swiglu_forward_ggml_split (const float gate, const float up, float *output, int tokens, int dim)

void	swiglu_forward_pytorch_bf16_storage (const float input, float output, int tokens, int dim)

void	swiglu_forward_q8_k (const float input, void output_q8, int tokens, int dim)

Detailed Description

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

CK-ENGINE KERNEL RULES:

NO malloc/free - memory via bump allocator, pointers passed in
NO OpenMP - parallelization at orchestrator/codegen layer
API must define: inputs, outputs, workspace, and memory layouts
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.

Macro Definition Documentation

◆ _GNU_SOURCE

#define _GNU_SOURCE

Definition at line 18 of file swiglu_kernels.c.

Function Documentation

◆ sigmoid_scalar_parity()

static float sigmoid_scalar_parity ( float x )

inlinestatic

Definition at line 39 of file swiglu_kernels.c.

{
    return 1.0f / (1.0f + expf(-x));
}

Referenced by swiglu_backward_exact(), swiglu_forward_exact(), and swiglu_forward_q8_k().

◆ 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 386 of file swiglu_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        swiglu_backward_exact(input, d_output, d_input, tokens, dim);
        return;
    }
 
    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 one_minus_s = _mm512_sub_ps(one, s);
            __m512 inner = _mm512_fmadd_ps(a, one_minus_s, one);      // 1 + a * (1 - s)
            __m512 silu_prime = _mm512_mul_ps(s, inner);              // s * (1 + a * (1 - s))
 
            // 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 one_minus_s = _mm256_sub_ps(one, s);
            __m256 inner = _mm256_fmadd_ps(a, one_minus_s, one);      // 1 + a * (1 - s)
            __m256 silu_prime = _mm256_mul_ps(s, inner);              // s * (1 + a * (1 - s))
 
            // 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 one_minus_s = _mm256_sub_ps(one, s);
            __m256 a_one_minus_s = _mm256_mul_ps(a, one_minus_s);
            __m256 inner = _mm256_add_ps(one, a_one_minus_s);         // 1 + a * (1 - s)
            __m256 silu_prime = _mm256_mul_ps(s, inner);              // s * (1 + a * (1 - s))
 
            // 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 silu_prime = s * (1.0f + a * (1.0f - s)); // silu'(a), PyTorch form
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = dA;
            dx_row[D + d] = dB;
        }
    }
}

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

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 719 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];
 
            float s = sigmoid_scalar_parity(a); // sigmoid(a)
            float silu = a * s;                 // silu(a)
            float silu_prime = s * (1.0f + a * (1.0f - s)); // silu'(a), PyTorch form
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = dA;
            dx_row[D + d] = dB;
        }
    }
}

References sigmoid_scalar_parity(), and silu().

Referenced by swiglu_backward().

◆ 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 237 of file swiglu_kernels.c.

{
    const char *fast_env = getenv("CK_SWIGLU_FAST");
    const char *exact_env = getenv("CK_SWIGLU_EXACT");
    if (ck_strict_parity_enabled() ||
        !(fast_env && atoi(fast_env) != 0) ||
        (exact_env && atoi(exact_env) != 0)) {
        swiglu_forward_exact(input, output, tokens, dim);
        return;
    }
 
    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__(), ck_strict_parity_enabled(), sigmoid_scalar(), silu(), and swiglu_forward_exact().

◆ 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 515 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
 
            float s = sigmoid_scalar_parity(a); // sigmoid(a)
            float silu = a * s;                 // silu(a)
            out_row[d] = silu * b;
        }
    }
}

References sigmoid_scalar_parity(), and silu().

Referenced by ck_mlp_swiglu_forward(), ck_mlp_swiglu_forward_ref(), and swiglu_forward().

◆ swiglu_forward_ggml()

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

Definition at line 538 of file swiglu_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *row = input + (size_t)t * (2 * dim);
        float *out_row = output + (size_t)t * dim;
        int d = 0;
 
#if defined(__AVX512F__) && defined(__AVX512DQ__)
        for (; d + 16 <= dim; d += 16) {
            const __m512 gate = _mm512_loadu_ps(row + d);
            const __m512 up = _mm512_loadu_ps(row + dim + d);
            const __m512 neg_gate = _mm512_sub_ps(_mm512_setzero_ps(), gate);
            const __m512 denom = _mm512_add_ps(
                _mm512_set1_ps(1.0f), ck_ggml_expf_avx512(neg_gate));
            const __m512 silu = _mm512_div_ps(gate, denom);
            _mm512_storeu_ps(out_row + d, _mm512_mul_ps(silu, up));
        }
#elif defined(__AVX2__) && defined(__FMA__)
        for (; d + 8 <= dim; d += 8) {
            const __m256 gate = _mm256_loadu_ps(row + d);
            const __m256 up = _mm256_loadu_ps(row + dim + d);
            const __m256 neg_gate = _mm256_sub_ps(_mm256_setzero_ps(), gate);
            const __m256 denom = _mm256_add_ps(
                _mm256_set1_ps(1.0f), ck_ggml_expf_avx2(neg_gate));
            const __m256 silu = _mm256_div_ps(gate, denom);
            _mm256_storeu_ps(out_row + d, _mm256_mul_ps(silu, up));
        }
#endif
        for (; d < dim; ++d) {
            const float gate = row[d];
            out_row[d] = (gate / (1.0f + expf(-gate))) * row[dim + d];
        }
    }
}

References silu().

Referenced by ck_moe_q4k_mixed_route_work(), ck_moe_q4k_q5k_bucket_work(), ck_moe_q4k_q5k_route_work(), moe_swiglu_expert_forward_q4k_q4k_workspace(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), moe_swiglu_expert_forward_q4k_q5k_workspace(), moe_swiglu_expert_forward_q4k_q6k_workspace(), moe_swiglu_expert_forward_q4k_q8_0_workspace(), moe_swiglu_shared_forward_q4k_q4k_workspace(), moe_swiglu_shared_forward_q4k_q6k_workspace(), and moe_swiglu_shared_forward_q8_0_gated_workspace().

◆ swiglu_forward_ggml_split()

void swiglu_forward_ggml_split	(	const float *	gate,
		const float *	up,
		float *	output,
		int	tokens,
		int	dim
	)

Definition at line 576 of file swiglu_kernels.c.

{
    if (!gate || !up || !output || tokens <= 0 || dim <= 0) {
        return;
    }
    for (int t = 0; t < tokens; ++t) {
        const float *gate_row = gate + (size_t)t * (size_t)dim;
        const float *up_row = up + (size_t)t * (size_t)dim;
        float *out_row = output + (size_t)t * (size_t)dim;
        int d = 0;
 
#if defined(__AVX512F__) && defined(__AVX512DQ__)
        for (; d + 16 <= dim; d += 16) {
            const __m512 gate_v = _mm512_loadu_ps(gate_row + d);
            const __m512 up_v = _mm512_loadu_ps(up_row + d);
            const __m512 neg_gate = _mm512_sub_ps(_mm512_setzero_ps(), gate_v);
            const __m512 denom = _mm512_add_ps(
                _mm512_set1_ps(1.0f), ck_ggml_expf_avx512(neg_gate));
            const __m512 silu = _mm512_div_ps(gate_v, denom);
            _mm512_storeu_ps(out_row + d, _mm512_mul_ps(silu, up_v));
        }
#elif defined(__AVX2__) && defined(__FMA__)
        for (; d + 8 <= dim; d += 8) {
            const __m256 gate_v = _mm256_loadu_ps(gate_row + d);
            const __m256 up_v = _mm256_loadu_ps(up_row + d);
            const __m256 neg_gate = _mm256_sub_ps(_mm256_setzero_ps(), gate_v);
            const __m256 denom = _mm256_add_ps(
                _mm256_set1_ps(1.0f), ck_ggml_expf_avx2(neg_gate));
            const __m256 silu = _mm256_div_ps(gate_v, denom);
            _mm256_storeu_ps(out_row + d, _mm256_mul_ps(silu, up_v));
        }
#endif
        for (; d < dim; ++d) {
            const float gate_v = gate_row[d];
            out_row[d] = (gate_v / (1.0f + expf(-gate_v))) * up_row[d];
        }
    }
}

References silu().

Referenced by ck_moe_q4k_q5k_bucket_work(), and ck_moe_shared_q4k_gated_workspace().

◆ swiglu_forward_pytorch_bf16_storage()

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

Definition at line 658 of file swiglu_kernels.c.

{
    if (!input || !output || tokens < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid PyTorch BF16 SwiGLU arguments\n");
        abort();
    }
 
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_swiglu_once, ck_bind_pytorch_swiglu_sleef);
    if (!ck_pytorch_swiglu_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16 SwiGLU requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
 
    for (int t = 0; t < tokens; ++t) {
        const float *row = input + (size_t)t * (size_t)(2 * dim);
        float *out_row = output + (size_t)t * (size_t)dim;
        int d = 0;
 
#if defined(__AVX512F__)
        for (; d + 16 <= dim; d += 16) {
            const __m512 gate = _mm512_loadu_ps(row + d);
            const __m512 denominator = _mm512_add_ps(
                _mm512_set1_ps(1.0f),
                ck_pytorch_swiglu_expf16(_mm512_sub_ps(_mm512_setzero_ps(), gate)));
            const __m512 silu = _mm512_div_ps(gate, denominator);
            float silu_lanes[16] __attribute__((aligned(64)));
            _mm512_store_ps(silu_lanes, silu);
            for (int lane = 0; lane < 16; ++lane) {
                const float silu_bf16 = bf16_to_float(float_to_bf16(silu_lanes[lane]));
                const float up_bf16 = bf16_to_float(float_to_bf16(row[dim + d + lane]));
                out_row[d + lane] = bf16_to_float(
                    float_to_bf16(silu_bf16 * up_bf16));
            }
        }
#endif
        for (; d < dim; ++d) {
            const float gate_bf16 = bf16_to_float(float_to_bf16(row[d]));
            const float up_bf16 = bf16_to_float(float_to_bf16(row[dim + d]));
            const float silu = gate_bf16 / (1.0f + expf(-gate_bf16));
            const float silu_bf16 = bf16_to_float(float_to_bf16(silu));
            out_row[d] = bf16_to_float(float_to_bf16(silu_bf16 * up_bf16));
        }
    }
}

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

Referenced by moe_swiglu_packed_expert_forward_bf16().

◆ swiglu_forward_q8_k()

void swiglu_forward_q8_k	(	const float *	input,
		void *	output_q8,
		int	tokens,
		int	dim
	)

Definition at line 319 of file swiglu_kernels.c.

{
    if (!input || !output_q8 || tokens <= 0 || dim <= 0) {
        return;
    }
    if ((dim % QK_K) != 0) {
        return;
    }
 
    const char *fast_env = getenv("CK_SWIGLU_FAST");
    const char *exact_env = getenv("CK_SWIGLU_EXACT");
    const int use_fast = !ck_strict_parity_enabled() &&
                         (fast_env && atoi(fast_env) != 0) &&
                         !(exact_env && atoi(exact_env) != 0);
 
    const int blocks_per_row = dim / QK_K;
    block_q8_K *q8 = (block_q8_K *)output_q8;
    float tmp[QK_K];
 
    for (int t = 0; t < tokens; ++t) {
        const float *row = input + (size_t)t * (size_t)(2 * dim);
        block_q8_K *q8_row = q8 + (size_t)t * (size_t)blocks_per_row;
 
        for (int block = 0; block < blocks_per_row; ++block) {
            const int base = block * QK_K;
            int d = 0;
 
#if defined(__AVX2__)
            if (use_fast) {
                for (; d + 8 <= QK_K; d += 8) {
                    const __m256 a = _mm256_loadu_ps(row + base + d);
                    const __m256 b = _mm256_loadu_ps(row + dim + base + d);
                    const __m256 s = sigmoid256_fast(a);
                    const __m256 y = _mm256_mul_ps(_mm256_mul_ps(a, s), b);
                    _mm256_storeu_ps(tmp + d, y);
                }
            }
#else
            (void)use_fast;
#endif
 
            for (; d < QK_K; ++d) {
                const float a = row[base + d];
                const float b = row[dim + base + d];
                const float s = use_fast ? sigmoid_scalar(a) : sigmoid_scalar_parity(a);
                tmp[d] = (a * s) * b;
            }
 
            quantize_row_q8_k(tmp, (void *)&q8_row[block], QK_K);
        }
    }
}

References ck_strict_parity_enabled(), QK_K, quantize_row_q8_k(), sigmoid_scalar(), and sigmoid_scalar_parity().

Macros

Functions

Detailed Description

CK-ENGINE KERNEL RULES:

Macro Definition Documentation

◆ _GNU_SOURCE

Function Documentation

◆ sigmoid_scalar_parity()

◆ swiglu_backward()

◆ swiglu_backward_exact()

◆ swiglu_forward()

◆ swiglu_forward_exact()

◆ swiglu_forward_ggml()

◆ swiglu_forward_ggml_split()

◆ swiglu_forward_pytorch_bf16_storage()

◆ swiglu_forward_q8_k()