hybrid_attention_kernels.c File Reference

#include "ckernel_engine.h"
#include "bf16_utils.h"
#include "ck_threadpool.h"
#include <dlfcn.h>
#include <limits.h>
#include <math.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

Go to the source code of this file.

Typedefs
typedef float(*	ck_hybrid_libm_f32_fn) (float)

Functions
void	attn_gate_sigmoid_mul_backward (const float *d_out, const float *x, const float *gate, float *d_x, float *d_gate, int rows, int num_heads, int state_dim)
void	attn_gate_sigmoid_mul_forward (const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
void	attn_gate_sigmoid_mul_pytorch_bf16_storage (const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
static void	attn_gate_sigmoid_mul_row_range (int begin, int end, void *opaque)
void	attn_gate_softplus_mul_forward (const float *x, const float *gate, float *out, int rows, int num_heads, int state_dim)
static void	ck_bind_hybrid_llama_libm (void)
static int	hybrid_ranges_overlap (const float *lhs, const float *rhs, size_t elements)
static float	hybrid_sigmoid (float x)
void	split_q_gate_backward (const float *d_q, const float *d_gate, float *d_packed_qg, int rows, int q_dim, int gate_dim, int group_dim)
void	split_q_gate_forward (const float *packed_qg, float *q, float *gate, int rows, int q_dim, int gate_dim, int group_dim)

Variables
static void *	ck_hybrid_libm_handle = NULL
static pthread_once_t	ck_hybrid_libm_once = PTHREAD_ONCE_INIT
static ck_hybrid_libm_f32_fn	ck_hybrid_llama_expf = NULL

Typedef Documentation

ck_hybrid_libm_f32_fn

typedef float(* ck_hybrid_libm_f32_fn) (float)

Definition at line 14 of file hybrid_attention_kernels.c.

Function Documentation

attn_gate_sigmoid_mul_backward()

void attn_gate_sigmoid_mul_backward	(	const float *	d_out,
		const float *	x,
		const float *	gate,
		float *	d_x,
		float *	d_gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 253 of file hybrid_attention_kernels.c.

                                                   {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *d_out_row = d_out + (size_t) row * (size_t) dim;
        const float *x_row = x + (size_t) row * (size_t) dim;
        const float *gate_row = gate + (size_t) row * (size_t) dim;
        float *d_x_row = d_x + (size_t) row * (size_t) dim;
        float *d_gate_row = d_gate + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float sig = hybrid_sigmoid(gate_row[col]);
            d_x_row[col] = d_out_row[col] * sig;
            d_gate_row[col] = d_out_row[col] * x_row[col] * sig * (1.0f - sig);
        }
    }
}

References hybrid_sigmoid().

attn_gate_sigmoid_mul_forward()

void attn_gate_sigmoid_mul_forward	(	const float *	x,
		const float *	gate,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 154 of file hybrid_attention_kernels.c.

                                                  {
    if (!x || !gate || !out || rows <= 0 || num_heads <= 0 ||
        state_dim <= 0 || num_heads > INT_MAX / state_dim) {
        return;
    }
    const int dim = num_heads * state_dim;
    const size_t elements = (size_t)rows * (size_t)dim;
    ck_attn_gate_sigmoid_mul_args_t args = {x, gate, out, dim};
    const size_t minimum_elements_per_thread = 32768u;
 
    const int partial_x_overlap = out != x &&
        hybrid_ranges_overlap(out, x, elements);
    const int gate_overlap = hybrid_ranges_overlap(out, gate, elements);
    if (partial_x_overlap || gate_overlap ||
        elements < 2u * minimum_elements_per_thread) {
        attn_gate_sigmoid_mul_row_range(0, rows, &args);
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    size_t useful_threads =
        (elements + minimum_elements_per_thread - 1u) /
        minimum_elements_per_thread;
    if ((size_t)active > useful_threads) {
        active = (int)useful_threads;
    }
    if (active > rows) {
        active = rows;
    }
 
    if (active > 1 && ck_threadpool_thread_id(pool) <= 0) {
        ck_threadpool_parallel_for_n(
            pool, active, 0, rows, 4,
            attn_gate_sigmoid_mul_row_range, &args);
    } else {
        attn_gate_sigmoid_mul_row_range(0, rows, &args);
    }
}

References attn_gate_sigmoid_mul_row_range(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), ck_threadpool_thread_id(), and hybrid_ranges_overlap().

Referenced by ck_test_attn_gate_sigmoid_mul().

attn_gate_sigmoid_mul_pytorch_bf16_storage()

void attn_gate_sigmoid_mul_pytorch_bf16_storage	(	const float *	x,
		const float *	gate,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 199 of file hybrid_attention_kernels.c.

                                                               {
    if (!x || !gate || !out || rows <= 0 || num_heads <= 0 || state_dim <= 0) {
        return;
    }
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t)row * (size_t)dim;
        const float *gate_row = gate + (size_t)row * (size_t)dim;
        float *out_row = out + (size_t)row * (size_t)dim;
        for (int col = 0; col < dim; col += 16) {
            const int width = dim - col < 16 ? dim - col : 16;
            float sigmoid[16];
            recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(
                gate_row + col, sigmoid, 1, width);
            for (int lane = 0; lane < width; ++lane) {
                const float x_bf16 = bf16_to_float(float_to_bf16(x_row[col + lane]));
                const float sigmoid_bf16 = bf16_to_float(float_to_bf16(sigmoid[lane]));
                out_row[col + lane] = bf16_to_float(float_to_bf16(
                    x_bf16 * sigmoid_bf16));
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), and recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output().

attn_gate_sigmoid_mul_row_range()

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

static

Definition at line 126 of file hybrid_attention_kernels.c.

                                                          {
    const ck_attn_gate_sigmoid_mul_args_t *args =
        (const ck_attn_gate_sigmoid_mul_args_t *)opaque;
    for (int row = begin; row < end; ++row) {
        const float *x_row =
            args->x + (size_t)row * (size_t)args->dim;
        const float *gate_row =
            args->gate + (size_t)row * (size_t)args->dim;
        float *out_row =
            args->out + (size_t)row * (size_t)args->dim;
        const int dim = args->dim;
        for (int col = 0; col < dim; ++col) {
            out_row[col] = x_row[col] * hybrid_sigmoid(gate_row[col]);
        }
    }
}

References end, and hybrid_sigmoid().

Referenced by attn_gate_sigmoid_mul_forward().

attn_gate_softplus_mul_forward()

void attn_gate_softplus_mul_forward	(	const float *	x,
		const float *	gate,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 228 of file hybrid_attention_kernels.c.

                                                   {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        const float *gate_row = gate + (size_t) row * (size_t) num_heads;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int head = 0; head < num_heads; ++head) {
            const float value = gate_row[head];
            const float scale = value > 20.0f
                ? value
                : log1pf(expf(value));
            const size_t base = (size_t) head * (size_t) state_dim;
            for (int col = 0; col < state_dim; ++col) {
                out_row[base + (size_t) col] =
                    x_row[base + (size_t) col] * scale;
            }
        }
    }
}

ck_bind_hybrid_llama_libm()

static void ck_bind_hybrid_llama_libm ( void  )

static

Definition at line 19 of file hybrid_attention_kernels.c.

                                            {
    ck_hybrid_libm_handle = dlopen("libm.so.6", RTLD_NOW | RTLD_LOCAL);
    if (ck_hybrid_libm_handle) {
        ck_hybrid_llama_expf =
            (ck_hybrid_libm_f32_fn)dlsym(ck_hybrid_libm_handle, "expf");
    }
    if (!ck_hybrid_llama_expf) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: llama.cpp attention gate "
                "requires expf from libm.so.6\n");
        abort();
    }
}

References ck_hybrid_libm_handle, and ck_hybrid_llama_expf.

Referenced by hybrid_sigmoid().

hybrid_ranges_overlap()

static int hybrid_ranges_overlap ( const float *  lhs, const float *  rhs, size_t  elements  )

static

Definition at line 145 of file hybrid_attention_kernels.c.

                                                  {
    const uintptr_t lhs_begin = (uintptr_t)lhs;
    const uintptr_t rhs_begin = (uintptr_t)rhs;
    const size_t bytes = elements * sizeof(float);
    return lhs_begin < rhs_begin + bytes && rhs_begin < lhs_begin + bytes;
}

Referenced by attn_gate_sigmoid_mul_forward().

hybrid_sigmoid()

static float hybrid_sigmoid ( float  x )

static

Definition at line 36 of file hybrid_attention_kernels.c.

                                     {
    pthread_once(&ck_hybrid_libm_once, ck_bind_hybrid_llama_libm);
    return 1.0f / (1.0f + ck_hybrid_llama_expf(-x));
}

References ck_bind_hybrid_llama_libm(), ck_hybrid_libm_once, and ck_hybrid_llama_expf.

Referenced by attn_gate_sigmoid_mul_backward(), and attn_gate_sigmoid_mul_row_range().

split_q_gate_backward()

void split_q_gate_backward	(	const float *	d_q,
		const float *	d_gate,
		float *	d_packed_qg,
		int	rows,
		int	q_dim,
		int	gate_dim,
		int	group_dim
	)

Definition at line 80 of file hybrid_attention_kernels.c.

                                          {
    const int packed_dim = q_dim + gate_dim;
    if (!d_q || !d_gate || !d_packed_qg || rows <= 0 || q_dim <= 0 || gate_dim <= 0) {
        return;
    }
    if (group_dim <= 0) {
        group_dim = q_dim;
    }
    const int q_groups = q_dim / group_dim;
    const int gate_group_dim = (q_groups > 0 && gate_dim % q_groups == 0) ? (gate_dim / q_groups) : gate_dim;
    for (int row = 0; row < rows; ++row) {
        const float *dq_src = d_q + (size_t) row * (size_t) q_dim;
        const float *dg_src = d_gate + (size_t) row * (size_t) gate_dim;
        float *dst = d_packed_qg + (size_t) row * (size_t) packed_dim;
        if (q_groups > 0 && q_groups * group_dim == q_dim && q_groups * gate_group_dim == gate_dim) {
            for (int group = 0; group < q_groups; ++group) {
                const size_t dst_group_off = (size_t) group * (size_t) (group_dim + gate_group_dim);
                memcpy(
                    dst + dst_group_off,
                    dq_src + (size_t) group * (size_t) group_dim,
                    (size_t) group_dim * sizeof(float));
                memcpy(
                    dst + dst_group_off + (size_t) group_dim,
                    dg_src + (size_t) group * (size_t) gate_group_dim,
                    (size_t) gate_group_dim * sizeof(float));
            }
        } else {
            memcpy(dst, dq_src, (size_t) q_dim * sizeof(float));
            memcpy(dst + q_dim, dg_src, (size_t) gate_dim * sizeof(float));
        }
    }
}

split_q_gate_forward()

void split_q_gate_forward	(	const float *	packed_qg,
		float *	q,
		float *	gate,
		int	rows,
		int	q_dim,
		int	gate_dim,
		int	group_dim
	)

Definition at line 41 of file hybrid_attention_kernels.c.

                                         {
    const int packed_dim = q_dim + gate_dim;
    if (!packed_qg || !q || !gate || rows <= 0 || q_dim <= 0 || gate_dim <= 0) {
        return;
    }
    if (group_dim <= 0) {
        group_dim = q_dim;
    }
    const int q_groups = q_dim / group_dim;
    const int gate_group_dim = (q_groups > 0 && gate_dim % q_groups == 0) ? (gate_dim / q_groups) : gate_dim;
    for (int row = 0; row < rows; ++row) {
        const float *src = packed_qg + (size_t) row * (size_t) packed_dim;
        float *q_dst = q + (size_t) row * (size_t) q_dim;
        float *gate_dst = gate + (size_t) row * (size_t) gate_dim;
        if (q_groups > 0 && q_groups * group_dim == q_dim && q_groups * gate_group_dim == gate_dim) {
            for (int group = 0; group < q_groups; ++group) {
                const size_t src_group_off = (size_t) group * (size_t) (group_dim + gate_group_dim);
                memcpy(
                    q_dst + (size_t) group * (size_t) group_dim,
                    src + src_group_off,
                    (size_t) group_dim * sizeof(float));
                memcpy(
                    gate_dst + (size_t) group * (size_t) gate_group_dim,
                    src + src_group_off + (size_t) group_dim,
                    (size_t) gate_group_dim * sizeof(float));
            }
        } else {
            memcpy(q_dst, src, (size_t) q_dim * sizeof(float));
            memcpy(gate_dst, src + q_dim, (size_t) gate_dim * sizeof(float));
        }
    }
}

Referenced by ck_test_split_q_gate().

Variable Documentation

ck_hybrid_libm_handle

void* ck_hybrid_libm_handle = NULL

static

Definition at line 16 of file hybrid_attention_kernels.c.

Referenced by ck_bind_hybrid_llama_libm().

ck_hybrid_libm_once

pthread_once_t ck_hybrid_libm_once = PTHREAD_ONCE_INIT

static

Definition at line 17 of file hybrid_attention_kernels.c.

Referenced by hybrid_sigmoid().

ck_hybrid_llama_expf

ck_hybrid_libm_f32_fn ck_hybrid_llama_expf = NULL

static

Definition at line 15 of file hybrid_attention_kernels.c.

Referenced by ck_bind_hybrid_llama_libm(), and hybrid_sigmoid().