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

#include "bf16_utils.h"
#include "ck_threadpool.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 void	ck_swiglu_dispatch_rows (ck_threadpool_t pool, int active, int tokens, int dim, ck_range_fn_t fn, void args, const float input, float output)

static void	ck_swiglu_exact_rows (int begin, int end, void *opaque)

const char *	ck_swiglu_ggml_provider_name (void)

static void	ck_swiglu_ggml_rows (int begin, int end, void *opaque)

static int	ck_swiglu_ranges_overlap (const float input, const float output, int tokens, int dim)

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

◆ ck_swiglu_dispatch_rows()

static void ck_swiglu_dispatch_rows	(	ck_threadpool_t *	pool,
		int	active,
		int	tokens,
		int	dim,
		ck_range_fn_t	fn,
		void *	args,
		const float *	input,
		float *	output
	)

static

Definition at line 581 of file swiglu_kernels.c.

{
    const size_t minimum_elements_per_thread = 65536u;
    const size_t elems = (size_t)tokens * (size_t)dim;
    if (!pool || active <= 1 || tokens < 2 ||
        elems < 2u * minimum_elements_per_thread ||
        ck_threadpool_thread_id(pool) > 0) {
        fn(0, tokens, args);
        return;
    }
 
    const size_t useful_threads =
        (elems + minimum_elements_per_thread - 1u) /
        minimum_elements_per_thread;
    if ((size_t)active > useful_threads) active = (int)useful_threads;
    if (active > tokens) active = tokens;
 
    if (!ck_swiglu_ranges_overlap(input, output, tokens, dim)) {
        ck_threadpool_parallel_for_n(
            pool, active, 0, tokens, 1, fn, args);
        return;
    }
    if (input != output) {
        fn(0, tokens, args);
        return;
    }
 
    /*
     * In-place [T, 2D] -> [T, D] compaction has a binary dependency:
     * output row t overlaps input row floor(t / 2). Process each parent
     * wave before its children; rows within a wave remain independent.
     */
    fn(0, 1, args);
    for (int begin = 1; begin < tokens;) {
        int end = begin > tokens - begin ? tokens : begin * 2;
        int wave_threads = active;
        const int wave_rows = end - begin;
        if (wave_threads > wave_rows) wave_threads = wave_rows;
        ck_threadpool_parallel_for_n(
            pool, wave_threads, begin, end, 1, fn, args);
        begin = end;
    }
}

References ck_swiglu_ranges_overlap(), ck_threadpool_parallel_for_n(), ck_threadpool_thread_id(), and end.

Referenced by swiglu_forward_exact(), and swiglu_forward_ggml().

◆ ck_swiglu_exact_rows()

static void ck_swiglu_exact_rows	(	int	begin,
		int	end,
		void *	opaque
	)

static

Definition at line 632 of file swiglu_kernels.c.

{
    const ck_swiglu_exact_args_t *args =
        (const ck_swiglu_exact_args_t *)opaque;
 
    for (int t = begin; t < end; ++t) {
        const float *row = args->input +
            (size_t)t * (size_t)(2 * args->dim);
        float *out_row = args->output + (size_t)t * (size_t)args->dim;
 
        for (int d = 0; d < args->dim; ++d) {
            float a = row[d];
            float b = row[args->dim + d];
            float s = sigmoid_scalar_parity(a);
            float silu = a * s;
            out_row[d] = silu * b;
        }
    }
}

References end, sigmoid_scalar_parity(), and silu().

Referenced by swiglu_forward_exact().

◆ ck_swiglu_ggml_provider_name()

const char * ck_swiglu_ggml_provider_name ( void )

Definition at line 804 of file swiglu_kernels.c.

{
#if defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
    if (ck_swiglu_ggml_ref_forced()) {
        return "swiglu_ggml_ref";
    }
#endif
#if defined(__AVX512F__) && defined(__AVX512DQ__)
    return "swiglu_ggml_avx512";
#elif defined(__AVX2__) && defined(__FMA__)
    return "swiglu_ggml_avx2";
#elif defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
    return "swiglu_ggml_neon";
#else
    return "swiglu_ggml_ref";
#endif
}

◆ ck_swiglu_ggml_rows()

static void ck_swiglu_ggml_rows	(	int	begin,
		int	end,
		void *	opaque
	)

static

Definition at line 674 of file swiglu_kernels.c.

{
    const ck_swiglu_ggml_args_t *args =
        (const ck_swiglu_ggml_args_t *)opaque;
    const int dim = args->dim;
#if defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
    const int force_ref = ck_swiglu_ggml_ref_forced();
#endif
    for (int t = begin; t < end; ++t) {
        const float *row = args->input + (size_t)t * (size_t)(2 * dim);
        float *out_row = args->output + (size_t)t * (size_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));
        }
#elif defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
        if (!force_ref) {
            for (; d + 4 <= dim; d += 4) {
                const float32x4_t gate = vld1q_f32(row + d);
                const float32x4_t up = vld1q_f32(row + dim + d);
                const float32x4_t neg_gate = vsubq_f32(vdupq_n_f32(0.0f), gate);
                const float32x4_t denom = vaddq_f32(
                    vdupq_n_f32(1.0f), ck_ggml_expf_neon(neg_gate));
                const float32x4_t silu = vdivq_f32(gate, denom);
                vst1q_f32(out_row + d, vmulq_f32(silu, up));
            }
        }
#endif
        for (; d < dim; ++d) {
            const float gate = row[d];
            out_row[d] = (gate / (1.0f + expf(-gate))) * row[dim + d];
        }
    }
}

References end, and silu().

Referenced by swiglu_forward_ggml().

◆ ck_swiglu_ranges_overlap()

static int ck_swiglu_ranges_overlap	(	const float *	input,
		const float *	output,
		int	tokens,
		int	dim
	)

static

Definition at line 567 of file swiglu_kernels.c.

{
    const uintptr_t input_begin = (uintptr_t)input;
    const uintptr_t input_end = input_begin +
        (size_t)tokens * (size_t)(2 * dim) * sizeof(float);
    const uintptr_t output_begin = (uintptr_t)output;
    const uintptr_t output_end = output_begin +
        (size_t)tokens * (size_t)dim * sizeof(float);
    return input_begin < output_end && output_begin < input_end;
}

Referenced by ck_swiglu_dispatch_rows().

◆ sigmoid_scalar_parity()

static float sigmoid_scalar_parity ( float x )

inlinestatic

Definition at line 44 of file swiglu_kernels.c.

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

Referenced by ck_swiglu_exact_rows(), swiglu_backward_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 432 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 922 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 283 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
	)

Definition at line 652 of file swiglu_kernels.c.

{
    if (!input || !output || tokens <= 0 || dim <= 0) {
        return;
    }
 
    ck_swiglu_exact_args_t args = {input, output, dim};
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    ck_swiglu_dispatch_rows(pool, active, tokens, dim,
                            ck_swiglu_exact_rows, &args, input, output);
}

References ck_swiglu_dispatch_rows(), ck_swiglu_exact_rows(), ck_threadpool_global(), and ck_threadpool_n_threads().

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

{
    if (!input || !output || tokens <= 0 || dim <= 0) {
        return;
    }
 
    ck_swiglu_ggml_args_t args = {input, output, dim};
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    ck_swiglu_dispatch_rows(pool, active, tokens, dim,
                            ck_swiglu_ggml_rows, &args, input, output);
}

References ck_swiglu_dispatch_rows(), ck_swiglu_ggml_rows(), ck_threadpool_global(), and ck_threadpool_n_threads().

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

{
    if (!gate || !up || !output || tokens <= 0 || dim <= 0) {
        return;
    }
#if defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
    const int force_ref = ck_swiglu_ggml_ref_forced();
#endif
    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));
        }
#elif defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
        if (!force_ref) {
            for (; d + 4 <= dim; d += 4) {
                const float32x4_t gate_v = vld1q_f32(gate_row + d);
                const float32x4_t up_v = vld1q_f32(up_row + d);
                const float32x4_t neg_gate = vsubq_f32(vdupq_n_f32(0.0f), gate_v);
                const float32x4_t denom = vaddq_f32(
                    vdupq_n_f32(1.0f), ck_ggml_expf_neon(neg_gate));
                const float32x4_t silu = vdivq_f32(gate_v, denom);
                vst1q_f32(out_row + d, vmulq_f32(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 861 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 365 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

◆ ck_swiglu_dispatch_rows()

◆ ck_swiglu_exact_rows()

◆ ck_swiglu_ggml_provider_name()

◆ ck_swiglu_ggml_rows()

◆ ck_swiglu_ranges_overlap()

◆ 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()