swiglu_kernels.c

Go to the documentation of this file.

/**
 * @file swiglu_kernels.c
 * @brief 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
 */
 
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif
 
#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>
 
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
#include <immintrin.h>
#endif
 
#if defined(__ARM_NEON) || defined(__aarch64__)
#include <arm_neon.h>
#endif
 
/*
 * PyTorch-parity sigmoid for strict/exact SwiGLU paths.
 * Keep this in fp32 to match ATen CPU opmath more closely.
 */
static inline float sigmoid_scalar_parity(float x)
{
    return 1.0f / (1.0f + expf(-x));
}
 
/* ========================================================================== */
/* Fast exp approximation for SIMD                                            */
/* ========================================================================== */
 
#if defined(__AVX512F__)
// AVX-512 fast exp approximation
static inline __m512 exp512_fast(__m512 x) {
    // Clamp to avoid overflow/underflow
    x = _mm512_max_ps(x, _mm512_set1_ps(-88.0f));
    x = _mm512_min_ps(x, _mm512_set1_ps(88.0f));
 
    // exp(x) = 2^(x * log2(e))
    const __m512 log2e = _mm512_set1_ps(1.4426950408889634f);
    __m512 z = _mm512_mul_ps(x, log2e);
 
    // Split into integer and fractional parts
    __m512 zf = _mm512_roundscale_ps(z, _MM_FROUND_TO_NEAREST_INT);
    __m512 f = _mm512_sub_ps(z, zf);
 
    // Polynomial for 2^f, f in [-0.5, 0.5]
    const __m512 c0 = _mm512_set1_ps(1.0f);
    const __m512 c1 = _mm512_set1_ps(0.6931471805599453f);
    const __m512 c2 = _mm512_set1_ps(0.2402265069591007f);
    const __m512 c3 = _mm512_set1_ps(0.05550410866482158f);
    const __m512 c4 = _mm512_set1_ps(0.009618129107628478f);
 
    __m512 poly = _mm512_fmadd_ps(f, c4, c3);
    poly = _mm512_fmadd_ps(f, poly, c2);
    poly = _mm512_fmadd_ps(f, poly, c1);
    poly = _mm512_fmadd_ps(f, poly, c0);
 
    // Scale by 2^n
    __m512i zi = _mm512_cvtps_epi32(zf);
    zi = _mm512_add_epi32(zi, _mm512_set1_epi32(127));
    zi = _mm512_slli_epi32(zi, 23);
    __m512 scale = _mm512_castsi512_ps(zi);
 
    return _mm512_mul_ps(poly, scale);
}
 
// AVX-512 sigmoid: 1 / (1 + exp(-x))
static inline __m512 sigmoid512_fast(__m512 x) {
    __m512 neg_x = _mm512_sub_ps(_mm512_setzero_ps(), x);
    __m512 exp_neg = exp512_fast(neg_x);
    __m512 one = _mm512_set1_ps(1.0f);
    return _mm512_div_ps(one, _mm512_add_ps(one, exp_neg));
}
#endif
 
#if defined(__AVX2__)
// AVX2 fast exp approximation (needs FMA and integer ops)
static inline __m256 exp256_fast(__m256 x) {
    // Clamp
    x = _mm256_max_ps(x, _mm256_set1_ps(-88.0f));
    x = _mm256_min_ps(x, _mm256_set1_ps(88.0f));
 
    // exp(x) = 2^(x * log2(e))
    const __m256 log2e = _mm256_set1_ps(1.4426950408889634f);
    __m256 z = _mm256_mul_ps(x, log2e);
 
    // Round to nearest integer
    __m256 zf = _mm256_round_ps(z, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
    __m256 f = _mm256_sub_ps(z, zf);
 
    // Polynomial for 2^f
    const __m256 c0 = _mm256_set1_ps(1.0f);
    const __m256 c1 = _mm256_set1_ps(0.6931471805599453f);
    const __m256 c2 = _mm256_set1_ps(0.2402265069591007f);
    const __m256 c3 = _mm256_set1_ps(0.05550410866482158f);
    const __m256 c4 = _mm256_set1_ps(0.009618129107628478f);
 
    __m256 poly = _mm256_fmadd_ps(f, c4, c3);
    poly = _mm256_fmadd_ps(f, poly, c2);
    poly = _mm256_fmadd_ps(f, poly, c1);
    poly = _mm256_fmadd_ps(f, poly, c0);
 
    // Scale by 2^n
    __m256i zi = _mm256_cvtps_epi32(zf);
    zi = _mm256_add_epi32(zi, _mm256_set1_epi32(127));
    zi = _mm256_slli_epi32(zi, 23);
    __m256 scale = _mm256_castsi256_ps(zi);
 
    return _mm256_mul_ps(poly, scale);
}
 
// AVX2 sigmoid
static inline __m256 sigmoid256_fast(__m256 x) {
    __m256 neg_x = _mm256_sub_ps(_mm256_setzero_ps(), x);
    __m256 exp_neg = exp256_fast(neg_x);
    __m256 one = _mm256_set1_ps(1.0f);
    return _mm256_div_ps(one, _mm256_add_ps(one, exp_neg));
}
 
/* GGML's parity path uses this specific vector exponential approximation. */
static inline __m256 ck_ggml_expf_avx2(__m256 x) {
    const __m256 r = _mm256_set1_ps(0x1.8p23f);
    const __m256 z = _mm256_fmadd_ps(x, _mm256_set1_ps(0x1.715476p+0f), r);
    const __m256 n = _mm256_sub_ps(z, r);
    const __m256 b = _mm256_fnmadd_ps(
        n,
        _mm256_set1_ps(0x1.7f7d1cp-20f),
        _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.62e4p-1f), x));
    const __m256i e = _mm256_slli_epi32(_mm256_castps_si256(z), 23);
    const __m256 k = _mm256_castsi256_ps(
        _mm256_add_epi32(e, _mm256_castps_si256(_mm256_set1_ps(1))));
    const __m256i c = _mm256_castps_si256(
        _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.0f), n),
                      _mm256_set1_ps(126), _CMP_GT_OQ));
    const __m256 u = _mm256_mul_ps(b, b);
    const __m256 j = _mm256_fmadd_ps(
        _mm256_fmadd_ps(
            _mm256_fmadd_ps(_mm256_set1_ps(0x1.0e4020p-7f), b,
                            _mm256_set1_ps(0x1.573e2ep-5f)),
            u,
            _mm256_fmadd_ps(_mm256_set1_ps(0x1.555e66p-3f), b,
                            _mm256_set1_ps(0x1.fffdb6p-2f))),
        u,
        _mm256_mul_ps(_mm256_set1_ps(0x1.ffffecp-1f), b));
    if (!_mm256_movemask_ps(_mm256_castsi256_ps(c))) {
        return _mm256_fmadd_ps(j, k, k);
    }
    const __m256i g = _mm256_and_si256(
        _mm256_castps_si256(_mm256_cmp_ps(n, _mm256_setzero_ps(), _CMP_LE_OQ)),
        _mm256_set1_epi32((int)0x82000000u));
    const __m256 s1 = _mm256_castsi256_ps(
        _mm256_add_epi32(g, _mm256_set1_epi32(0x7f000000u)));
    const __m256 s2 = _mm256_castsi256_ps(_mm256_sub_epi32(e, g));
    const __m256i d = _mm256_castps_si256(
        _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.0f), n),
                      _mm256_set1_ps(192), _CMP_GT_OQ));
    return _mm256_or_ps(
        _mm256_and_ps(_mm256_castsi256_ps(d), _mm256_mul_ps(s1, s1)),
        _mm256_andnot_ps(
            _mm256_castsi256_ps(d),
            _mm256_or_ps(
                _mm256_and_ps(_mm256_castsi256_ps(c),
                              _mm256_mul_ps(_mm256_fmadd_ps(s2, j, s2), s1)),
                _mm256_andnot_ps(_mm256_castsi256_ps(c),
                                 _mm256_fmadd_ps(k, j, k)))));
}
#endif
 
#if defined(__AVX512F__) && defined(__AVX512DQ__)
/* Keep the production parity provider aligned with llama.cpp's AVX-512
 * exponential approximation and instruction grouping. */
static inline __m512 ck_ggml_expf_avx512(__m512 x) {
    const __m512 r = _mm512_set1_ps(0x1.8p23f);
    const __m512 z = _mm512_fmadd_ps(x, _mm512_set1_ps(0x1.715476p+0f), r);
    const __m512 n = _mm512_sub_ps(z, r);
    const __m512 b = _mm512_fnmadd_ps(
        n, _mm512_set1_ps(0x1.7f7d1cp-20f),
        _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.62e4p-1f), x));
    const __mmask16 d = _mm512_cmp_ps_mask(
        _mm512_abs_ps(n), _mm512_set1_ps(192.0f), _CMP_GT_OQ);
    const __m512 u = _mm512_mul_ps(b, b);
    const __m512 j = _mm512_fmadd_ps(
        _mm512_fmadd_ps(
            _mm512_fmadd_ps(
                _mm512_set1_ps(0x1.0e4020p-7f), b,
                _mm512_set1_ps(0x1.573e2ep-5f)),
            u,
            _mm512_fmadd_ps(
                _mm512_set1_ps(0x1.555e66p-3f), b,
                _mm512_set1_ps(0x1.fffdb6p-2f))),
        u,
        _mm512_fmadd_ps(
            _mm512_set1_ps(0x1.ffffecp-1f), b,
            _mm512_set1_ps(1.0f)));
    const __m512 res = _mm512_scalef_ps(j, n);
    if (_mm512_kortestz(d, d)) {
        return res;
    }
    const __m512 zero = _mm512_setzero_ps();
    const __m512 alt = _mm512_mask_blend_ps(
        _mm512_cmp_ps_mask(n, zero, _CMP_LE_OQ),
        _mm512_set1_ps(INFINITY),
        zero);
    return _mm512_mask_blend_ps(d, res, alt);
}
#endif
 
#if defined(__ARM_NEON) && (defined(__aarch64__) || defined(__ARM_FEATURE_FMA))
/* NEON port of the same ggml vector exponential the AVX2 provider above
 * implements; this is llama.cpp's own AArch64 ggml_v_expf (arm limited
 * optimized-routines polynomial, max error 1.45358 + 0.5 ulps, flushes above
 * 88.38 / below -103.97). On AArch64 this is the instruction-level oracle
 * ggml itself uses, so the numerical contract is identical to the AVX2
 * branch: approximation, not libm expf. Uses vfmaq_f32 (fused), so it
 * requires FP FMA (baseline on AArch64); 32-bit NEON without FMA keeps the
 * scalar expf tail. The any-lane predicate uses vgetq_lane_u64 instead of
 * the AArch64-only vpaddd_u64 so the source stays portable NEON. */
static inline float32x4_t ck_ggml_expf_neon(float32x4_t x) {
    const float32x4_t r = vdupq_n_f32(0x1.8p23f);
    const float32x4_t z = vfmaq_f32(r, x, vdupq_n_f32(0x1.715476p+0f));
    const float32x4_t n = vsubq_f32(z, r);
    const float32x4_t b = vfmsq_f32(vfmsq_f32(x, n, vdupq_n_f32(0x1.62e4p-1f)), n,
                                    vdupq_n_f32(0x1.7f7d1cp-20f));
    const uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_f32(z), 23);
    const float32x4_t k = vreinterpretq_f32_u32(vaddq_u32(e, vreinterpretq_u32_f32(vdupq_n_f32(1))));
    const uint32x4_t c = vcagtq_f32(n, vdupq_n_f32(126));
    const float32x4_t u = vmulq_f32(b, b);
    const float32x4_t j = vfmaq_f32(
        vmulq_f32(vdupq_n_f32(0x1.ffffecp-1f), b),
        vfmaq_f32(vfmaq_f32(vdupq_n_f32(0x1.fffdb6p-2f), vdupq_n_f32(0x1.555e66p-3f), b),
                  vfmaq_f32(vdupq_n_f32(0x1.573e2ep-5f), vdupq_n_f32(0x1.0e4020p-7f), b), u), u);
    const uint64x2_t c64 = vreinterpretq_u64_u32(c);
    if (!(vgetq_lane_u64(c64, 0) | vgetq_lane_u64(c64, 1))) {
        return vfmaq_f32(k, j, k);
    }
    const uint32x4_t d = vandq_u32(vclezq_f32(n), vdupq_n_u32(0x82000000));
    const float32x4_t s1 = vreinterpretq_f32_u32(vaddq_u32(d, vdupq_n_u32(0x7f000000)));
    const float32x4_t s2 = vreinterpretq_f32_u32(vsubq_u32(e, d));
    return vbslq_f32(vcagtq_f32(n, vdupq_n_f32(192)), vmulq_f32(s1, s1),
                     vbslq_f32(c, vmulq_f32(vfmaq_f32(s2, s2, j), s1), vfmaq_f32(k, k, j)));
}
 
static inline int ck_swiglu_ggml_ref_forced(void) {
    const char *env = getenv("CK_DEBUG_SWIGLU_GGML_REF");
    return env && atoi(env) != 0;
}
#endif
 
/**
 * SwiGLU forward pass
 * @test test_swiglu.py::TestSwiGLUForward::test_forward_tokens
 * @test test_swiglu.py::TestSwiGLUForward::test_forward_single
 * @test test_mlp.py::TestMLPForward::test_swiglu_mlp
 * @test test_fused_swiglu_decode.py::TestFusedSwiGLUDecode::test_fused_swiglu_decode
 * @test 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
 */
void swiglu_forward(const float *input,
                    float *output,
                    int tokens,
                    int dim)
{
    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;
        }
    }
}
 
void swiglu_forward_q8_k(const float *input,
                         void *output_q8,
                         int tokens,
                         int dim)
{
    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);
        }
    }
}
 
/**
 * SwiGLU backward pass
 * @test test_swiglu.py::TestSwiGLUBackward::test_backward_tokens
 * @test test_swiglu.py::TestSwiGLUBackward::test_backward_single
 * @test 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
 */
void swiglu_backward(const float *input,
                     const float *d_output,
                     float *d_input,
                     int tokens,
                     int dim)
{
    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;
        }
    }
}
// ============================================================================
// Exact versions using standard library expf (slower but accurate)
// ============================================================================
 
/**
 * SwiGLU forward pass (exact version using stdlib sigmoid)
 * @test test_swiglu.py::TestSwiGLUForward::test_exact_vs_fast
 * @test test_swiglu.py::TestSwiGLUForward::test_exact_single
 *
 * Uses standard library expf for numerical accuracy reference.
 *
 * After changes: make test
 */
typedef struct {
    const float *input;
    float *output;
    int dim;
} ck_swiglu_exact_args_t;
 
static int ck_swiglu_ranges_overlap(const float *input,
                                    const float *output,
                                    int tokens,
                                    int dim)
{
    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;
}
 
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)
{
    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;
    }
}
 
static void ck_swiglu_exact_rows(int begin, int end, void *opaque)
{
    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;
        }
    }
}
 
void swiglu_forward_exact(const float *input,
                          float *output,
                          int tokens,
                          int dim)
{
    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);
}
 
typedef struct {
    const float *input;
    float *output;
    int dim;
} ck_swiglu_ggml_args_t;
 
static void ck_swiglu_ggml_rows(int begin, int end, void *opaque)
{
    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];
        }
    }
}
 
void swiglu_forward_ggml(const float *input,
                         float *output,
                         int tokens,
                         int dim)
{
    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);
}
 
void swiglu_forward_ggml_split(const float *gate,
                               const float *up,
                               float *output,
                               int tokens,
                               int dim)
{
    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];
        }
    }
}
 
/* Truthful provider-name helper for the ggml SwiGLU providers
 * (swiglu_forward_ggml / swiglu_forward_ggml_split), same pattern as
 * ck_q4_k_q8_k_provider_name. */
const char *ck_swiglu_ggml_provider_name(void)
{
#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
}
 
#if defined(__AVX512F__)
typedef __m512 (*ck_sleef_expf16_fn)(__m512);
 
static ck_sleef_expf16_fn ck_pytorch_swiglu_expf16 = NULL;
static void *ck_pytorch_swiglu_sleef_handle = NULL;
static pthread_once_t ck_pytorch_swiglu_once = PTHREAD_ONCE_INIT;
 
static void ck_bind_pytorch_swiglu_sleef(void)
{
    const char *library = getenv("CK_SLEEF_LIBRARY");
    if (library && *library) {
        ck_pytorch_swiglu_sleef_handle = dlopen(library, RTLD_NOW | RTLD_LOCAL);
        if (ck_pytorch_swiglu_sleef_handle) {
            ck_pytorch_swiglu_expf16 = (ck_sleef_expf16_fn)dlsym(
                ck_pytorch_swiglu_sleef_handle, "Sleef_expf16_u10");
        }
    } else {
        ck_pytorch_swiglu_expf16 =
            (ck_sleef_expf16_fn)dlsym(RTLD_DEFAULT, "Sleef_expf16_u10");
    }
}
#endif
 
/*
 * Match the ATen BF16 Qwen MLP expression, not the GGML FP32 expression:
 *
 *   silu_bf16 = aten::silu(gate_bf16)
 *   out_bf16  = silu_bf16 * up_bf16
 *
 * PyTorch 2.8's reduced-floating CPU SiLU converts BF16 lanes to FP32, uses
 * its SLEEF vector exponential, evaluates x / (1 + exp(-x)), and rounds the
 * SiLU result to BF16. The following TensorIterator multiplication converts
 * both BF16 operands to FP32 and performs a second BF16 store. Keeping the
 * SiLU intermediate in FP32 is close, but is a different storage contract and
 * changes long teacher-forced generation after repeated MLP blocks.
 *
 * Source oracle: PyTorch aten/src/ATen/native/cpu/Activation.cpp,
 * silu_kernel(), commit a1cb3cc05d46d198467bebbb6e8fba50a325d4e7.
 */
void swiglu_forward_pytorch_bf16_storage(const float *input,
                                         float *output,
                                         int tokens,
                                         int dim)
{
    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));
        }
    }
}
 
/**
 * SwiGLU backward pass (exact version using stdlib sigmoid)
 * @test test_swiglu.py::TestSwiGLUBackward::test_exact_vs_fast
 * @test test_swiglu.py::TestSwiGLUBackward::test_exact_single
 *
 * Uses standard library expf for numerical accuracy reference.
 *
 * After changes: make test
 */
void swiglu_backward_exact(const float *input,
                           const float *d_output,
                           float *d_input,
                           int tokens,
                           int dim)
{
    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;
        }
    }
}