deltanet_kernels.c File Reference

FP32 Gated DeltaNet kernels for Qwen3.5-style recurrent attention. More...

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

Go to the source code of this file.

Macros
#define	CK_DELTANET_LLAMA_CHUNK_MAX_DIM   256
#define	CK_DELTANET_LLAMA_CHUNK_SIZE   64
#define	CK_DELTANET_MAX_STACK_DIM   4096
#define	CK_DELTANET_NOINLINE

Typedefs
typedef float(*	ck_deltanet_libm_f32_fn) (float)
typedef void(*	ck_deltanet_mkl_vsexp_fn) (int, const float *, float *)

Functions
static void	ck_bind_deltanet_llama_libm (void)
static void	ck_bind_deltanet_pytorch_primitives (void)
static int	ck_deltanet_ceil_log2 (int value)
static int	ck_deltanet_force_ref (void)
static float	ck_deltanet_llama_sigmoidf (float x)
static void	ck_deltanet_pytorch_gate_values (const float *g, const float *beta, float *gate_values, float *beta_values, int num_heads)
static void	ck_deltanet_pytorch_outer_sum (const float *matrix, const float *row_weights, float *output, int state_dim)
static float	ck_deltanet_sigmoidf (float x)
void	gated_deltanet_autoregressive_backward (const float *d_out, const float *d_state_out, const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, const float *state_out, float *d_q, float *d_k, float *d_v, float *d_g, float *d_beta, float *d_state_in, int num_heads, int state_dim, float norm_eps)
void	gated_deltanet_autoregressive_backward_ref (const float *d_out, const float *d_state_out, const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, const float *state_out, float *d_q, float *d_k, float *d_v, float *d_g, float *d_beta, float *d_state_in, int num_heads, int state_dim, float norm_eps)
void	gated_deltanet_autoregressive_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)
void	gated_deltanet_autoregressive_forward_ref (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int state_dim, float norm_eps)
const char *	gated_deltanet_impl_name (void)
void	gated_deltanet_llama_avx2_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps)
void	gated_deltanet_llama_avx2_forward_head_range (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps, int head_begin, int head_end)
static void	gated_deltanet_llama_avx2_grouped_forward_impl (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps, int head_begin, int head_end, int pytorch_bf16_boundaries)
static void	gated_deltanet_llama_avx2_grouped_forward_transposed_impl (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps, int head_begin, int head_end)
void	gated_deltanet_llama_avx2_prefill_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)
void	gated_deltanet_llama_chunk64_head_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int head, int state_dim)
void	gated_deltanet_llama_chunk64_prefill_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)
void	gated_deltanet_prefill_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int state_dim, float norm_eps)
void	gated_deltanet_pytorch_gate_values_debug (const float *g, const float *beta, float *gate_values, float *beta_values, int num_heads)
void	gated_deltanet_pytorch_grouped_bf16_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int num_heads, int group_count, int state_dim, float norm_eps)
void	gated_deltanet_pytorch_grouped_bf16_forward_debug (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, float *decayed_state, float *memory, float *delta, int num_heads, int group_count, int state_dim, float norm_eps)
static void	gated_deltanet_pytorch_grouped_bf16_forward_impl (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, float *debug_decayed_state, float *debug_memory, float *debug_delta, int num_heads, int group_count, int state_dim, float norm_eps)
void	gated_deltanet_pytorch_grouped_bf16_prefill_forward (const float *q, const float *k, const float *v, const float *g, const float *beta, const float *state_in, float *state_out, float *out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)

Variables
static void *	ck_deltanet_libm_handle = NULL
static pthread_once_t	ck_deltanet_libm_once = PTHREAD_ONCE_INIT
static ck_deltanet_libm_f32_fn	ck_deltanet_llama_expf = NULL
static void *	ck_deltanet_mkl_handle = NULL
static pthread_once_t	ck_deltanet_pytorch_primitives_once = PTHREAD_ONCE_INIT
static ck_deltanet_mkl_vsexp_fn	ck_deltanet_pytorch_vsexp = NULL

Detailed Description

FP32 Gated DeltaNet kernels for Qwen3.5-style recurrent attention.

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

This file implements the single-token recurrent update used by the qwen3next / Gated DeltaNet path in llama.cpp.

Per head, matching llama.cpp qwen35/qwen3next autoregressive DeltaNet: q_scaled = q / sqrt(state_dim) // q and k arrive pre-normalized k_hat = k beta_s = sigmoid(beta) gate = exp(g) S = gate * S_prev kv_mem = S^T * k_hat delta = (v - kv_mem) * beta_s S_new = S + outer(k_hat, delta) out = S_new^T * q_scaled

Design:

• *_ref is the scalar reference implementation.
• *_avx keeps a simple 1-row vector walk.
• *_avx2 precomputes scaled q rows and unrolls the state sweep in row pairs to reduce loop overhead and layout churn.
• The public dispatcher selects the best compiled ISA unless strict parity is enabled, in which case it falls back to *_ref.

Definition in file deltanet_kernels.c.

Macro Definition Documentation

CK_DELTANET_LLAMA_CHUNK_MAX_DIM

#define CK_DELTANET_LLAMA_CHUNK_MAX_DIM   256

Definition at line 54 of file deltanet_kernels.c.

CK_DELTANET_LLAMA_CHUNK_SIZE

#define CK_DELTANET_LLAMA_CHUNK_SIZE   64

Definition at line 53 of file deltanet_kernels.c.

CK_DELTANET_MAX_STACK_DIM

#define CK_DELTANET_MAX_STACK_DIM   4096

Definition at line 52 of file deltanet_kernels.c.

CK_DELTANET_NOINLINE

#define CK_DELTANET_NOINLINE

Definition at line 59 of file deltanet_kernels.c.

Typedef Documentation

ck_deltanet_libm_f32_fn

typedef float(* ck_deltanet_libm_f32_fn) (float)

Definition at line 62 of file deltanet_kernels.c.

ck_deltanet_mkl_vsexp_fn

typedef void(* ck_deltanet_mkl_vsexp_fn) (int, const float *, float *)

Definition at line 99 of file deltanet_kernels.c.

Function Documentation

ck_bind_deltanet_llama_libm()

static void ck_bind_deltanet_llama_libm ( void  )

static

Definition at line 67 of file deltanet_kernels.c.

{
    ck_deltanet_libm_handle = dlopen("libm.so.6", RTLD_NOW | RTLD_LOCAL);
    if (ck_deltanet_libm_handle) {
        ck_deltanet_llama_expf =
            (ck_deltanet_libm_f32_fn)dlsym(ck_deltanet_libm_handle, "expf");
    }
    if (!ck_deltanet_llama_expf) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: llama.cpp DeltaNet requires "
                "expf from libm.so.6\n");
        abort();
    }
}

References ck_deltanet_libm_handle, and ck_deltanet_llama_expf.

Referenced by ck_deltanet_llama_sigmoidf(), gated_deltanet_llama_avx2_grouped_forward_impl(), and gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

ck_bind_deltanet_pytorch_primitives()

static void ck_bind_deltanet_pytorch_primitives ( void  )

static

Definition at line 104 of file deltanet_kernels.c.

{
    const char *mkl_library = getenv("CK_MKL_LIBRARY");
    if (mkl_library && *mkl_library) {
        ck_deltanet_mkl_handle = dlopen(mkl_library, RTLD_NOW | RTLD_LOCAL);
        if (ck_deltanet_mkl_handle) {
            ck_deltanet_pytorch_vsexp =
                (ck_deltanet_mkl_vsexp_fn)dlsym(
                    ck_deltanet_mkl_handle, "vsExp");
        }
    } else {
        ck_deltanet_pytorch_vsexp =
            (ck_deltanet_mkl_vsexp_fn)dlsym(RTLD_DEFAULT, "vsExp");
    }
 
#if defined(__AVX512F__)
    const char *library = getenv("CK_SLEEF_LIBRARY");
    if (library && *library) {
        ck_deltanet_sleef_handle = dlopen(library, RTLD_NOW | RTLD_LOCAL);
        if (ck_deltanet_sleef_handle) {
            ck_deltanet_pytorch_expf16 =
                (ck_deltanet_sleef_expf16_fn)dlsym(
                    ck_deltanet_sleef_handle, "Sleef_expf16_u10");
        }
    } else {
        ck_deltanet_pytorch_expf16 =
            (ck_deltanet_sleef_expf16_fn)dlsym(
                RTLD_DEFAULT, "Sleef_expf16_u10");
    }
#endif
}

References ck_deltanet_mkl_handle, ck_deltanet_pytorch_vsexp, and RTLD_DEFAULT.

Referenced by ck_deltanet_pytorch_gate_values().

ck_deltanet_ceil_log2()

static int ck_deltanet_ceil_log2 ( int  value )

static

Definition at line 812 of file deltanet_kernels.c.

{
    int result = 0;
    int power = 1;
    while (power < value) {
        power <<= 1;
        ++result;
    }
    return result;
}

Referenced by ck_deltanet_pytorch_outer_sum().

ck_deltanet_force_ref()

static int ck_deltanet_force_ref ( void  )

static

Definition at line 1884 of file deltanet_kernels.c.

{
    const char *env = getenv("CK_DELTANET_FORCE_REF");
    return env && atoi(env) != 0;
}

Referenced by gated_deltanet_autoregressive_forward(), and gated_deltanet_impl_name().

ck_deltanet_llama_sigmoidf()

static float ck_deltanet_llama_sigmoidf ( float  x )

inlinestatic

Definition at line 87 of file deltanet_kernels.c.

{
    pthread_once(&ck_deltanet_libm_once, ck_bind_deltanet_llama_libm);
    return 1.0f / (1.0f + ck_deltanet_llama_expf(-x));
}

References ck_bind_deltanet_llama_libm(), ck_deltanet_libm_once, and ck_deltanet_llama_expf.

Referenced by gated_deltanet_llama_avx2_grouped_forward_impl(), and gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

ck_deltanet_pytorch_gate_values()

static void ck_deltanet_pytorch_gate_values ( const float *  g, const float *  beta, float *  gate_values, float *  beta_values, int  num_heads  )

static

Definition at line 136 of file deltanet_kernels.c.

{
    pthread_once(
        &ck_deltanet_pytorch_primitives_once,
        ck_bind_deltanet_pytorch_primitives);
    if (!ck_deltanet_pytorch_vsexp) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch DeltaNet requires "
                "MKL vsExp; set CK_MKL_LIBRARY\n");
        abort();
    }
    ck_deltanet_pytorch_vsexp(num_heads, g, gate_values);
 
    int h = 0;
#if defined(__AVX512F__)
    if (ck_deltanet_pytorch_expf16) {
        const __m512 one = _mm512_set1_ps(1.0f);
        for (; h + 15 < num_heads; h += 16) {
            const __m512 bv = _mm512_loadu_ps(beta + h);
            const __m512 beta_exp = ck_deltanet_pytorch_expf16(
                _mm512_sub_ps(_mm512_setzero_ps(), bv));
            const __m512 beta_sigmoid = _mm512_div_ps(
                one, _mm512_add_ps(one, beta_exp));
            float beta_lanes[16];
            _mm512_storeu_ps(beta_lanes, beta_sigmoid);
            for (int lane = 0; lane < 16; ++lane) {
                beta_values[h + lane] = bf16_to_float(
                    float_to_bf16(beta_lanes[lane]));
            }
        }
    }
#endif
    for (; h < num_heads; ++h) {
        beta_values[h] = bf16_to_float(float_to_bf16(
            ck_deltanet_sigmoidf(beta[h])));
    }
}

References bf16_to_float(), ck_bind_deltanet_pytorch_primitives(), ck_deltanet_pytorch_primitives_once, ck_deltanet_pytorch_vsexp, ck_deltanet_sigmoidf(), and float_to_bf16().

Referenced by gated_deltanet_pytorch_gate_values_debug(), and gated_deltanet_pytorch_grouped_bf16_forward_impl().

ck_deltanet_pytorch_outer_sum()

static void ck_deltanet_pytorch_outer_sum ( const float *  matrix, const float *  row_weights, float *  output, int  state_dim  )

static

Definition at line 826 of file deltanet_kernels.c.

{
    const int num_levels = 4;
    int level_power = ck_deltanet_ceil_log2(state_dim) / num_levels;
    if (level_power < 4) {
        level_power = 4;
    }
    const int level_step = 1 << level_power;
    const int level_mask = level_step - 1;
    int col = 0;
 
#if defined(__AVX512F__)
    /* PyTorch vectorized_outer_sum reduces four adjacent vectors together. */
    for (; col + 63 < state_dim; col += 64) {
        __m512 acc[4][4];
        for (int level = 0; level < num_levels; ++level) {
            for (int block = 0; block < 4; ++block) {
                acc[level][block] = _mm512_setzero_ps();
            }
        }
 
        int i = 0;
        for (; i + level_step <= state_dim;) {
            for (int j = 0; j < level_step; ++j, ++i) {
                const float *row = matrix + (size_t)i * (size_t)state_dim + col;
                const __m512 weight = _mm512_set1_ps(row_weights[i]);
                for (int block = 0; block < 4; ++block) {
                    const __m512 product = _mm512_mul_ps(
                        _mm512_loadu_ps(row + block * 16), weight);
                    acc[0][block] = _mm512_add_ps(acc[0][block], product);
                }
            }
 
            for (int level = 1; level < num_levels; ++level) {
                for (int block = 0; block < 4; ++block) {
                    acc[level][block] = _mm512_add_ps(
                        acc[level][block], acc[level - 1][block]);
                    acc[level - 1][block] = _mm512_setzero_ps();
                }
                const int mask = level_mask << (level * level_power);
                if ((i & mask) != 0) {
                    break;
                }
            }
        }
 
        for (; i < state_dim; ++i) {
            const float *row = matrix + (size_t)i * (size_t)state_dim + col;
            const __m512 weight = _mm512_set1_ps(row_weights[i]);
            for (int block = 0; block < 4; ++block) {
                const __m512 product = _mm512_mul_ps(
                    _mm512_loadu_ps(row + block * 16), weight);
                acc[0][block] = _mm512_add_ps(acc[0][block], product);
            }
        }
 
        for (int level = 1; level < num_levels; ++level) {
            for (int block = 0; block < 4; ++block) {
                acc[0][block] = _mm512_add_ps(
                    acc[0][block], acc[level][block]);
            }
        }
        for (int block = 0; block < 4; ++block) {
            _mm512_storeu_ps(output + col + block * 16, acc[0][block]);
        }
    }
#elif defined(__AVX2__)
    for (; col + 31 < state_dim; col += 32) {
        __m256 acc[4][4];
        for (int level = 0; level < num_levels; ++level) {
            for (int block = 0; block < 4; ++block) {
                acc[level][block] = _mm256_setzero_ps();
            }
        }
 
        int i = 0;
        for (; i + level_step <= state_dim;) {
            for (int j = 0; j < level_step; ++j, ++i) {
                const float *row = matrix + (size_t)i * (size_t)state_dim + col;
                const __m256 weight = _mm256_set1_ps(row_weights[i]);
                for (int block = 0; block < 4; ++block) {
                    const __m256 product = _mm256_mul_ps(
                        _mm256_loadu_ps(row + block * 8), weight);
                    acc[0][block] = _mm256_add_ps(acc[0][block], product);
                }
            }
            for (int level = 1; level < num_levels; ++level) {
                for (int block = 0; block < 4; ++block) {
                    acc[level][block] = _mm256_add_ps(
                        acc[level][block], acc[level - 1][block]);
                    acc[level - 1][block] = _mm256_setzero_ps();
                }
                const int mask = level_mask << (level * level_power);
                if ((i & mask) != 0) {
                    break;
                }
            }
        }
        for (; i < state_dim; ++i) {
            const float *row = matrix + (size_t)i * (size_t)state_dim + col;
            const __m256 weight = _mm256_set1_ps(row_weights[i]);
            for (int block = 0; block < 4; ++block) {
                const __m256 product = _mm256_mul_ps(
                    _mm256_loadu_ps(row + block * 8), weight);
                acc[0][block] = _mm256_add_ps(acc[0][block], product);
            }
        }
        for (int level = 1; level < num_levels; ++level) {
            for (int block = 0; block < 4; ++block) {
                acc[0][block] = _mm256_add_ps(
                    acc[0][block], acc[level][block]);
            }
        }
        for (int block = 0; block < 4; ++block) {
            _mm256_storeu_ps(output + col + block * 8, acc[0][block]);
        }
    }
#endif
 
    /* Production Qwen dimensions are covered above. Keep a deterministic
     * scalar fallback for uncommon tail widths. */
    for (; col < state_dim; ++col) {
        float sum = 0.0f;
        for (int row = 0; row < state_dim; ++row) {
            sum += matrix[(size_t)row * (size_t)state_dim + col] *
                   row_weights[row];
        }
        output[col] = sum;
    }
}

References ck_deltanet_ceil_log2(), and mask.

Referenced by gated_deltanet_pytorch_grouped_bf16_forward_impl().

ck_deltanet_sigmoidf()

static float ck_deltanet_sigmoidf ( float  x )

inlinestatic

Definition at line 82 of file deltanet_kernels.c.

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

Referenced by ck_deltanet_pytorch_gate_values(), gated_deltanet_autoregressive_backward_ref(), gated_deltanet_autoregressive_forward_ref(), and gated_deltanet_llama_avx2_grouped_forward_impl().

gated_deltanet_autoregressive_backward()

void gated_deltanet_autoregressive_backward	(	const float *	d_out,
		const float *	d_state_out,
		const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		const float *	state_out,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_g,
		float *	d_beta,
		float *	d_state_in,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1988 of file deltanet_kernels.c.

{
    if (!d_out || !d_state_out || !q || !k || !v || !g || !beta || !state_in || !state_out ||
        !d_q || !d_k || !d_v || !d_g || !d_beta || !d_state_in) {
        return;
    }
    if (num_heads <= 0 || state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM) {
        return;
    }
 
    gated_deltanet_autoregressive_backward_ref(
        d_out,
        d_state_out,
        q,
        k,
        v,
        g,
        beta,
        state_in,
        state_out,
        d_q,
        d_k,
        d_v,
        d_g,
        d_beta,
        d_state_in,
        num_heads,
        state_dim,
        norm_eps);
}

References CK_DELTANET_MAX_STACK_DIM, and gated_deltanet_autoregressive_backward_ref().

gated_deltanet_autoregressive_backward_ref()

void gated_deltanet_autoregressive_backward_ref	(	const float *	d_out,
		const float *	d_state_out,
		const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		const float *	state_out,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_g,
		float *	d_beta,
		float *	d_state_in,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1339 of file deltanet_kernels.c.

{
    const float q_scale = 1.0f / sqrtf((float)state_dim);
    const size_t vec_stride = (size_t)state_dim;
    const size_t state_stride = (size_t)state_dim * (size_t)state_dim;
 
    float q_hat[CK_DELTANET_MAX_STACK_DIM];
    float k_hat[CK_DELTANET_MAX_STACK_DIM];
    float kv_mem[CK_DELTANET_MAX_STACK_DIM];
    float delta[CK_DELTANET_MAX_STACK_DIM];
    float d_q_hat[CK_DELTANET_MAX_STACK_DIM];
    float d_k_hat[CK_DELTANET_MAX_STACK_DIM];
    float d_mem[CK_DELTANET_MAX_STACK_DIM];
 
    for (int h = 0; h < num_heads; ++h) {
        const float *d_out_head = d_out + (size_t)h * vec_stride;
        const float *d_state_out_head = d_state_out + (size_t)h * state_stride;
        const float *q_head = q + (size_t)h * vec_stride;
        const float *k_head = k + (size_t)h * vec_stride;
        const float *v_head = v + (size_t)h * vec_stride;
        const float *state_prev = state_in + (size_t)h * state_stride;
        const float *state_cur = state_out + (size_t)h * state_stride;
        float *d_q_head = d_q + (size_t)h * vec_stride;
        float *d_k_head = d_k + (size_t)h * vec_stride;
        float *d_v_head = d_v + (size_t)h * vec_stride;
        float *d_state_prev = d_state_in + (size_t)h * state_stride;
 
        const float beta_s = ck_deltanet_sigmoidf(beta[h]);
        const float gate = expf(g[h]);
 
        float qk_dot = 0.0f;
        float out_delta_dot = 0.0f;
        float beta_acc = 0.0f;
        float gate_acc = 0.0f;
 
        for (int i = 0; i < state_dim; ++i) {
            q_hat[i] = q_head[i] * q_scale;
            k_hat[i] = k_head[i];
            kv_mem[i] = 0.0f;
            d_q_hat[i] = 0.0f;
            d_k_hat[i] = 0.0f;
            d_mem[i] = 0.0f;
            d_v_head[i] = 0.0f;
            qk_dot += q_hat[i] * k_hat[i];
        }
 
        for (int col = 0; col < state_dim; ++col) {
            float mem = 0.0f;
            for (int row = 0; row < state_dim; ++row) {
                mem += (state_prev[(size_t)row * (size_t)state_dim + (size_t)col] * gate) * k_hat[row];
            }
            kv_mem[col] = mem;
            delta[col] = (v_head[col] - mem) * beta_s;
            out_delta_dot += d_out_head[col] * delta[col];
        }
 
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_off = (size_t)row * (size_t)state_dim;
            float dq_acc = 0.0f;
            float dk_acc = q_hat[row] * out_delta_dot;
            for (int col = 0; col < state_dim; ++col) {
                const float d_state_direct = d_state_out_head[row_off + (size_t)col];
                dq_acc += state_cur[row_off + (size_t)col] * d_out_head[col];
                dk_acc += d_state_direct * delta[col];
            }
            d_q_hat[row] = dq_acc;
            d_k_hat[row] = dk_acc;
        }
 
        for (int col = 0; col < state_dim; ++col) {
            float d_delta_acc = d_out_head[col] * qk_dot;
            for (int row = 0; row < state_dim; ++row) {
                d_delta_acc += d_state_out_head[(size_t)row * (size_t)state_dim + (size_t)col] * k_hat[row];
            }
 
            d_v_head[col] = beta_s * d_delta_acc;
            d_mem[col] = -beta_s * d_delta_acc;
            beta_acc += d_delta_acc * (v_head[col] - kv_mem[col]);
        }
 
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_off = (size_t)row * (size_t)state_dim;
            float s_dm_acc = 0.0f;
            for (int col = 0; col < state_dim; ++col) {
                s_dm_acc += (state_prev[row_off + (size_t)col] * gate) * d_mem[col];
            }
            d_k_hat[row] += s_dm_acc;
        }
 
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_off = (size_t)row * (size_t)state_dim;
            for (int col = 0; col < state_dim; ++col) {
                const float d_state_total = d_state_out_head[row_off + (size_t)col]
                                          + q_hat[row] * d_out_head[col]
                                          + k_hat[row] * d_mem[col];
                d_state_prev[row_off + (size_t)col] = gate * d_state_total;
                gate_acc += d_state_total * state_prev[row_off + (size_t)col];
            }
        }
 
        for (int i = 0; i < state_dim; ++i) {
            d_q_head[i] = d_q_hat[i] * q_scale;
            d_k_head[i] = d_k_hat[i];
        }
 
        d_g[h] = gate_acc * gate;
        d_beta[h] = beta_acc * beta_s * (1.0f - beta_s);
    }
}

References CK_DELTANET_MAX_STACK_DIM, and ck_deltanet_sigmoidf().

Referenced by gated_deltanet_autoregressive_backward().

gated_deltanet_autoregressive_forward()

void gated_deltanet_autoregressive_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1906 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
        return;
    }
    if (num_heads <= 0 || state_dim <= 0) {
        return;
    }
 
    /*
     * q and k arrive pre-normalized by recurrent_qk_l2_norm, so the
     * ISA-specialized kernels can follow the same contract as the scalar ref.
     */
    if (ck_strict_parity_enabled() || ck_deltanet_force_ref()) {
        gated_deltanet_autoregressive_forward_ref(
            q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
        return;
    }
#if defined(__AVX512F__)
    gated_deltanet_autoregressive_forward_avx512(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#elif defined(__AVX2__)
    gated_deltanet_autoregressive_forward_avx2(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#elif defined(__AVX__)
    gated_deltanet_autoregressive_forward_avx(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#else
    gated_deltanet_autoregressive_forward_ref(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#endif
}

References ck_deltanet_force_ref(), ck_strict_parity_enabled(), and gated_deltanet_autoregressive_forward_ref().

Referenced by ck_test_gated_deltanet_autoregressive(), and gated_deltanet_prefill_forward().

gated_deltanet_autoregressive_forward_ref()

void gated_deltanet_autoregressive_forward_ref	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1280 of file deltanet_kernels.c.

{
    const float q_scale = 1.0f / sqrtf((float)state_dim);
    const size_t vec_stride = (size_t)state_dim;
    const size_t state_stride = (size_t)state_dim * (size_t)state_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        const float *q_head = q + (size_t)h * vec_stride;
        const float *k_head = k + (size_t)h * vec_stride;
        const float *v_head = v + (size_t)h * vec_stride;
        const float *state_prev = state_in + (size_t)h * state_stride;
        float *state_cur = state_out + (size_t)h * state_stride;
        float *out_head = out + (size_t)h * vec_stride;
 
        const float beta_s = ck_deltanet_sigmoidf(beta[h]);
        const float gate = expf(g[h]);
 
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_off = (size_t)row * (size_t)state_dim;
            for (int col = 0; col < state_dim; ++col) {
                state_cur[row_off + (size_t)col] = state_prev[row_off + (size_t)col] * gate;
            }
        }
 
        for (int col = 0; col < state_dim; ++col) {
            float kv_mem = 0.0f;
            for (int row = 0; row < state_dim; ++row) {
                const float k_hat = k_head[row];
                kv_mem += state_cur[(size_t)row * (size_t)state_dim + (size_t)col] * k_hat;
            }
 
            const float delta = (v_head[col] - kv_mem) * beta_s;
            for (int row = 0; row < state_dim; ++row) {
                const float k_hat = k_head[row];
                state_cur[(size_t)row * (size_t)state_dim + (size_t)col] += k_hat * delta;
            }
        }
 
        for (int col = 0; col < state_dim; ++col) {
            float acc = 0.0f;
            for (int row = 0; row < state_dim; ++row) {
                const float q_hat = q_head[row] * q_scale;
                acc += state_cur[(size_t)row * (size_t)state_dim + (size_t)col] * q_hat;
            }
            out_head[col] = acc;
        }
    }
}

References ck_deltanet_sigmoidf().

Referenced by gated_deltanet_autoregressive_forward(), and gated_deltanet_llama_avx2_grouped_forward_impl().

gated_deltanet_impl_name()

const char * gated_deltanet_impl_name

(

void

)

Definition at line 1890 of file deltanet_kernels.c.

{
    if (ck_strict_parity_enabled() || ck_deltanet_force_ref()) {
        return "REF";
    }
#if defined(__AVX512F__)
    return "AVX512";
#elif defined(__AVX2__)
    return "AVX2";
#elif defined(__AVX__)
    return "AVX";
#else
    return "REF";
#endif
}

References ck_deltanet_force_ref(), and ck_strict_parity_enabled().

gated_deltanet_llama_avx2_forward()

void gated_deltanet_llama_avx2_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 767 of file deltanet_kernels.c.

{
    gated_deltanet_llama_avx2_grouped_forward_transposed_impl(
        q, k, v, g, beta, state_in, state_out, out,
        num_heads, group_count, state_dim, norm_eps, 0, num_heads);
}

References gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

Referenced by gated_deltanet_llama_avx2_prefill_forward().

gated_deltanet_llama_avx2_forward_head_range()

void gated_deltanet_llama_avx2_forward_head_range	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps,
		int	head_begin,
		int	head_end
	)

Definition at line 790 of file deltanet_kernels.c.

{
    gated_deltanet_llama_avx2_grouped_forward_transposed_impl(
        q, k, v, g, beta, state_in, state_out, out,
        num_heads, group_count, state_dim, norm_eps,
        head_begin, head_end);
}

References gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

gated_deltanet_llama_avx2_grouped_forward_impl()

static void gated_deltanet_llama_avx2_grouped_forward_impl ( const float *  q, const float *  k, const float *  v, const float *  g, const float *  beta, const float *  state_in, float *  state_out, float *  out, int  num_heads, int  group_count, int  state_dim, float  norm_eps, int  head_begin, int  head_end, int  pytorch_bf16_boundaries  )

static

Definition at line 574 of file deltanet_kernels.c.

{
#if defined(__AVX2__)
    (void) norm_eps;
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        num_heads <= 0 || num_heads > CK_DELTANET_MAX_STACK_DIM ||
        group_count <= 0 || num_heads % group_count != 0 ||
        state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM ||
        head_begin < 0 || head_end < head_begin || head_end > num_heads) {
        return;
    }
    const float scale = ck_deltanet_llama_scale(state_dim);
    const size_t vector_stride = (size_t) state_dim;
    const size_t state_stride = (size_t) state_dim * (size_t) state_dim;
    float column[CK_DELTANET_MAX_STACK_DIM];
    float q_scaled[CK_DELTANET_MAX_STACK_DIM];
 
    for (int h = head_begin; h < head_end; ++h) {
        /* llama.cpp ggml_repeat_4d tiles compact Q/K heads (0..G-1 repeated),
         * while the PyTorch Qwen3-Next reference uses repeat_interleave so
         * each compact head owns H/G adjacent value heads.  These layouts
         * are distinct numerical contracts even though their buffer shapes
         * are identical. */
        const int group = pytorch_bf16_boundaries
            ? h / (num_heads / group_count)
            : h % group_count;
        const float *q_head = q + (size_t) group * vector_stride;
        const float *k_head = k + (size_t) group * vector_stride;
        const float *v_head = v + (size_t) h * vector_stride;
        const float *state_prev = state_in + (size_t) h * state_stride;
        float *state_cur = state_out + (size_t) h * state_stride;
        float *out_head = out + (size_t) h * vector_stride;
        float gate;
        float beta_s;
        if (pytorch_bf16_boundaries) {
            gate = expf(g[h]);
            beta_s = ck_deltanet_sigmoidf(beta[h]);
        } else {
            pthread_once(&ck_deltanet_libm_once, ck_bind_deltanet_llama_libm);
            gate = ck_deltanet_llama_expf(g[h]);
            beta_s = ck_deltanet_llama_sigmoidf(beta[h]);
        }
        if (pytorch_bf16_boundaries) {
            beta_s = bf16_to_float(float_to_bf16(beta_s));
            for (int row = 0; row < state_dim; ++row) {
                q_scaled[row] = q_head[row] * scale;
            }
        }
 
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_offset = (size_t) row * (size_t) state_dim;
            for (int col = 0; col < state_dim; ++col) {
                state_cur[row_offset + (size_t) col] =
                    state_prev[row_offset + (size_t) col] * gate;
            }
        }
 
        for (int col = 0; col < state_dim; ++col) {
            for (int row = 0; row < state_dim; ++row) {
                column[row] = state_cur[(size_t) row * (size_t) state_dim + (size_t) col];
            }
            const float memory = ck_deltanet_llama_avx2_dot(column, k_head, state_dim);
            const float delta = (v_head[col] - memory) * beta_s;
            for (int row = 0; row < state_dim; ++row) {
                const size_t offset = (size_t) row * (size_t) state_dim + (size_t) col;
#if defined(__FMA__)
                const float updated = fmaf(k_head[row], delta, state_cur[offset]);
#else
                const float updated = state_cur[offset] + k_head[row] * delta;
#endif
                state_cur[offset] = updated;
                column[row] = updated;
            }
            if (pytorch_bf16_boundaries) {
                out_head[col] =
                    ck_deltanet_llama_avx2_dot(column, q_scaled, state_dim);
            } else {
                out_head[col] =
                    ck_deltanet_llama_avx2_dot(column, q_head, state_dim) * scale;
            }
        }
    }
#else
    if (group_count != num_heads) {
        return;
    }
    gated_deltanet_autoregressive_forward_ref(
        q, k, v, g, beta, state_in, state_out, out,
        num_heads, state_dim, norm_eps);
#endif
}

References bf16_to_float(), ck_bind_deltanet_llama_libm(), ck_deltanet_libm_once, ck_deltanet_llama_expf, ck_deltanet_llama_sigmoidf(), CK_DELTANET_MAX_STACK_DIM, ck_deltanet_sigmoidf(), float_to_bf16(), and gated_deltanet_autoregressive_forward_ref().

gated_deltanet_llama_avx2_grouped_forward_transposed_impl()

static void gated_deltanet_llama_avx2_grouped_forward_transposed_impl ( const float *  q, const float *  k, const float *  v, const float *  g, const float *  beta, const float *  state_in, float *  state_out, float *  out, int  num_heads, int  group_count, int  state_dim, float  norm_eps, int  head_begin, int  head_end  )

static

Definition at line 687 of file deltanet_kernels.c.

{
#if defined(__AVX2__)
    (void) norm_eps;
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        num_heads <= 0 || group_count <= 0 || num_heads % group_count != 0 ||
        state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM ||
        head_begin < 0 || head_end < head_begin || head_end > num_heads) {
        return;
    }
    const float scale = ck_deltanet_llama_scale(state_dim);
    const size_t vector_stride = (size_t) state_dim;
    const size_t state_stride = (size_t) state_dim * (size_t) state_dim;
 
    pthread_once(&ck_deltanet_libm_once, ck_bind_deltanet_llama_libm);
    for (int h = head_begin; h < head_end; ++h) {
        const int group = h % group_count;
        const float *q_head = q + (size_t) group * vector_stride;
        const float *k_head = k + (size_t) group * vector_stride;
        const float *v_head = v + (size_t) h * vector_stride;
        const float *state_prev = state_in + (size_t) h * state_stride;
        float *state_cur = state_out + (size_t) h * state_stride;
        float *out_head = out + (size_t) h * vector_stride;
        const float gate = ck_deltanet_llama_expf(g[h]);
        const float beta_s = ck_deltanet_llama_sigmoidf(beta[h]);
 
        for (int col = 0; col < state_dim; ++col) {
            const float *prev_col = state_prev + (size_t) col * vector_stride;
            float *cur_col = state_cur + (size_t) col * vector_stride;
            int row = 0;
            const __m256 gate8 = _mm256_set1_ps(gate);
            for (; row + 7 < state_dim; row += 8) {
                const __m256 scaled = _mm256_mul_ps(
                    _mm256_loadu_ps(prev_col + row), gate8);
                _mm256_storeu_ps(cur_col + row, scaled);
            }
            for (; row < state_dim; ++row) {
                cur_col[row] = prev_col[row] * gate;
            }
 
            const float memory =
                ck_deltanet_llama_avx2_dot(cur_col, k_head, state_dim);
            const float delta = (v_head[col] - memory) * beta_s;
            row = 0;
            const __m256 delta8 = _mm256_set1_ps(delta);
            for (; row + 7 < state_dim; row += 8) {
                const __m256 updated = _mm256_fmadd_ps(
                    _mm256_loadu_ps(k_head + row), delta8,
                    _mm256_loadu_ps(cur_col + row));
                _mm256_storeu_ps(cur_col + row, updated);
            }
            for (; row < state_dim; ++row) {
                cur_col[row] = fmaf(k_head[row], delta, cur_col[row]);
            }
            out_head[col] =
                ck_deltanet_llama_avx2_dot(cur_col, q_head, state_dim) * scale;
        }
    }
#else
    (void)q; (void)k; (void)v; (void)g; (void)beta;
    (void)state_in; (void)state_out; (void)out;
    (void)num_heads; (void)group_count; (void)state_dim; (void)norm_eps;
    (void)head_begin; (void)head_end;
#endif
}

References ck_bind_deltanet_llama_libm(), ck_deltanet_libm_once, ck_deltanet_llama_expf, ck_deltanet_llama_sigmoidf(), and CK_DELTANET_MAX_STACK_DIM.

Referenced by gated_deltanet_llama_avx2_forward(), and gated_deltanet_llama_avx2_forward_head_range().

gated_deltanet_llama_avx2_prefill_forward()

void gated_deltanet_llama_avx2_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1131 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || state_dim <= 0) {
        return;
    }
    const size_t qk_stride = (size_t) group_count * (size_t) state_dim;
    const size_t value_stride = (size_t) num_heads * (size_t) state_dim;
    const size_t gate_stride = (size_t) num_heads;
    for (int row = 0; row < rows; ++row) {
        gated_deltanet_llama_avx2_forward(
            q + (size_t) row * qk_stride,
            k + (size_t) row * qk_stride,
            v + (size_t) row * value_stride,
            g + (size_t) row * gate_stride,
            beta + (size_t) row * gate_stride,
            row == 0 ? state_in : state_out,
            state_out,
            out + (size_t) row * value_stride,
            num_heads, group_count, state_dim, norm_eps);
    }
}

References gated_deltanet_llama_avx2_forward().

Referenced by gated_deltanet_llama_chunk64_prefill_forward().

gated_deltanet_llama_chunk64_head_forward()

void gated_deltanet_llama_chunk64_head_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	head,
		int	state_dim
	)

Definition at line 1202 of file deltanet_kernels.c.

{
#if defined(__AVX2__)
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || head < 0 || head >= num_heads ||
        state_dim <= 0 || state_dim > CK_DELTANET_LLAMA_CHUNK_MAX_DIM) {
        return;
    }
    gated_deltanet_llama_chunk64_head(
        q, k, v, g, beta, state_in, state_out, out,
        rows, num_heads, group_count, head, state_dim);
#else
    (void)q;
    (void)k;
    (void)v;
    (void)g;
    (void)beta;
    (void)state_in;
    (void)state_out;
    (void)out;
    (void)rows;
    (void)num_heads;
    (void)group_count;
    (void)head;
    (void)state_dim;
#endif
}

References CK_DELTANET_LLAMA_CHUNK_MAX_DIM.

gated_deltanet_llama_chunk64_prefill_forward()

void gated_deltanet_llama_chunk64_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1167 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || state_dim <= 0) {
        return;
    }
#if defined(__AVX2__)
    (void)norm_eps;
    if (state_dim <= CK_DELTANET_LLAMA_CHUNK_MAX_DIM) {
        for (int head = 0; head < num_heads; ++head) {
            gated_deltanet_llama_chunk64_head(
                q, k, v, g, beta, state_in, state_out, out,
                rows, num_heads, group_count, head, state_dim);
        }
        return;
    }
#endif
    gated_deltanet_llama_avx2_prefill_forward(
        q, k, v, g, beta, state_in, state_out, out,
        rows, num_heads, group_count, state_dim, norm_eps);
}

References CK_DELTANET_LLAMA_CHUNK_MAX_DIM, and gated_deltanet_llama_avx2_prefill_forward().

gated_deltanet_prefill_forward()

void gated_deltanet_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1949 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
        return;
    }
    if (rows <= 0 || num_heads <= 0 || state_dim <= 0) {
        return;
    }
 
    const size_t vector_stride = (size_t)num_heads * (size_t)state_dim;
    const size_t gate_stride = (size_t)num_heads;
    for (int row = 0; row < rows; ++row) {
        const float *row_state_in = row == 0 ? state_in : state_out;
        gated_deltanet_autoregressive_forward(
            q + (size_t)row * vector_stride,
            k + (size_t)row * vector_stride,
            v + (size_t)row * vector_stride,
            g + (size_t)row * gate_stride,
            beta + (size_t)row * gate_stride,
            row_state_in,
            state_out,
            out + (size_t)row * vector_stride,
            num_heads,
            state_dim,
            norm_eps);
    }
}

References gated_deltanet_autoregressive_forward().

gated_deltanet_pytorch_gate_values_debug()

void gated_deltanet_pytorch_gate_values_debug	(	const float *	g,
		const float *	beta,
		float *	gate_values,
		float *	beta_values,
		int	num_heads
	)

Definition at line 178 of file deltanet_kernels.c.

{
    if (!g || !beta || !gate_values || !beta_values || num_heads <= 0 ||
        num_heads > CK_DELTANET_MAX_STACK_DIM) {
        return;
    }
    ck_deltanet_pytorch_gate_values(
        g, beta, gate_values, beta_values, num_heads);
}

References CK_DELTANET_MAX_STACK_DIM, and ck_deltanet_pytorch_gate_values().

gated_deltanet_pytorch_grouped_bf16_forward()

void gated_deltanet_pytorch_grouped_bf16_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1089 of file deltanet_kernels.c.

{
    gated_deltanet_pytorch_grouped_bf16_forward_impl(
        q, k, v, g, beta, state_in, state_out, out,
        NULL, NULL, NULL,
        num_heads, group_count, state_dim, norm_eps);
}

References gated_deltanet_pytorch_grouped_bf16_forward_impl().

Referenced by gated_deltanet_pytorch_grouped_bf16_prefill_forward().

gated_deltanet_pytorch_grouped_bf16_forward_debug()

void gated_deltanet_pytorch_grouped_bf16_forward_debug	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		float *	decayed_state,
		float *	memory,
		float *	delta,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1108 of file deltanet_kernels.c.

{
    gated_deltanet_pytorch_grouped_bf16_forward_impl(
        q, k, v, g, beta, state_in, state_out, out,
        decayed_state, memory, delta,
        num_heads, group_count, state_dim, norm_eps);
}

References gated_deltanet_pytorch_grouped_bf16_forward_impl().

gated_deltanet_pytorch_grouped_bf16_forward_impl()

static void gated_deltanet_pytorch_grouped_bf16_forward_impl ( const float *  q, const float *  k, const float *  v, const float *  g, const float *  beta, const float *  state_in, float *  state_out, float *  out, float *  debug_decayed_state, float *  debug_memory, float *  debug_delta, int  num_heads, int  group_count, int  state_dim, float  norm_eps  )

static

Definition at line 963 of file deltanet_kernels.c.

{
    (void)norm_eps;
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        num_heads <= 0 || group_count <= 0 || num_heads % group_count != 0 ||
        state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM) {
        return;
    }
 
    const size_t vector_stride = (size_t)state_dim;
    const size_t state_stride = vector_stride * vector_stride;
    const int heads_per_group = num_heads / group_count;
    const float sqrt_dim = sqrtf((float)state_dim);
    float gate_values[CK_DELTANET_MAX_STACK_DIM];
    float beta_values[CK_DELTANET_MAX_STACK_DIM];
    float memory[CK_DELTANET_MAX_STACK_DIM];
    float delta[CK_DELTANET_MAX_STACK_DIM];
    float q_scaled[CK_DELTANET_MAX_STACK_DIM];
    ck_deltanet_pytorch_gate_values(
        g, beta, gate_values, beta_values, num_heads);
 
    for (int h = 0; h < num_heads; ++h) {
        const int group = h / heads_per_group;
        const float *q_head = q + (size_t)group * vector_stride;
        const float *k_head = k + (size_t)group * vector_stride;
        const float *v_head = v + (size_t)h * vector_stride;
        const float *state_prev = state_in + (size_t)h * state_stride;
        float *state_cur = state_out + (size_t)h * state_stride;
        float *out_head = out + (size_t)h * vector_stride;
 
        for (int col = 0; col < state_dim; ++col) {
            q_scaled[col] = q_head[col] / sqrt_dim;
        }
 
        /* Materialize the decayed state before the separately ordered sum. */
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_offset = (size_t)row * vector_stride;
            int col = 0;
#if defined(__AVX2__)
            const __m256 gate8 = _mm256_set1_ps(gate_values[h]);
            for (; col + 7 < state_dim; col += 8) {
                const __m256 state = _mm256_mul_ps(
                    _mm256_loadu_ps(state_prev + row_offset + (size_t)col),
                    gate8);
                _mm256_storeu_ps(state_cur + row_offset + (size_t)col, state);
            }
#endif
            for (; col < state_dim; ++col) {
                const size_t offset = row_offset + (size_t)col;
                const float state = state_prev[offset] * gate_values[h];
                state_cur[offset] = state;
            }
        }
 
        ck_deltanet_pytorch_outer_sum(
            state_cur, k_head, memory, state_dim);
 
        if (debug_decayed_state) {
            memcpy(
                debug_decayed_state + (size_t)h * state_stride,
                state_cur,
                state_stride * sizeof(float));
        }
        if (debug_memory) {
            memcpy(
                debug_memory + (size_t)h * vector_stride,
                memory,
                vector_stride * sizeof(float));
        }
 
        for (int col = 0; col < state_dim; ++col) {
            delta[col] = (v_head[col] - memory[col]) * beta_values[h];
        }
        if (debug_delta) {
            memcpy(
                debug_delta + (size_t)h * vector_stride,
                delta,
                vector_stride * sizeof(float));
        }
 
        for (int row = 0; row < state_dim; ++row) {
            const size_t row_offset = (size_t)row * vector_stride;
            const float key = k_head[row];
            int col = 0;
#if defined(__AVX2__)
            const __m256 key8 = _mm256_set1_ps(key);
            for (; col + 7 < state_dim; col += 8) {
                const __m256 update = _mm256_mul_ps(
                    key8, _mm256_loadu_ps(delta + col));
                const __m256 state = _mm256_add_ps(
                    _mm256_loadu_ps(state_cur + row_offset + (size_t)col),
                    update);
                _mm256_storeu_ps(state_cur + row_offset + (size_t)col, state);
            }
#endif
            for (; col < state_dim; ++col) {
                const size_t offset = row_offset + (size_t)col;
                const float state = state_cur[offset] + key * delta[col];
                state_cur[offset] = state;
            }
        }
 
        ck_deltanet_pytorch_outer_sum(
            state_cur, q_scaled, out_head, state_dim);
 
        for (int col = 0; col < state_dim; ++col) {
            out_head[col] = bf16_to_float(float_to_bf16(out_head[col]));
        }
    }
}

References bf16_to_float(), CK_DELTANET_MAX_STACK_DIM, ck_deltanet_pytorch_gate_values(), ck_deltanet_pytorch_outer_sum(), and float_to_bf16().

Referenced by gated_deltanet_pytorch_grouped_bf16_forward(), and gated_deltanet_pytorch_grouped_bf16_forward_debug().

gated_deltanet_pytorch_grouped_bf16_prefill_forward()

void gated_deltanet_pytorch_grouped_bf16_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1243 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || state_dim <= 0) {
        return;
    }
    const size_t qk_stride = (size_t)group_count * (size_t)state_dim;
    const size_t value_stride = (size_t)num_heads * (size_t)state_dim;
    const size_t gate_stride = (size_t)num_heads;
    for (int row = 0; row < rows; ++row) {
        gated_deltanet_pytorch_grouped_bf16_forward(
            q + (size_t)row * qk_stride,
            k + (size_t)row * qk_stride,
            v + (size_t)row * value_stride,
            g + (size_t)row * gate_stride,
            beta + (size_t)row * gate_stride,
            row == 0 ? state_in : state_out,
            state_out,
            out + (size_t)row * value_stride,
            num_heads, group_count, state_dim, norm_eps);
    }
}

References gated_deltanet_pytorch_grouped_bf16_forward().

Variable Documentation

ck_deltanet_libm_handle

void* ck_deltanet_libm_handle = NULL

static

Definition at line 64 of file deltanet_kernels.c.

Referenced by ck_bind_deltanet_llama_libm().

ck_deltanet_libm_once

pthread_once_t ck_deltanet_libm_once = PTHREAD_ONCE_INIT

static

Definition at line 65 of file deltanet_kernels.c.

Referenced by ck_deltanet_llama_sigmoidf(), gated_deltanet_llama_avx2_grouped_forward_impl(), and gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

ck_deltanet_llama_expf

ck_deltanet_libm_f32_fn ck_deltanet_llama_expf = NULL

static

Definition at line 63 of file deltanet_kernels.c.

Referenced by ck_bind_deltanet_llama_libm(), ck_deltanet_llama_sigmoidf(), gated_deltanet_llama_avx2_grouped_forward_impl(), and gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

ck_deltanet_mkl_handle

void* ck_deltanet_mkl_handle = NULL

static

Definition at line 101 of file deltanet_kernels.c.

Referenced by ck_bind_deltanet_pytorch_primitives().

ck_deltanet_pytorch_primitives_once

pthread_once_t ck_deltanet_pytorch_primitives_once = PTHREAD_ONCE_INIT

static

Definition at line 102 of file deltanet_kernels.c.

Referenced by ck_deltanet_pytorch_gate_values().

ck_deltanet_pytorch_vsexp

ck_deltanet_mkl_vsexp_fn ck_deltanet_pytorch_vsexp = NULL

static

Definition at line 100 of file deltanet_kernels.c.

Referenced by ck_bind_deltanet_pytorch_primitives(), and ck_deltanet_pytorch_gate_values().