recurrent_split_kernels.c File Reference

#include "ckernel_engine.h"
#include <string.h>

Go to the source code of this file.

Functions
void	recurrent_split_conv_qkv_backward (const float *d_q, const float *d_k, const float *d_v, float *d_packed_qkv, int rows, int q_dim, int k_dim, int v_dim)
void	recurrent_split_conv_qkv_forward (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
void	recurrent_split_qkv_backward (const float *d_q, const float *d_k, const float *d_v, float *d_packed_qkv, int rows, int q_dim, int k_dim, int v_dim)
void	recurrent_split_qkv_forward (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim)
void	split_qkv_packed_head_major_forward (const float *packed_qkv, float *q, float *k, float *v, int rows, int q_dim, int k_dim, int v_dim, int num_heads, int num_kv_heads)

Function Documentation

recurrent_split_conv_qkv_backward()

void recurrent_split_conv_qkv_backward	(	const float *	d_q,
		const float *	d_k,
		const float *	d_v,
		float *	d_packed_qkv,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 108 of file recurrent_split_kernels.c.

                                                  {
    recurrent_split_qkv_backward(d_q, d_k, d_v, d_packed_qkv, rows, q_dim, k_dim, v_dim);
}

References recurrent_split_qkv_backward().

recurrent_split_conv_qkv_forward()

void recurrent_split_conv_qkv_forward	(	const float *	packed_qkv,
		float *	q,
		float *	k,
		float *	v,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 97 of file recurrent_split_kernels.c.

                                                 {
    recurrent_split_qkv_forward(packed_qkv, q, k, v, rows, q_dim, k_dim, v_dim);
}

References recurrent_split_qkv_forward().

Referenced by ck_test_recurrent_split_conv_qkv().

recurrent_split_qkv_backward()

void recurrent_split_qkv_backward	(	const float *	d_q,
		const float *	d_k,
		const float *	d_v,
		float *	d_packed_qkv,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 77 of file recurrent_split_kernels.c.

                                             {
    const int packed_dim = q_dim + k_dim + v_dim;
    for (int row = 0; row < rows; ++row) {
        const float *dq_src = d_q + (size_t) row * (size_t) q_dim;
        const float *dk_src = d_k + (size_t) row * (size_t) k_dim;
        const float *dv_src = d_v + (size_t) row * (size_t) v_dim;
        float *dst = d_packed_qkv + (size_t) row * (size_t) packed_dim;
        memcpy(dst, dq_src, (size_t) q_dim * sizeof(float));
        memcpy(dst + q_dim, dk_src, (size_t) k_dim * sizeof(float));
        memcpy(dst + q_dim + k_dim, dv_src, (size_t) v_dim * sizeof(float));
    }
}

Referenced by recurrent_split_conv_qkv_backward().

recurrent_split_qkv_forward()

void recurrent_split_qkv_forward	(	const float *	packed_qkv,
		float *	q,
		float *	k,
		float *	v,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 5 of file recurrent_split_kernels.c.

                                            {
    const int packed_dim = q_dim + k_dim + v_dim;
    for (int row = 0; row < rows; ++row) {
        const float *src = packed_qkv + (size_t) row * (size_t) packed_dim;
        float *q_dst = q + (size_t) row * (size_t) q_dim;
        float *k_dst = k + (size_t) row * (size_t) k_dim;
        float *v_dst = v + (size_t) row * (size_t) v_dim;
        memcpy(q_dst, src, (size_t) q_dim * sizeof(float));
        memcpy(k_dst, src + q_dim, (size_t) k_dim * sizeof(float));
        memcpy(v_dst, src + q_dim + k_dim, (size_t) v_dim * sizeof(float));
    }
}

Referenced by ck_test_recurrent_split_qkv(), and recurrent_split_conv_qkv_forward().

split_qkv_packed_head_major_forward()

void split_qkv_packed_head_major_forward	(	const float *	packed_qkv,
		float *	q,
		float *	k,
		float *	v,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim,
		int	num_heads,
		int	num_kv_heads
	)

Definition at line 25 of file recurrent_split_kernels.c.

                                                           {
    if (!packed_qkv || !q || !k || !v || rows <= 0 || q_dim <= 0 || k_dim <= 0 || v_dim <= 0 ||
        num_heads <= 0 || num_kv_heads <= 0) {
        return;
    }
 
    const int q_head_dim = q_dim / num_heads;
    const int k_head_dim = k_dim / num_kv_heads;
    const int v_head_dim = v_dim / num_kv_heads;
    if (q_head_dim <= 0 || k_head_dim <= 0 || v_head_dim <= 0) {
        return;
    }
    if (q_head_dim * num_heads != q_dim || k_head_dim * num_kv_heads != k_dim || v_head_dim * num_kv_heads != v_dim) {
        return;
    }
 
    const int packed_dim = q_dim + k_dim + v_dim;
    const size_t q_head_stride = (size_t) rows * (size_t) q_head_dim;
    const size_t k_head_stride = (size_t) rows * (size_t) k_head_dim;
    const size_t v_head_stride = (size_t) rows * (size_t) v_head_dim;
 
    for (int row = 0; row < rows; ++row) {
        const float *src = packed_qkv + (size_t) row * (size_t) packed_dim;
 
        for (int head = 0; head < num_heads; ++head) {
            const float *src_q = src + (size_t) head * (size_t) q_head_dim;
            float *dst_q = q + (size_t) head * q_head_stride + (size_t) row * (size_t) q_head_dim;
            memcpy(dst_q, src_q, (size_t) q_head_dim * sizeof(float));
        }
 
        const float *src_k_base = src + (size_t) q_dim;
        const float *src_v_base = src + (size_t) q_dim + (size_t) k_dim;
        for (int head = 0; head < num_kv_heads; ++head) {
            const float *src_k = src_k_base + (size_t) head * (size_t) k_head_dim;
            const float *src_v = src_v_base + (size_t) head * (size_t) v_head_dim;
            float *dst_k = k + (size_t) head * k_head_stride + (size_t) row * (size_t) k_head_dim;
            float *dst_v = v + (size_t) head * v_head_stride + (size_t) row * (size_t) v_head_dim;
            memcpy(dst_k, src_k, (size_t) k_head_dim * sizeof(float));
            memcpy(dst_v, src_v, (size_t) v_head_dim * sizeof(float));
        }
    }
}