hybrid_attention_kernels.c File Reference

#include "ckernel_engine.h"
#include "bf16_utils.h"
#include <dlfcn.h>
#include <math.h>
#include <pthread.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)
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 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 11 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 187 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 116 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) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            out_row[col] = x_row[col] * hybrid_sigmoid(gate_row[col]);
        }
    }
}

References hybrid_sigmoid().

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 133 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_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 162 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 16 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_sigmoid()

static float hybrid_sigmoid ( float  x )

static

Definition at line 33 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_forward().

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 77 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 38 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 13 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 14 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 12 of file hybrid_attention_kernels.c.

Referenced by ck_bind_hybrid_llama_libm(), and hybrid_sigmoid().