← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
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 _GNU_SOURCE
 
#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

◆ _GNU_SOURCE

#define _GNU_SOURCE

Definition at line 38 of file deltanet_kernels.c.

◆ CK_DELTANET_LLAMA_CHUNK_MAX_DIM

#define CK_DELTANET_LLAMA_CHUNK_MAX_DIM   256

Definition at line 58 of file deltanet_kernels.c.

◆ CK_DELTANET_LLAMA_CHUNK_SIZE

#define CK_DELTANET_LLAMA_CHUNK_SIZE   64

Definition at line 57 of file deltanet_kernels.c.

◆ CK_DELTANET_MAX_STACK_DIM

#define CK_DELTANET_MAX_STACK_DIM   4096

Definition at line 56 of file deltanet_kernels.c.

◆ CK_DELTANET_NOINLINE

#define CK_DELTANET_NOINLINE

Definition at line 63 of file deltanet_kernels.c.

Typedef Documentation

◆ ck_deltanet_libm_f32_fn

typedef float(* ck_deltanet_libm_f32_fn) (float)

Definition at line 66 of file deltanet_kernels.c.

◆ ck_deltanet_mkl_vsexp_fn

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

Definition at line 103 of file deltanet_kernels.c.

Function Documentation

◆ ck_bind_deltanet_llama_libm()

static void ck_bind_deltanet_llama_libm ( void  )
static

Definition at line 71 of file deltanet_kernels.c.

72{
73 ck_deltanet_libm_handle = dlopen("libm.so.6", RTLD_NOW | RTLD_LOCAL);
77 }
79 fprintf(stderr,
80 "HARD KERNEL CONTRACT FAULT: llama.cpp DeltaNet requires "
81 "expf from libm.so.6\n");
82 abort();
83 }
84}
static ck_deltanet_libm_f32_fn ck_deltanet_llama_expf
static void * ck_deltanet_libm_handle
float(* ck_deltanet_libm_f32_fn)(float)

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 108 of file deltanet_kernels.c.

109{
110 const char *mkl_library = getenv("CK_MKL_LIBRARY");
111 if (mkl_library && *mkl_library) {
112 ck_deltanet_mkl_handle = dlopen(mkl_library, RTLD_NOW | RTLD_LOCAL);
116 ck_deltanet_mkl_handle, "vsExp");
117 }
118 } else {
120 (ck_deltanet_mkl_vsexp_fn)dlsym(RTLD_DEFAULT, "vsExp");
121 }
122
123#if defined(__AVX512F__)
124 const char *library = getenv("CK_SLEEF_LIBRARY");
125 if (library && *library) {
126 ck_deltanet_sleef_handle = dlopen(library, RTLD_NOW | RTLD_LOCAL);
127 if (ck_deltanet_sleef_handle) {
128 ck_deltanet_pytorch_expf16 =
129 (ck_deltanet_sleef_expf16_fn)dlsym(
130 ck_deltanet_sleef_handle, "Sleef_expf16_u10");
131 }
132 } else {
133 ck_deltanet_pytorch_expf16 =
134 (ck_deltanet_sleef_expf16_fn)dlsym(
135 RTLD_DEFAULT, "Sleef_expf16_u10");
136 }
137#endif
138}
#define RTLD_DEFAULT
static ck_deltanet_mkl_vsexp_fn ck_deltanet_pytorch_vsexp
void(* ck_deltanet_mkl_vsexp_fn)(int, const float *, float *)
static void * ck_deltanet_mkl_handle

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 862 of file deltanet_kernels.c.

863{
864 int result = 0;
865 int power = 1;
866 while (power < value) {
867 power <<= 1;
868 ++result;
869 }
870 return result;
871}

Referenced by ck_deltanet_pytorch_outer_sum().

◆ ck_deltanet_force_ref()

static int ck_deltanet_force_ref ( void  )
static

Definition at line 1934 of file deltanet_kernels.c.

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

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 91 of file deltanet_kernels.c.

92{
94 return 1.0f / (1.0f + ck_deltanet_llama_expf(-x));
95}
static pthread_once_t ck_deltanet_libm_once
static void ck_bind_deltanet_llama_libm(void)

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 140 of file deltanet_kernels.c.

145{
146 pthread_once(
150 fprintf(stderr,
151 "HARD KERNEL CONTRACT FAULT: PyTorch DeltaNet requires "
152 "MKL vsExp; set CK_MKL_LIBRARY\n");
153 abort();
154 }
155 ck_deltanet_pytorch_vsexp(num_heads, g, gate_values);
156
157 int h = 0;
158#if defined(__AVX512F__)
159 if (ck_deltanet_pytorch_expf16) {
160 const __m512 one = _mm512_set1_ps(1.0f);
161 for (; h + 15 < num_heads; h += 16) {
162 const __m512 bv = _mm512_loadu_ps(beta + h);
163 const __m512 beta_exp = ck_deltanet_pytorch_expf16(
164 _mm512_sub_ps(_mm512_setzero_ps(), bv));
165 const __m512 beta_sigmoid = _mm512_div_ps(
166 one, _mm512_add_ps(one, beta_exp));
167 float beta_lanes[16];
168 _mm512_storeu_ps(beta_lanes, beta_sigmoid);
169 for (int lane = 0; lane < 16; ++lane) {
170 beta_values[h + lane] = bf16_to_float(
171 float_to_bf16(beta_lanes[lane]));
172 }
173 }
174 }
175#endif
176 for (; h < num_heads; ++h) {
177 beta_values[h] = bf16_to_float(float_to_bf16(
178 ck_deltanet_sigmoidf(beta[h])));
179 }
180}
static uint16_t float_to_bf16(float f)
Definition bf16_utils.h:90
static float bf16_to_float(uint16_t v)
Definition bf16_utils.h:38
static void ck_bind_deltanet_pytorch_primitives(void)
static pthread_once_t ck_deltanet_pytorch_primitives_once
static float ck_deltanet_sigmoidf(float x)

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 876 of file deltanet_kernels.c.

880{
881 const int num_levels = 4;
882 int level_power = ck_deltanet_ceil_log2(state_dim) / num_levels;
883 if (level_power < 4) {
884 level_power = 4;
885 }
886 const int level_step = 1 << level_power;
887 const int level_mask = level_step - 1;
888 int col = 0;
889
890#if defined(__AVX512F__)
891 /* PyTorch vectorized_outer_sum reduces four adjacent vectors together. */
892 for (; col + 63 < state_dim; col += 64) {
893 __m512 acc[4][4];
894 for (int level = 0; level < num_levels; ++level) {
895 for (int block = 0; block < 4; ++block) {
896 acc[level][block] = _mm512_setzero_ps();
897 }
898 }
899
900 int i = 0;
901 for (; i + level_step <= state_dim;) {
902 for (int j = 0; j < level_step; ++j, ++i) {
903 const float *row = matrix + (size_t)i * (size_t)state_dim + col;
904 const __m512 weight = _mm512_set1_ps(row_weights[i]);
905 for (int block = 0; block < 4; ++block) {
906 const __m512 product = _mm512_mul_ps(
907 _mm512_loadu_ps(row + block * 16), weight);
908 acc[0][block] = _mm512_add_ps(acc[0][block], product);
909 }
910 }
911
912 for (int level = 1; level < num_levels; ++level) {
913 for (int block = 0; block < 4; ++block) {
914 acc[level][block] = _mm512_add_ps(
915 acc[level][block], acc[level - 1][block]);
916 acc[level - 1][block] = _mm512_setzero_ps();
917 }
918 const int mask = level_mask << (level * level_power);
919 if ((i & mask) != 0) {
920 break;
921 }
922 }
923 }
924
925 for (; i < state_dim; ++i) {
926 const float *row = matrix + (size_t)i * (size_t)state_dim + col;
927 const __m512 weight = _mm512_set1_ps(row_weights[i]);
928 for (int block = 0; block < 4; ++block) {
929 const __m512 product = _mm512_mul_ps(
930 _mm512_loadu_ps(row + block * 16), weight);
931 acc[0][block] = _mm512_add_ps(acc[0][block], product);
932 }
933 }
934
935 for (int level = 1; level < num_levels; ++level) {
936 for (int block = 0; block < 4; ++block) {
937 acc[0][block] = _mm512_add_ps(
938 acc[0][block], acc[level][block]);
939 }
940 }
941 for (int block = 0; block < 4; ++block) {
942 _mm512_storeu_ps(output + col + block * 16, acc[0][block]);
943 }
944 }
945#elif defined(__AVX2__)
946 for (; col + 31 < state_dim; col += 32) {
947 __m256 acc[4][4];
948 for (int level = 0; level < num_levels; ++level) {
949 for (int block = 0; block < 4; ++block) {
950 acc[level][block] = _mm256_setzero_ps();
951 }
952 }
953
954 int i = 0;
955 for (; i + level_step <= state_dim;) {
956 for (int j = 0; j < level_step; ++j, ++i) {
957 const float *row = matrix + (size_t)i * (size_t)state_dim + col;
958 const __m256 weight = _mm256_set1_ps(row_weights[i]);
959 for (int block = 0; block < 4; ++block) {
960 const __m256 product = _mm256_mul_ps(
961 _mm256_loadu_ps(row + block * 8), weight);
962 acc[0][block] = _mm256_add_ps(acc[0][block], product);
963 }
964 }
965 for (int level = 1; level < num_levels; ++level) {
966 for (int block = 0; block < 4; ++block) {
967 acc[level][block] = _mm256_add_ps(
968 acc[level][block], acc[level - 1][block]);
969 acc[level - 1][block] = _mm256_setzero_ps();
970 }
971 const int mask = level_mask << (level * level_power);
972 if ((i & mask) != 0) {
973 break;
974 }
975 }
976 }
977 for (; i < state_dim; ++i) {
978 const float *row = matrix + (size_t)i * (size_t)state_dim + col;
979 const __m256 weight = _mm256_set1_ps(row_weights[i]);
980 for (int block = 0; block < 4; ++block) {
981 const __m256 product = _mm256_mul_ps(
982 _mm256_loadu_ps(row + block * 8), weight);
983 acc[0][block] = _mm256_add_ps(acc[0][block], product);
984 }
985 }
986 for (int level = 1; level < num_levels; ++level) {
987 for (int block = 0; block < 4; ++block) {
988 acc[0][block] = _mm256_add_ps(
989 acc[0][block], acc[level][block]);
990 }
991 }
992 for (int block = 0; block < 4; ++block) {
993 _mm256_storeu_ps(output + col + block * 8, acc[0][block]);
994 }
995 }
996#endif
997
998 /* Production Qwen dimensions are covered above. Keep a deterministic
999 * scalar fallback for uncommon tail widths. */
1000 for (; col < state_dim; ++col) {
1001 float sum = 0.0f;
1002 for (int row = 0; row < state_dim; ++row) {
1003 sum += matrix[(size_t)row * (size_t)state_dim + col] *
1004 row_weights[row];
1005 }
1006 output[col] = sum;
1007 }
1008}
static int product(size_t a, size_t b, size_t *out)
static int ck_deltanet_ceil_log2(int value)
int32_t int32_t int32_t int32_t int32_t mask
Definition tokenizer.h:234

References ck_deltanet_ceil_log2(), mask, and product().

Referenced by gated_deltanet_pytorch_grouped_bf16_forward_impl().

◆ ck_deltanet_sigmoidf()

static float ck_deltanet_sigmoidf ( float  x)
inlinestatic

◆ 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 2038 of file deltanet_kernels.c.

2056{
2057 if (!d_out || !d_state_out || !q || !k || !v || !g || !beta || !state_in || !state_out ||
2058 !d_q || !d_k || !d_v || !d_g || !d_beta || !d_state_in) {
2059 return;
2060 }
2061 if (num_heads <= 0 || state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM) {
2062 return;
2063 }
2064
2066 d_out,
2067 d_state_out,
2068 q,
2069 k,
2070 v,
2071 g,
2072 beta,
2073 state_in,
2074 state_out,
2075 d_q,
2076 d_k,
2077 d_v,
2078 d_g,
2079 d_beta,
2080 d_state_in,
2081 num_heads,
2082 state_dim,
2083 norm_eps);
2084}
#define CK_DELTANET_MAX_STACK_DIM
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)

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 1389 of file deltanet_kernels.c.

1407{
1408 const float q_scale = 1.0f / sqrtf((float)state_dim);
1409 const size_t vec_stride = (size_t)state_dim;
1410 const size_t state_stride = (size_t)state_dim * (size_t)state_dim;
1411
1412 float q_hat[CK_DELTANET_MAX_STACK_DIM];
1413 float k_hat[CK_DELTANET_MAX_STACK_DIM];
1414 float kv_mem[CK_DELTANET_MAX_STACK_DIM];
1415 float delta[CK_DELTANET_MAX_STACK_DIM];
1416 float d_q_hat[CK_DELTANET_MAX_STACK_DIM];
1417 float d_k_hat[CK_DELTANET_MAX_STACK_DIM];
1418 float d_mem[CK_DELTANET_MAX_STACK_DIM];
1419
1420 for (int h = 0; h < num_heads; ++h) {
1421 const float *d_out_head = d_out + (size_t)h * vec_stride;
1422 const float *d_state_out_head = d_state_out + (size_t)h * state_stride;
1423 const float *q_head = q + (size_t)h * vec_stride;
1424 const float *k_head = k + (size_t)h * vec_stride;
1425 const float *v_head = v + (size_t)h * vec_stride;
1426 const float *state_prev = state_in + (size_t)h * state_stride;
1427 const float *state_cur = state_out + (size_t)h * state_stride;
1428 float *d_q_head = d_q + (size_t)h * vec_stride;
1429 float *d_k_head = d_k + (size_t)h * vec_stride;
1430 float *d_v_head = d_v + (size_t)h * vec_stride;
1431 float *d_state_prev = d_state_in + (size_t)h * state_stride;
1432
1433 const float beta_s = ck_deltanet_sigmoidf(beta[h]);
1434 const float gate = expf(g[h]);
1435
1436 float qk_dot = 0.0f;
1437 float out_delta_dot = 0.0f;
1438 float beta_acc = 0.0f;
1439 float gate_acc = 0.0f;
1440
1441 for (int i = 0; i < state_dim; ++i) {
1442 q_hat[i] = q_head[i] * q_scale;
1443 k_hat[i] = k_head[i];
1444 kv_mem[i] = 0.0f;
1445 d_q_hat[i] = 0.0f;
1446 d_k_hat[i] = 0.0f;
1447 d_mem[i] = 0.0f;
1448 d_v_head[i] = 0.0f;
1449 qk_dot += q_hat[i] * k_hat[i];
1450 }
1451
1452 for (int col = 0; col < state_dim; ++col) {
1453 float mem = 0.0f;
1454 for (int row = 0; row < state_dim; ++row) {
1455 mem += (state_prev[(size_t)row * (size_t)state_dim + (size_t)col] * gate) * k_hat[row];
1456 }
1457 kv_mem[col] = mem;
1458 delta[col] = (v_head[col] - mem) * beta_s;
1459 out_delta_dot += d_out_head[col] * delta[col];
1460 }
1461
1462 for (int row = 0; row < state_dim; ++row) {
1463 const size_t row_off = (size_t)row * (size_t)state_dim;
1464 float dq_acc = 0.0f;
1465 float dk_acc = q_hat[row] * out_delta_dot;
1466 for (int col = 0; col < state_dim; ++col) {
1467 const float d_state_direct = d_state_out_head[row_off + (size_t)col];
1468 dq_acc += state_cur[row_off + (size_t)col] * d_out_head[col];
1469 dk_acc += d_state_direct * delta[col];
1470 }
1471 d_q_hat[row] = dq_acc;
1472 d_k_hat[row] = dk_acc;
1473 }
1474
1475 for (int col = 0; col < state_dim; ++col) {
1476 float d_delta_acc = d_out_head[col] * qk_dot;
1477 for (int row = 0; row < state_dim; ++row) {
1478 d_delta_acc += d_state_out_head[(size_t)row * (size_t)state_dim + (size_t)col] * k_hat[row];
1479 }
1480
1481 d_v_head[col] = beta_s * d_delta_acc;
1482 d_mem[col] = -beta_s * d_delta_acc;
1483 beta_acc += d_delta_acc * (v_head[col] - kv_mem[col]);
1484 }
1485
1486 for (int row = 0; row < state_dim; ++row) {
1487 const size_t row_off = (size_t)row * (size_t)state_dim;
1488 float s_dm_acc = 0.0f;
1489 for (int col = 0; col < state_dim; ++col) {
1490 s_dm_acc += (state_prev[row_off + (size_t)col] * gate) * d_mem[col];
1491 }
1492 d_k_hat[row] += s_dm_acc;
1493 }
1494
1495 for (int row = 0; row < state_dim; ++row) {
1496 const size_t row_off = (size_t)row * (size_t)state_dim;
1497 for (int col = 0; col < state_dim; ++col) {
1498 const float d_state_total = d_state_out_head[row_off + (size_t)col]
1499 + q_hat[row] * d_out_head[col]
1500 + k_hat[row] * d_mem[col];
1501 d_state_prev[row_off + (size_t)col] = gate * d_state_total;
1502 gate_acc += d_state_total * state_prev[row_off + (size_t)col];
1503 }
1504 }
1505
1506 for (int i = 0; i < state_dim; ++i) {
1507 d_q_head[i] = d_q_hat[i] * q_scale;
1508 d_k_head[i] = d_k_hat[i];
1509 }
1510
1511 d_g[h] = gate_acc * gate;
1512 d_beta[h] = beta_acc * beta_s * (1.0f - beta_s);
1513 }
1514}

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 1956 of file deltanet_kernels.c.

1967{
1968 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
1969 return;
1970 }
1971 if (num_heads <= 0 || state_dim <= 0) {
1972 return;
1973 }
1974
1975 /*
1976 * q and k arrive pre-normalized by recurrent_qk_l2_norm, so the
1977 * ISA-specialized kernels can follow the same contract as the scalar ref.
1978 */
1981 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1982 return;
1983 }
1984#if defined(__AVX512F__)
1985 gated_deltanet_autoregressive_forward_avx512(
1986 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1987#elif defined(__AVX2__)
1988 gated_deltanet_autoregressive_forward_avx2(
1989 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1990#elif defined(__AVX__)
1991 gated_deltanet_autoregressive_forward_avx(
1992 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1993#else
1995 q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
1996#endif
1997}
int ck_strict_parity_enabled(void)
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)
static int ck_deltanet_force_ref(void)

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 1330 of file deltanet_kernels.c.

1341{
1342 const float q_scale = 1.0f / sqrtf((float)state_dim);
1343 const size_t vec_stride = (size_t)state_dim;
1344 const size_t state_stride = (size_t)state_dim * (size_t)state_dim;
1345
1346 for (int h = 0; h < num_heads; ++h) {
1347 const float *q_head = q + (size_t)h * vec_stride;
1348 const float *k_head = k + (size_t)h * vec_stride;
1349 const float *v_head = v + (size_t)h * vec_stride;
1350 const float *state_prev = state_in + (size_t)h * state_stride;
1351 float *state_cur = state_out + (size_t)h * state_stride;
1352 float *out_head = out + (size_t)h * vec_stride;
1353
1354 const float beta_s = ck_deltanet_sigmoidf(beta[h]);
1355 const float gate = expf(g[h]);
1356
1357 for (int row = 0; row < state_dim; ++row) {
1358 const size_t row_off = (size_t)row * (size_t)state_dim;
1359 for (int col = 0; col < state_dim; ++col) {
1360 state_cur[row_off + (size_t)col] = state_prev[row_off + (size_t)col] * gate;
1361 }
1362 }
1363
1364 for (int col = 0; col < state_dim; ++col) {
1365 float kv_mem = 0.0f;
1366 for (int row = 0; row < state_dim; ++row) {
1367 const float k_hat = k_head[row];
1368 kv_mem += state_cur[(size_t)row * (size_t)state_dim + (size_t)col] * k_hat;
1369 }
1370
1371 const float delta = (v_head[col] - kv_mem) * beta_s;
1372 for (int row = 0; row < state_dim; ++row) {
1373 const float k_hat = k_head[row];
1374 state_cur[(size_t)row * (size_t)state_dim + (size_t)col] += k_hat * delta;
1375 }
1376 }
1377
1378 for (int col = 0; col < state_dim; ++col) {
1379 float acc = 0.0f;
1380 for (int row = 0; row < state_dim; ++row) {
1381 const float q_hat = q_head[row] * q_scale;
1382 acc += state_cur[(size_t)row * (size_t)state_dim + (size_t)col] * q_hat;
1383 }
1384 out_head[col] = acc;
1385 }
1386 }
1387}

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 1940 of file deltanet_kernels.c.

1941{
1943 return "REF";
1944 }
1945#if defined(__AVX512F__)
1946 return "AVX512";
1947#elif defined(__AVX2__)
1948 return "AVX2";
1949#elif defined(__AVX__)
1950 return "AVX";
1951#else
1952 return "REF";
1953#endif
1954}

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 817 of file deltanet_kernels.c.

829{
831 q, k, v, g, beta, state_in, state_out, out,
832 num_heads, group_count, state_dim, norm_eps, 0, num_heads);
833}
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)

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 840 of file deltanet_kernels.c.

855{
857 q, k, v, g, beta, state_in, state_out, out,
858 num_heads, group_count, state_dim, norm_eps,
859 head_begin, head_end);
860}

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 578 of file deltanet_kernels.c.

594{
595#if defined(__AVX2__)
596 (void) norm_eps;
597 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
598 num_heads <= 0 || num_heads > CK_DELTANET_MAX_STACK_DIM ||
599 group_count <= 0 || num_heads % group_count != 0 ||
600 state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM ||
601 head_begin < 0 || head_end < head_begin || head_end > num_heads) {
602 return;
603 }
604 const float scale = ck_deltanet_llama_scale(state_dim);
605 const size_t vector_stride = (size_t) state_dim;
606 const size_t state_stride = (size_t) state_dim * (size_t) state_dim;
607 float column[CK_DELTANET_MAX_STACK_DIM];
608 float q_scaled[CK_DELTANET_MAX_STACK_DIM];
609
610 for (int h = head_begin; h < head_end; ++h) {
611 /* llama.cpp ggml_repeat_4d tiles compact Q/K heads (0..G-1 repeated),
612 * while the PyTorch Qwen3-Next reference uses repeat_interleave so
613 * each compact head owns H/G adjacent value heads. These layouts
614 * are distinct numerical contracts even though their buffer shapes
615 * are identical. */
616 const int group = pytorch_bf16_boundaries
617 ? h / (num_heads / group_count)
618 : h % group_count;
619 const float *q_head = q + (size_t) group * vector_stride;
620 const float *k_head = k + (size_t) group * vector_stride;
621 const float *v_head = v + (size_t) h * vector_stride;
622 const float *state_prev = state_in + (size_t) h * state_stride;
623 float *state_cur = state_out + (size_t) h * state_stride;
624 float *out_head = out + (size_t) h * vector_stride;
625 float gate;
626 float beta_s;
627 if (pytorch_bf16_boundaries) {
628 gate = expf(g[h]);
629 beta_s = ck_deltanet_sigmoidf(beta[h]);
630 } else {
632 gate = ck_deltanet_llama_expf(g[h]);
633 beta_s = ck_deltanet_llama_sigmoidf(beta[h]);
634 }
635 if (pytorch_bf16_boundaries) {
636 beta_s = bf16_to_float(float_to_bf16(beta_s));
637 for (int row = 0; row < state_dim; ++row) {
638 q_scaled[row] = q_head[row] * scale;
639 }
640 }
641
642 for (int row = 0; row < state_dim; ++row) {
643 const size_t row_offset = (size_t) row * (size_t) state_dim;
644 for (int col = 0; col < state_dim; ++col) {
645 state_cur[row_offset + (size_t) col] =
646 state_prev[row_offset + (size_t) col] * gate;
647 }
648 }
649
650 for (int col = 0; col < state_dim; ++col) {
651 for (int row = 0; row < state_dim; ++row) {
652 column[row] = state_cur[(size_t) row * (size_t) state_dim + (size_t) col];
653 }
654 const float memory = ck_deltanet_llama_avx2_dot(column, k_head, state_dim);
655 const float delta = (v_head[col] - memory) * beta_s;
656 for (int row = 0; row < state_dim; ++row) {
657 const size_t offset = (size_t) row * (size_t) state_dim + (size_t) col;
658#if defined(__FMA__)
659 const float updated = fmaf(k_head[row], delta, state_cur[offset]);
660#else
661 const float updated = state_cur[offset] + k_head[row] * delta;
662#endif
663 state_cur[offset] = updated;
664 column[row] = updated;
665 }
666 if (pytorch_bf16_boundaries) {
667 out_head[col] =
668 ck_deltanet_llama_avx2_dot(column, q_scaled, state_dim);
669 } else {
670 out_head[col] =
671 ck_deltanet_llama_avx2_dot(column, q_head, state_dim) * scale;
672 }
673 }
674 }
675#else
676 if (group_count != num_heads) {
677 return;
678 }
680 q, k, v, g, beta, state_in, state_out, out,
681 num_heads, state_dim, norm_eps);
682#endif
683}
static float ck_deltanet_llama_sigmoidf(float x)

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 691 of file deltanet_kernels.c.

706{
707#if defined(__AVX2__)
708 (void) norm_eps;
709 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
710 num_heads <= 0 || group_count <= 0 || num_heads % group_count != 0 ||
711 state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM ||
712 head_begin < 0 || head_end < head_begin || head_end > num_heads) {
713 return;
714 }
715 const float scale = ck_deltanet_llama_scale(state_dim);
716 const size_t vector_stride = (size_t) state_dim;
717 const size_t state_stride = (size_t) state_dim * (size_t) state_dim;
718
720 for (int h = head_begin; h < head_end; ++h) {
721 const int group = h % group_count;
722 const float *q_head = q + (size_t) group * vector_stride;
723 const float *k_head = k + (size_t) group * vector_stride;
724 const float *v_head = v + (size_t) h * vector_stride;
725 const float *state_prev = state_in + (size_t) h * state_stride;
726 float *state_cur = state_out + (size_t) h * state_stride;
727 float *out_head = out + (size_t) h * vector_stride;
728 const float gate = ck_deltanet_llama_expf(g[h]);
729 const float beta_s = ck_deltanet_llama_sigmoidf(beta[h]);
730
731 for (int col = 0; col < state_dim; ++col) {
732 const float *prev_col = state_prev + (size_t) col * vector_stride;
733 float *cur_col = state_cur + (size_t) col * vector_stride;
734 int row = 0;
735 const __m256 gate8 = _mm256_set1_ps(gate);
736 for (; row + 7 < state_dim; row += 8) {
737 const __m256 scaled = _mm256_mul_ps(
738 _mm256_loadu_ps(prev_col + row), gate8);
739 _mm256_storeu_ps(cur_col + row, scaled);
740 }
741 for (; row < state_dim; ++row) {
742 cur_col[row] = prev_col[row] * gate;
743 }
744
745 const float memory =
746 ck_deltanet_llama_avx2_dot(cur_col, k_head, state_dim);
747 const float delta = (v_head[col] - memory) * beta_s;
748 row = 0;
749 const __m256 delta8 = _mm256_set1_ps(delta);
750 for (; row + 7 < state_dim; row += 8) {
751 const __m256 updated = _mm256_fmadd_ps(
752 _mm256_loadu_ps(k_head + row), delta8,
753 _mm256_loadu_ps(cur_col + row));
754 _mm256_storeu_ps(cur_col + row, updated);
755 }
756 for (; row < state_dim; ++row) {
757 cur_col[row] = fmaf(k_head[row], delta, cur_col[row]);
758 }
759 out_head[col] =
760 ck_deltanet_llama_avx2_dot(cur_col, q_head, state_dim) * scale;
761 }
762 }
763#else
764 /* Portable scalar path (ARM and other non-AVX2 targets): identical
765 * transposed state layout and update order as the AVX2 path above.
766 * Without this the recurrent core silently no-ops on ARM and qwen35-
767 * class models emit garbage (state never updates, out stays stale). */
768 (void) norm_eps;
769 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
770 num_heads <= 0 || group_count <= 0 || num_heads % group_count != 0 ||
771 state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM ||
772 head_begin < 0 || head_end < head_begin || head_end > num_heads) {
773 return;
774 }
775 /* ck_deltanet_llama_scale() is inside an AVX2-only block; use its
776 * portable non-SSE equivalent here. */
777 const float scale = 1.0f / sqrtf((float) state_dim);
778 const size_t vector_stride = (size_t) state_dim;
779 const size_t state_stride = (size_t) state_dim * (size_t) state_dim;
780
782 for (int h = head_begin; h < head_end; ++h) {
783 const int group = h % group_count;
784 const float *q_head = q + (size_t) group * vector_stride;
785 const float *k_head = k + (size_t) group * vector_stride;
786 const float *v_head = v + (size_t) h * vector_stride;
787 const float *state_prev = state_in + (size_t) h * state_stride;
788 float *state_cur = state_out + (size_t) h * state_stride;
789 float *out_head = out + (size_t) h * vector_stride;
790 const float gate = ck_deltanet_llama_expf(g[h]);
791 const float beta_s = ck_deltanet_llama_sigmoidf(beta[h]);
792
793 for (int col = 0; col < state_dim; ++col) {
794 const float *prev_col = state_prev + (size_t) col * vector_stride;
795 float *cur_col = state_cur + (size_t) col * vector_stride;
796 float memory = 0.0f;
797 float dot = 0.0f;
798 for (int row = 0; row < state_dim; ++row) {
799 cur_col[row] = prev_col[row] * gate;
800 }
801 for (int row = 0; row < state_dim; ++row) {
802 memory += cur_col[row] * k_head[row];
803 }
804 const float delta = (v_head[col] - memory) * beta_s;
805 for (int row = 0; row < state_dim; ++row) {
806 cur_col[row] = fmaf(k_head[row], delta, cur_col[row]);
807 }
808 for (int row = 0; row < state_dim; ++row) {
809 dot += cur_col[row] * q_head[row];
810 }
811 out_head[col] = dot * scale;
812 }
813 }
814#endif
815}

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 1181 of file deltanet_kernels.c.

1194{
1195 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
1196 rows <= 0 || num_heads <= 0 || group_count <= 0 ||
1197 num_heads % group_count != 0 || state_dim <= 0) {
1198 return;
1199 }
1200 const size_t qk_stride = (size_t) group_count * (size_t) state_dim;
1201 const size_t value_stride = (size_t) num_heads * (size_t) state_dim;
1202 const size_t gate_stride = (size_t) num_heads;
1203 for (int row = 0; row < rows; ++row) {
1205 q + (size_t) row * qk_stride,
1206 k + (size_t) row * qk_stride,
1207 v + (size_t) row * value_stride,
1208 g + (size_t) row * gate_stride,
1209 beta + (size_t) row * gate_stride,
1210 row == 0 ? state_in : state_out,
1211 state_out,
1212 out + (size_t) row * value_stride,
1213 num_heads, group_count, state_dim, norm_eps);
1214 }
1215}
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)

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 1252 of file deltanet_kernels.c.

1265{
1266#if defined(__AVX2__)
1267 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
1268 rows <= 0 || num_heads <= 0 || group_count <= 0 ||
1269 num_heads % group_count != 0 || head < 0 || head >= num_heads ||
1270 state_dim <= 0 || state_dim > CK_DELTANET_LLAMA_CHUNK_MAX_DIM) {
1271 return;
1272 }
1273 gated_deltanet_llama_chunk64_head(
1274 q, k, v, g, beta, state_in, state_out, out,
1275 rows, num_heads, group_count, head, state_dim);
1276#else
1277 (void)q;
1278 (void)k;
1279 (void)v;
1280 (void)g;
1281 (void)beta;
1282 (void)state_in;
1283 (void)state_out;
1284 (void)out;
1285 (void)rows;
1286 (void)num_heads;
1287 (void)group_count;
1288 (void)head;
1289 (void)state_dim;
1290#endif
1291}
#define CK_DELTANET_LLAMA_CHUNK_MAX_DIM

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 1217 of file deltanet_kernels.c.

1230{
1231 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
1232 rows <= 0 || num_heads <= 0 || group_count <= 0 ||
1233 num_heads % group_count != 0 || state_dim <= 0) {
1234 return;
1235 }
1236#if defined(__AVX2__)
1237 (void)norm_eps;
1238 if (state_dim <= CK_DELTANET_LLAMA_CHUNK_MAX_DIM) {
1239 for (int head = 0; head < num_heads; ++head) {
1240 gated_deltanet_llama_chunk64_head(
1241 q, k, v, g, beta, state_in, state_out, out,
1242 rows, num_heads, group_count, head, state_dim);
1243 }
1244 return;
1245 }
1246#endif
1248 q, k, v, g, beta, state_in, state_out, out,
1249 rows, num_heads, group_count, state_dim, norm_eps);
1250}
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)

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 1999 of file deltanet_kernels.c.

2011{
2012 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
2013 return;
2014 }
2015 if (rows <= 0 || num_heads <= 0 || state_dim <= 0) {
2016 return;
2017 }
2018
2019 const size_t vector_stride = (size_t)num_heads * (size_t)state_dim;
2020 const size_t gate_stride = (size_t)num_heads;
2021 for (int row = 0; row < rows; ++row) {
2022 const float *row_state_in = row == 0 ? state_in : state_out;
2024 q + (size_t)row * vector_stride,
2025 k + (size_t)row * vector_stride,
2026 v + (size_t)row * vector_stride,
2027 g + (size_t)row * gate_stride,
2028 beta + (size_t)row * gate_stride,
2029 row_state_in,
2030 state_out,
2031 out + (size_t)row * vector_stride,
2032 num_heads,
2033 state_dim,
2034 norm_eps);
2035 }
2036}
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)

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 182 of file deltanet_kernels.c.

187{
188 if (!g || !beta || !gate_values || !beta_values || num_heads <= 0 ||
189 num_heads > CK_DELTANET_MAX_STACK_DIM) {
190 return;
191 }
193 g, beta, gate_values, beta_values, num_heads);
194}
static void ck_deltanet_pytorch_gate_values(const float *g, const float *beta, float *gate_values, float *beta_values, int 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 1139 of file deltanet_kernels.c.

1151{
1153 q, k, v, g, beta, state_in, state_out, out,
1154 NULL, NULL, NULL,
1155 num_heads, group_count, state_dim, norm_eps);
1156}
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)

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 1158 of file deltanet_kernels.c.

1174{
1176 q, k, v, g, beta, state_in, state_out, out,
1177 decayed_state, memory, delta,
1178 num_heads, group_count, state_dim, norm_eps);
1179}

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 1013 of file deltanet_kernels.c.

1029{
1030 (void)norm_eps;
1031 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
1032 num_heads <= 0 || group_count <= 0 || num_heads % group_count != 0 ||
1033 state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM) {
1034 return;
1035 }
1036
1037 const size_t vector_stride = (size_t)state_dim;
1038 const size_t state_stride = vector_stride * vector_stride;
1039 const int heads_per_group = num_heads / group_count;
1040 const float sqrt_dim = sqrtf((float)state_dim);
1041 float gate_values[CK_DELTANET_MAX_STACK_DIM];
1042 float beta_values[CK_DELTANET_MAX_STACK_DIM];
1043 float memory[CK_DELTANET_MAX_STACK_DIM];
1044 float delta[CK_DELTANET_MAX_STACK_DIM];
1045 float q_scaled[CK_DELTANET_MAX_STACK_DIM];
1047 g, beta, gate_values, beta_values, num_heads);
1048
1049 for (int h = 0; h < num_heads; ++h) {
1050 const int group = h / heads_per_group;
1051 const float *q_head = q + (size_t)group * vector_stride;
1052 const float *k_head = k + (size_t)group * vector_stride;
1053 const float *v_head = v + (size_t)h * vector_stride;
1054 const float *state_prev = state_in + (size_t)h * state_stride;
1055 float *state_cur = state_out + (size_t)h * state_stride;
1056 float *out_head = out + (size_t)h * vector_stride;
1057
1058 for (int col = 0; col < state_dim; ++col) {
1059 q_scaled[col] = q_head[col] / sqrt_dim;
1060 }
1061
1062 /* Materialize the decayed state before the separately ordered sum. */
1063 for (int row = 0; row < state_dim; ++row) {
1064 const size_t row_offset = (size_t)row * vector_stride;
1065 int col = 0;
1066#if defined(__AVX2__)
1067 const __m256 gate8 = _mm256_set1_ps(gate_values[h]);
1068 for (; col + 7 < state_dim; col += 8) {
1069 const __m256 state = _mm256_mul_ps(
1070 _mm256_loadu_ps(state_prev + row_offset + (size_t)col),
1071 gate8);
1072 _mm256_storeu_ps(state_cur + row_offset + (size_t)col, state);
1073 }
1074#endif
1075 for (; col < state_dim; ++col) {
1076 const size_t offset = row_offset + (size_t)col;
1077 const float state = state_prev[offset] * gate_values[h];
1078 state_cur[offset] = state;
1079 }
1080 }
1081
1083 state_cur, k_head, memory, state_dim);
1084
1085 if (debug_decayed_state) {
1086 memcpy(
1087 debug_decayed_state + (size_t)h * state_stride,
1088 state_cur,
1089 state_stride * sizeof(float));
1090 }
1091 if (debug_memory) {
1092 memcpy(
1093 debug_memory + (size_t)h * vector_stride,
1094 memory,
1095 vector_stride * sizeof(float));
1096 }
1097
1098 for (int col = 0; col < state_dim; ++col) {
1099 delta[col] = (v_head[col] - memory[col]) * beta_values[h];
1100 }
1101 if (debug_delta) {
1102 memcpy(
1103 debug_delta + (size_t)h * vector_stride,
1104 delta,
1105 vector_stride * sizeof(float));
1106 }
1107
1108 for (int row = 0; row < state_dim; ++row) {
1109 const size_t row_offset = (size_t)row * vector_stride;
1110 const float key = k_head[row];
1111 int col = 0;
1112#if defined(__AVX2__)
1113 const __m256 key8 = _mm256_set1_ps(key);
1114 for (; col + 7 < state_dim; col += 8) {
1115 const __m256 update = _mm256_mul_ps(
1116 key8, _mm256_loadu_ps(delta + col));
1117 const __m256 state = _mm256_add_ps(
1118 _mm256_loadu_ps(state_cur + row_offset + (size_t)col),
1119 update);
1120 _mm256_storeu_ps(state_cur + row_offset + (size_t)col, state);
1121 }
1122#endif
1123 for (; col < state_dim; ++col) {
1124 const size_t offset = row_offset + (size_t)col;
1125 const float state = state_cur[offset] + key * delta[col];
1126 state_cur[offset] = state;
1127 }
1128 }
1129
1131 state_cur, q_scaled, out_head, state_dim);
1132
1133 for (int col = 0; col < state_dim; ++col) {
1134 out_head[col] = bf16_to_float(float_to_bf16(out_head[col]));
1135 }
1136 }
1137}
static void ck_deltanet_pytorch_outer_sum(const float *matrix, const float *row_weights, float *output, int state_dim)

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 1293 of file deltanet_kernels.c.

1307{
1308 if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
1309 rows <= 0 || num_heads <= 0 || group_count <= 0 ||
1310 num_heads % group_count != 0 || state_dim <= 0) {
1311 return;
1312 }
1313 const size_t qk_stride = (size_t)group_count * (size_t)state_dim;
1314 const size_t value_stride = (size_t)num_heads * (size_t)state_dim;
1315 const size_t gate_stride = (size_t)num_heads;
1316 for (int row = 0; row < rows; ++row) {
1318 q + (size_t)row * qk_stride,
1319 k + (size_t)row * qk_stride,
1320 v + (size_t)row * value_stride,
1321 g + (size_t)row * gate_stride,
1322 beta + (size_t)row * gate_stride,
1323 row == 0 ? state_in : state_out,
1324 state_out,
1325 out + (size_t)row * value_stride,
1326 num_heads, group_count, state_dim, norm_eps);
1327 }
1328}
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)

References gated_deltanet_pytorch_grouped_bf16_forward().

Variable Documentation

◆ ck_deltanet_libm_handle

void* ck_deltanet_libm_handle = NULL
static

Definition at line 68 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

◆ ck_deltanet_llama_expf

◆ ck_deltanet_mkl_handle

void* ck_deltanet_mkl_handle = NULL
static

Definition at line 105 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 106 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