forked from tinygrad/tinygrad
qazalandGitHub
939f28d571
* work * fused qkv_norm * work * speed * not that yet * test cleanup * just clone * remove .realize() * cleanup tests
70 lines
2.3 KiB
C++
70 lines
2.3 KiB
C++
#include <hip/hip_runtime.h>
|
|
#include <hip/hip_bf16.h>
|
|
|
|
#ifndef ATTN_B
|
|
#define ATTN_B 2
|
|
#endif
|
|
#ifndef ATTN_N
|
|
#define ATTN_N 8192
|
|
#endif
|
|
#ifndef ATTN_H
|
|
#define ATTN_H 32
|
|
#endif
|
|
#ifndef ATTN_H_KV
|
|
#define ATTN_H_KV 8
|
|
#endif
|
|
#ifndef ATTN_D
|
|
#define ATTN_D 128
|
|
#endif
|
|
#ifndef THREADS_PER_BLOCK
|
|
#define THREADS_PER_BLOCK 256
|
|
#endif
|
|
|
|
constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
|
|
constexpr int HALF_D = ATTN_D / 2;
|
|
constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
|
|
|
|
extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
|
|
fused_qkv_rope_forward(
|
|
__hip_bfloat16* __restrict__ q,
|
|
__hip_bfloat16* __restrict__ k,
|
|
__hip_bfloat16* __restrict__ v,
|
|
const __hip_bfloat16* __restrict__ xqkv,
|
|
const __hip_bfloat16* __restrict__ freqs_cis) {
|
|
const int b = blockIdx.x;
|
|
const int n = blockIdx.y;
|
|
const int bn = b * ATTN_N + n;
|
|
const int packed_bn = bn * ATTN_H_KV * PACKED_D;
|
|
const int q_bn = bn * ATTN_H * ATTN_D;
|
|
const int kv_bn = bn * ATTN_H_KV * ATTN_D;
|
|
|
|
if (threadIdx.x < HALF_D) {
|
|
const int pair = threadIdx.x;
|
|
const int even = pair << 1;
|
|
const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
|
|
const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
|
|
|
|
for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
|
|
const int base = packed_bn + kvh * PACKED_D;
|
|
|
|
for (int rep = 0; rep < GROUP_SIZE; rep++) {
|
|
const int qbase = base + rep * ATTN_D;
|
|
const int h = kvh * GROUP_SIZE + rep;
|
|
const float a = static_cast<float>(xqkv[qbase + even]);
|
|
const float bb = static_cast<float>(xqkv[qbase + even + 1]);
|
|
const int out = q_bn + h * ATTN_D + even;
|
|
q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
|
|
q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
|
|
}
|
|
|
|
const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
|
|
const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
|
|
const int out = kv_bn + kvh * ATTN_D + even;
|
|
k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
|
|
k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
|
|
v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
|
|
v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
|
|
}
|
|
}
|
|
}
|