mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-07 22:47:54 +02:00
Compare commits
15 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| a8cfdbb9e4 | |||
| 6f8895feec | |||
| ee445f93d8 | |||
| f36e5c348b | |||
| 74976e1aef | |||
| 9abce7473a | |||
| cb295bf596 | |||
| bfdf581b8b | |||
| 20a04b2206 | |||
| 3b4fca11ac | |||
| 86961efd56 | |||
| d80e878501 | |||
| 48719618e8 | |||
| d06ddd3589 | |||
| 898b08854d |
+22
-1
@@ -55,6 +55,10 @@
|
||||
#include <pwd.h>
|
||||
#endif
|
||||
|
||||
#if defined(_AIX)
|
||||
#include <sys/systemcfg.h>
|
||||
#endif
|
||||
|
||||
#if defined(_MSC_VER)
|
||||
#pragma warning(disable: 4244 4267) // possible loss of data
|
||||
#endif
|
||||
@@ -72,7 +76,16 @@ common_time_meas::~common_time_meas() {
|
||||
//
|
||||
|
||||
int32_t common_cpu_get_num_physical_cores() {
|
||||
#ifdef __linux__
|
||||
#if defined(_AIX)
|
||||
int32_t logical_cpus = _system_configuration.ncpus;
|
||||
int32_t smt_threads = _system_configuration.smt_threads;
|
||||
if (smt_threads > 0) {
|
||||
return static_cast<int32_t>(logical_cpus / smt_threads);
|
||||
}
|
||||
if (logical_cpus > 0) {
|
||||
return static_cast<int32_t>(logical_cpus);
|
||||
}
|
||||
#elif defined(__linux__)
|
||||
// enumerate the set of thread siblings, num entries is num cores
|
||||
std::unordered_set<std::string> siblings;
|
||||
for (uint32_t cpu=0; cpu < UINT32_MAX; ++cpu) {
|
||||
@@ -202,6 +215,14 @@ int32_t common_cpu_get_num_math() {
|
||||
}
|
||||
}
|
||||
}
|
||||
#elif defined(__powerpc64__) || defined(__powerpc__)
|
||||
int32_t smt_factor = 1;
|
||||
int phy_cpus = common_cpu_get_num_physical_cores();
|
||||
int logical_cpus = sysconf(_SC_NPROCESSORS_ONLN);
|
||||
if (phy_cpus > 0 && logical_cpus > phy_cpus) {
|
||||
smt_factor = logical_cpus / phy_cpus;
|
||||
}
|
||||
return phy_cpus * std::min(smt_factor, 2);
|
||||
#endif
|
||||
return common_cpu_get_num_physical_cores();
|
||||
}
|
||||
|
||||
+3
-6
@@ -270,13 +270,10 @@ The environment variable [`CUDA_SCALE_LAUNCH_QUEUES`](https://docs.nvidia.com/cu
|
||||
|
||||
Consider setting `CUDA_SCALE_LAUNCH_QUEUES=4x`, which increases the CUDA command buffer to 4 times its default size. This optimization is particularly beneficial for **Multi-GPU setups with pipeline parallelism**, where it significantly improves prompt processing throughput by allowing more operations to be enqueued across GPUs.
|
||||
|
||||
#### GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F
|
||||
#### GGML_CUDA_CUBLAS_COMPUTE_TYPE
|
||||
|
||||
Use `GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F` environment variable to use FP32 compute type on all GPUs in FP16 cuBLAS for preventing possible numerical overflows in exchange for slower prompt processing (small impact on RTX PRO/Datacenter products and significant on GeForce products).
|
||||
|
||||
#### GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F
|
||||
|
||||
Use `GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F` environment variable to force use FP16 compute type (instead of default FP32) in FP16 cuBLAS for V100, CDNA and RDNA4.
|
||||
Override default, speed-optimized compute types for cuBLAS matrix multiplications.
|
||||
Legal values: `auto`, `f16`, `fp16`, `bf16`, `f32`, `fp32`.
|
||||
|
||||
### Unified Memory
|
||||
|
||||
|
||||
@@ -30,9 +30,6 @@ GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int de
|
||||
// conduct allreduce operation between devices
|
||||
GGML_BACKEND_API bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends);
|
||||
|
||||
// split tensor buffer that splits matrices by rows across multiple devices
|
||||
GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split);
|
||||
|
||||
// pinned host buffer for use with the CPU backend for faster copies between CPU and GPU
|
||||
GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void);
|
||||
|
||||
|
||||
@@ -812,10 +812,10 @@ void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
|
||||
const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
|
||||
const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
|
||||
const float32x4_t nvsc = {
|
||||
ggml_ue4m3_to_fp32(x[ib].d[0]),
|
||||
ggml_ue4m3_to_fp32(x[ib].d[1]),
|
||||
ggml_ue4m3_to_fp32(x[ib].d[2]),
|
||||
ggml_ue4m3_to_fp32(x[ib].d[3])
|
||||
GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]),
|
||||
GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]),
|
||||
GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]),
|
||||
GGML_CPU_UE4M3_TO_FP32(x[ib].d[3])
|
||||
};
|
||||
const float32x4_t scales = vmulq_f32(nvsc, (float32x4_t){dy0, dy0, dy1, dy1});
|
||||
|
||||
|
||||
@@ -2321,24 +2321,28 @@ class tinyBLAS_Q0_PPC {
|
||||
}
|
||||
|
||||
void matmul(int64_t m, int64_t n) {
|
||||
#if defined(_AIX) || defined(__BIG_ENDIAN__)
|
||||
mnpack(0, m, 0, n);
|
||||
#else
|
||||
const int64_t mc = 64;
|
||||
const int64_t kc = 64;
|
||||
int64_t mc = 64;
|
||||
int64_t nc = 64;
|
||||
int64_t kc = 64;
|
||||
int64_t n_chunk = 64;
|
||||
#if defined(_AIX) || defined(__BIG_ENDIAN__)
|
||||
mc = 32;
|
||||
nc = 32;
|
||||
kc = 32;
|
||||
n_chunk = 32
|
||||
#endif
|
||||
int64_t n_aligned = 0;
|
||||
if (n % 64 == 0) {
|
||||
if (n % n_chunk == 0) {
|
||||
n_aligned = n;
|
||||
} else if (n == 4) {
|
||||
n_aligned = 4;
|
||||
} else if (n < 64) {
|
||||
} else if (n < n_chunk) {
|
||||
n_aligned = (n / 8) * 8;
|
||||
} else {
|
||||
n_aligned = (n / 64) * 64;
|
||||
n_aligned = (n / n_chunk) * n_chunk;
|
||||
}
|
||||
if (n_aligned > 0) {
|
||||
if (n_aligned % 64 == 0) nc = 64;
|
||||
if (n_aligned % n_chunk == 0) nc = n_chunk;
|
||||
else if (n_aligned == n) nc = n;
|
||||
else if (n_aligned % 32 == 0) nc = 32;
|
||||
else if (n_aligned % 24 == 0) nc = 24;
|
||||
@@ -2354,7 +2358,6 @@ class tinyBLAS_Q0_PPC {
|
||||
} else {
|
||||
mnpack(0, m, 0, n);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
private:
|
||||
@@ -3195,16 +3198,19 @@ class tinyBLAS_PPC {
|
||||
}
|
||||
|
||||
void matmul(int64_t m, int64_t n) {
|
||||
int64_t mc = 256;
|
||||
int64_t nc = 256;
|
||||
int64_t kc = 256;
|
||||
#if defined(_AIX) || defined(__BIG_ENDIAN__)
|
||||
mnpack(0, m, 0, n);
|
||||
#else
|
||||
int64_t mc = 256; int64_t nc = 256; int64_t kc = 256;
|
||||
mc = 128;
|
||||
nc = 128;
|
||||
kc = 128;
|
||||
#endif
|
||||
if (m % mc == 0 && n % nc == 0 && k % kc == 0) {
|
||||
matmul_tiled(m, n, mc, nc, kc);
|
||||
} else {
|
||||
mnpack(0, m, 0, n);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
@@ -131,8 +131,8 @@ extern float ggml_table_f32_ue4m3[1 << 8];
|
||||
#define GGML_CPU_E8M0_TO_FP32_HALF(x) GGML_E8M0_TO_FP32_HALF(x)
|
||||
#endif
|
||||
|
||||
// Use lookup table for UE4M3 on x86 (faster than bit manipulation)
|
||||
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
|
||||
// Use lookup table for UE4M3 on x86 and ARM (faster than bit manipulation)
|
||||
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__ARM_NEON)
|
||||
#define GGML_CPU_UE4M3_TO_FP32(x) ggml_table_f32_ue4m3[(uint8_t)(x)]
|
||||
#else
|
||||
#define GGML_CPU_UE4M3_TO_FP32(x) ggml_ue4m3_to_fp32(x)
|
||||
|
||||
@@ -104,8 +104,8 @@ static __global__ void dequantize_block_q4_0(const void * __restrict__ vx, dst_t
|
||||
const uint8_t * q = x->qs + 4*il;
|
||||
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
y[l+ 0] = d * (q[l] & 0xF) + dm;
|
||||
y[l+16] = d * (q[l] >> 4) + dm;
|
||||
y[l+ 0] = ggml_cuda_cast<dst_t>(d * (q[l] & 0xF) + dm);
|
||||
y[l+16] = ggml_cuda_cast<dst_t>(d * (q[l] >> 4) + dm);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -131,8 +131,8 @@ static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t
|
||||
const uint8_t * q = x->qs + 4*il;
|
||||
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
y[l+ 0] = d.x * (q[l] & 0xF) + d.y;
|
||||
y[l+16] = d.x * (q[l] >> 4) + d.y;
|
||||
y[l+ 0] = ggml_cuda_cast<dst_t>(d.x * (q[l] & 0xF) + d.y);
|
||||
y[l+16] = ggml_cuda_cast<dst_t>(d.x * (q[l] >> 4) + d.y);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -154,10 +154,10 @@ static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t
|
||||
|
||||
float dall = __low2half(x[i].dm);
|
||||
float dmin = __high2half(x[i].dm);
|
||||
y[l+ 0] = dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4);
|
||||
y[l+32] = dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4);
|
||||
y[l+64] = dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4);
|
||||
y[l+96] = dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4);
|
||||
y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4));
|
||||
y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4));
|
||||
y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4));
|
||||
y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4));
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
@@ -188,7 +188,9 @@ static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t
|
||||
const uint8_t * q = x[i].qs + 32*n;
|
||||
const uint8_t * hm = x[i].hmask;
|
||||
|
||||
for (int l = l0; l < l0+4; ++l) y[l] = dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4));
|
||||
for (int l = l0; l < l0+4; ++l) {
|
||||
y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
|
||||
}
|
||||
}
|
||||
|
||||
static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
|
||||
@@ -226,8 +228,8 @@ static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t
|
||||
get_scale_min_k4(is + 1, x[i].scales, sc, m);
|
||||
const float d2 = dall * sc; const float m2 = dmin * m;
|
||||
for (int l = 0; l < n; ++l) {
|
||||
y[l + 0] = d1 * (q[l] & 0xF) - m1;
|
||||
y[l +32] = d2 * (q[l] >> 4) - m2;
|
||||
y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
|
||||
y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >> 4) - m2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -258,11 +260,11 @@ static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t
|
||||
const float d2 = dall * sc; const float m2 = dmin * m;
|
||||
|
||||
uint8_t hm = 1 << (2*il);
|
||||
y[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1;
|
||||
y[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1;
|
||||
y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
|
||||
y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
|
||||
hm <<= 1;
|
||||
y[32] = d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2;
|
||||
y[33] = d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2;
|
||||
y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2);
|
||||
y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2);
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
@@ -285,10 +287,10 @@ static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t
|
||||
const uint8_t qh = x[i].qh[32*ip + il];
|
||||
const int8_t * sc = x[i].scales + is;
|
||||
|
||||
y[ 0] = d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32);
|
||||
y[32] = d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32);
|
||||
y[64] = d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32);
|
||||
y[96] = d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32);
|
||||
y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
|
||||
y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
|
||||
y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32));
|
||||
y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32));
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
@@ -307,7 +309,9 @@ static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, ds
|
||||
const uint32_t aux32 = q2[2] | (q2[3] << 16);
|
||||
const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.25f;
|
||||
const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
|
||||
for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
|
||||
}
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
@@ -324,7 +328,9 @@ static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst
|
||||
const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
|
||||
const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
|
||||
const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
|
||||
for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
|
||||
}
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
@@ -340,7 +346,9 @@ static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_
|
||||
const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
|
||||
const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
|
||||
const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
|
||||
for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
|
||||
}
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
@@ -361,8 +369,8 @@ static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, ds
|
||||
const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.5f;
|
||||
const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
|
||||
y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
|
||||
y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
|
||||
y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -382,8 +390,8 @@ static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_
|
||||
const float d = (float)x[i].d * (1 + 2*((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf));
|
||||
const uint8_t signs = x[i].signs[4*ib + il];
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
|
||||
y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
|
||||
y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
|
||||
y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -404,7 +412,7 @@ static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_
|
||||
grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
|
||||
grid32[0] &= 0x0f0f0f0f;
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
y[j] = d * (q[j] + delta);
|
||||
y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -429,7 +437,7 @@ static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_
|
||||
grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
|
||||
grid32[0] &= 0x0f0f0f0f;
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
y[j] = d * (q[j] + delta);
|
||||
y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -446,8 +454,8 @@ static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst
|
||||
const uint8_t * q4 = x[ib].qs + 4*il;
|
||||
const float d = (float)x[ib].d;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
|
||||
y[j+16] = d * kvalues_iq4nl[q4[j] >> 4];
|
||||
y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
|
||||
y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >> 4]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -463,8 +471,8 @@ static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst
|
||||
const uint8_t * q4 = x[i].qs + 16*ib + 4*il;
|
||||
const float d = (float)x[i].d * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
|
||||
y[j+16] = d * kvalues_iq4nl[q4[j] >> 4];
|
||||
y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
|
||||
y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >> 4]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -481,8 +489,8 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_
|
||||
const uint8_t * q4 = x[ib].qs + 4*il;
|
||||
const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
y[j+ 0] = d * kvalues_mxfp4[q4[j] & 0xf]*0.5f;
|
||||
y[j+16] = d * kvalues_mxfp4[q4[j] >> 4]*0.5f;
|
||||
y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
|
||||
y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >> 4]*0.5f);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -700,6 +708,50 @@ static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k,
|
||||
|
||||
to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_Q1_0:
|
||||
return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
|
||||
case GGML_TYPE_Q4_0:
|
||||
return dequantize_row_q4_0_cuda;
|
||||
case GGML_TYPE_Q4_1:
|
||||
return dequantize_row_q4_1_cuda;
|
||||
case GGML_TYPE_Q5_0:
|
||||
return dequantize_block_cont_cuda<QK5_0, QR5_0, dequantize_q5_0>;
|
||||
case GGML_TYPE_Q5_1:
|
||||
return dequantize_block_cont_cuda<QK5_1, QR5_1, dequantize_q5_1>;
|
||||
case GGML_TYPE_Q8_0:
|
||||
return dequantize_block_cont_cuda<QK8_0, QR8_0, dequantize_q8_0>;
|
||||
case GGML_TYPE_Q2_K:
|
||||
return dequantize_row_q2_K_cuda;
|
||||
case GGML_TYPE_Q3_K:
|
||||
return dequantize_row_q3_K_cuda;
|
||||
case GGML_TYPE_Q4_K:
|
||||
return dequantize_row_q4_K_cuda;
|
||||
case GGML_TYPE_Q5_K:
|
||||
return dequantize_row_q5_K_cuda;
|
||||
case GGML_TYPE_Q6_K:
|
||||
return dequantize_row_q6_K_cuda;
|
||||
case GGML_TYPE_IQ2_XXS:
|
||||
return dequantize_row_iq2_xxs_cuda;
|
||||
case GGML_TYPE_IQ2_XS:
|
||||
return dequantize_row_iq2_xs_cuda;
|
||||
case GGML_TYPE_IQ2_S:
|
||||
return dequantize_row_iq2_s_cuda;
|
||||
case GGML_TYPE_IQ3_XXS:
|
||||
return dequantize_row_iq3_xxs_cuda;
|
||||
case GGML_TYPE_IQ1_S:
|
||||
return dequantize_row_iq1_s_cuda;
|
||||
case GGML_TYPE_IQ1_M:
|
||||
return dequantize_row_iq1_m_cuda;
|
||||
case GGML_TYPE_IQ4_NL:
|
||||
return dequantize_row_iq4_nl_cuda;
|
||||
case GGML_TYPE_IQ4_XS:
|
||||
return dequantize_row_iq4_xs_cuda;
|
||||
case GGML_TYPE_IQ3_S:
|
||||
return dequantize_row_iq3_s_cuda;
|
||||
case GGML_TYPE_MXFP4:
|
||||
return dequantize_row_mxfp4_cuda;
|
||||
case GGML_TYPE_NVFP4:
|
||||
return dequantize_row_nvfp4_cuda;
|
||||
case GGML_TYPE_F32:
|
||||
return convert_unary_cont_cuda<float>;
|
||||
case GGML_TYPE_F16:
|
||||
|
||||
+22
-16
@@ -337,6 +337,26 @@ enum best_fattn_kernel {
|
||||
BEST_FATTN_KERNEL_MMA_F16 = 400,
|
||||
};
|
||||
|
||||
static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_F16:
|
||||
return true;
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
#ifndef GGML_CUDA_FA_ALL_QUANTS
|
||||
return false;
|
||||
#endif // GGML_CUDA_FA_ALL_QUANTS
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_BF16:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
|
||||
#ifndef FLASH_ATTN_AVAILABLE
|
||||
GGML_UNUSED(device); GGML_UNUSED(dst);
|
||||
@@ -427,22 +447,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
|
||||
}
|
||||
#endif // GGML_CUDA_FA_ALL_QUANTS
|
||||
|
||||
switch (K->type) {
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_F16:
|
||||
break;
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
#ifndef GGML_CUDA_FA_ALL_QUANTS
|
||||
return BEST_FATTN_KERNEL_NONE;
|
||||
#endif // GGML_CUDA_FA_ALL_QUANTS
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_BF16:
|
||||
break;
|
||||
default:
|
||||
return BEST_FATTN_KERNEL_NONE;
|
||||
if (!ggml_cuda_fattn_kv_type_supported(K->type) || !ggml_cuda_fattn_kv_type_supported(V->type)) {
|
||||
return BEST_FATTN_KERNEL_NONE;
|
||||
}
|
||||
|
||||
if (mask && mask->ne[2] != 1) {
|
||||
|
||||
+182
-1124
File diff suppressed because it is too large
Load Diff
@@ -278,6 +278,9 @@ int get_mmvq_mmid_max_batch(ggml_type type, int cc) {
|
||||
}
|
||||
|
||||
bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
|
||||
if (!ggml_is_quantized(type)) {
|
||||
return false;
|
||||
}
|
||||
if (GGML_CUDA_CC_IS_CDNA(cc)) {
|
||||
if (GGML_CUDA_CC_IS_CDNA1(cc)) {
|
||||
switch (type) {
|
||||
|
||||
@@ -155,3 +155,5 @@ if (GGML_HIP_RCCL)
|
||||
endif()
|
||||
|
||||
target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas)
|
||||
|
||||
target_compile_options(ggml-hip PRIVATE "$<$<COMPILE_LANGUAGE:HIP>:-ffast-math>")
|
||||
|
||||
@@ -1800,6 +1800,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_COL2IM_1D);
|
||||
|
||||
GGML_ASSERT(ggml_is_contiguous(op->src[0]));
|
||||
GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16);
|
||||
|
||||
char base[256];
|
||||
char name[256];
|
||||
|
||||
snprintf(base, 256, "kernel_col2im_1d_%s", ggml_type_name(op->src[0]->type));
|
||||
snprintf(name, 256, "%s", base);
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_CONV_TRANSPOSE_2D);
|
||||
|
||||
|
||||
@@ -150,6 +150,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rope
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_im2col (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_upscale (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
|
||||
@@ -1157,6 +1157,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
||||
(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32) &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32;
|
||||
case GGML_OP_COL2IM_1D:
|
||||
return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16) &&
|
||||
op->type == op->src[0]->type &&
|
||||
ggml_is_contiguous(op->src[0]) &&
|
||||
ggml_is_contiguous(op);
|
||||
case GGML_OP_CONV_3D:
|
||||
return ggml_is_contiguous(op->src[0]) &&
|
||||
ggml_is_contiguous(op->src[1]) &&
|
||||
|
||||
@@ -603,6 +603,16 @@ typedef struct {
|
||||
uint64_t nb1;
|
||||
} ggml_metal_kargs_conv_transpose_1d;
|
||||
|
||||
typedef struct {
|
||||
int32_t T_in;
|
||||
int32_t T_out;
|
||||
int32_t OC;
|
||||
int32_t K;
|
||||
int32_t K_OC;
|
||||
int32_t s0;
|
||||
int32_t p0;
|
||||
} ggml_metal_kargs_col2im_1d;
|
||||
|
||||
typedef struct {
|
||||
int32_t IC;
|
||||
int32_t IH;
|
||||
|
||||
@@ -395,6 +395,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
|
||||
{
|
||||
n_fuse = ggml_metal_op_conv_transpose_2d(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_COL2IM_1D:
|
||||
{
|
||||
n_fuse = ggml_metal_op_col2im_1d(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_CONV_3D:
|
||||
{
|
||||
n_fuse = ggml_metal_op_conv_3d(ctx, idx);
|
||||
@@ -3854,6 +3858,47 @@ int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_col2im_1d(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_library_t lib = ctx->lib;
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
|
||||
const int32_t s0 = ((const int32_t *)(op->op_params))[0];
|
||||
const int32_t OC = ((const int32_t *)(op->op_params))[1];
|
||||
const int32_t p0 = ((const int32_t *)(op->op_params))[2];
|
||||
|
||||
const int32_t K_OC = (int32_t) op->src[0]->ne[0];
|
||||
const int32_t T_in = (int32_t) op->src[0]->ne[1];
|
||||
const int32_t K = K_OC / OC;
|
||||
const int32_t T_out = (int32_t) op->ne[0];
|
||||
|
||||
ggml_metal_kargs_col2im_1d args = {
|
||||
/*.T_in =*/ T_in,
|
||||
/*.T_out =*/ T_out,
|
||||
/*.OC =*/ OC,
|
||||
/*.K =*/ K,
|
||||
/*.K_OC =*/ K_OC,
|
||||
/*.s0 =*/ s0,
|
||||
/*.p0 =*/ p0,
|
||||
};
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_col2im_1d(lib, op);
|
||||
|
||||
const int total = T_out * OC;
|
||||
const int nth = 256;
|
||||
const int ntg = (total + nth - 1) / nth;
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
|
||||
@@ -78,6 +78,7 @@ int ggml_metal_op_conv_2d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_conv_3d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_col2im_1d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_upscale (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_pad (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_pad_reflect_1d (ggml_metal_op_t ctx, int idx);
|
||||
|
||||
@@ -4977,6 +4977,49 @@ kernel void kernel_conv_transpose_1d<half>(
|
||||
uint3 tgpg[[threadgroups_per_grid]]);
|
||||
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_col2im_1d(
|
||||
constant ggml_metal_kargs_col2im_1d & args,
|
||||
device const T * col,
|
||||
device T * dst,
|
||||
uint tgpig [[threadgroup_position_in_grid]],
|
||||
uint tpitg [[thread_position_in_threadgroup]],
|
||||
uint ntg [[threads_per_threadgroup]]) {
|
||||
|
||||
const int idx = tgpig * ntg + tpitg;
|
||||
if (idx >= args.T_out * args.OC) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int t_out = idx % args.T_out;
|
||||
const int oc = idx / args.T_out;
|
||||
const int t_abs = t_out + args.p0; // absolute position in uncropped signal
|
||||
|
||||
int t_in_min = (t_abs - args.K + args.s0) / args.s0; // ceil((t_abs - K + 1) / s0)
|
||||
if (t_in_min < 0) {
|
||||
t_in_min = 0;
|
||||
}
|
||||
int t_in_max = t_abs / args.s0;
|
||||
if (t_in_max >= args.T_in) {
|
||||
t_in_max = args.T_in - 1;
|
||||
}
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int t_in = t_in_min; t_in <= t_in_max; t_in++) {
|
||||
const int k = t_abs - t_in * args.s0;
|
||||
sum += float(col[(oc * args.K + k) + t_in * args.K_OC]);
|
||||
}
|
||||
|
||||
dst[t_out + oc * args.T_out] = T(sum);
|
||||
}
|
||||
|
||||
template [[host_name("kernel_col2im_1d_f32")]] kernel void kernel_col2im_1d<float>(constant ggml_metal_kargs_col2im_1d &, device const float *, device float *, uint, uint, uint);
|
||||
template [[host_name("kernel_col2im_1d_f16")]] kernel void kernel_col2im_1d<half>(constant ggml_metal_kargs_col2im_1d &, device const half *, device half *, uint, uint, uint);
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfloat>(constant ggml_metal_kargs_col2im_1d &, device const bfloat *, device bfloat *, uint, uint, uint);
|
||||
#endif
|
||||
|
||||
|
||||
typedef void (conv_transpose_2d_t)(
|
||||
constant ggml_metal_kargs_conv_transpose_2d & args,
|
||||
device const float * src0,
|
||||
|
||||
@@ -20,6 +20,7 @@ static const ggml_opencl_fa_dim g_fa_dims_adreno_default[] = {
|
||||
{192, 128, 16, 16, 1, 0},
|
||||
{192, 192, 16, 16, 1, 0},
|
||||
{256, 256, 16, 16, 16, 0},
|
||||
{512, 512, 8, 16, 64, 0},
|
||||
};
|
||||
|
||||
struct ggml_opencl_fa_dim_table {
|
||||
|
||||
+1250
-115
File diff suppressed because it is too large
Load Diff
@@ -10,7 +10,12 @@
|
||||
#define DK_VEC (DK/4)
|
||||
#define DV_VEC (DV/4)
|
||||
#define WG_SIZE (BLOCK_M)
|
||||
#define Q1_WG_SIZE 64
|
||||
// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
|
||||
// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
|
||||
#ifndef FA_SG
|
||||
#define FA_SG 64
|
||||
#endif
|
||||
#define Q1_WG_SIZE FA_SG
|
||||
|
||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||
|
||||
@@ -11,7 +11,12 @@
|
||||
#define DK_VEC (DK/4)
|
||||
#define DV_VEC (DV/4)
|
||||
#define WG_SIZE (BLOCK_M)
|
||||
#define Q1_WG_SIZE 64
|
||||
// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
|
||||
// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
|
||||
#ifndef FA_SG
|
||||
#define FA_SG 64
|
||||
#endif
|
||||
#define Q1_WG_SIZE FA_SG
|
||||
|
||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||
@@ -114,6 +119,15 @@ __kernel void flash_attn_f32(
|
||||
__local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
|
||||
|
||||
for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
|
||||
#if FA_SG < 64
|
||||
// WAR on l_k/l_v: threads with my_query_row >= n_q skip the compute below
|
||||
// (continue) and would race ahead to reload the tiles while active threads
|
||||
// still read them. A single 64-wide Adreno subgroup (WG == sg) runs lockstep
|
||||
// and hides this; a WG that spans multiple narrower subgroups (Intel sg=32)
|
||||
// corrupts the result. All threads reach this each iteration (no-op on the
|
||||
// first), so it does not diverge with the continue. Compiled out at sg=64.
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
#endif
|
||||
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
|
||||
const int row = i / DK_VEC;
|
||||
const int col = i % DK_VEC;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -27,7 +27,11 @@
|
||||
|
||||
#define DK_VEC (DK/4)
|
||||
#define DV_VEC (DV/4)
|
||||
#define Q1_WG_SIZE 64
|
||||
|
||||
#ifndef FA_SG
|
||||
#define FA_SG 64
|
||||
#endif
|
||||
#define Q1_WG_SIZE FA_SG
|
||||
|
||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||
@@ -365,6 +369,263 @@ __kernel void flash_attn_f32_q4_0_q1(
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef cl_intel_subgroups
|
||||
#pragma OPENCL EXTENSION cl_intel_subgroups : enable
|
||||
#else
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#endif
|
||||
|
||||
#ifdef cl_qcom_reqd_sub_group_size
|
||||
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
|
||||
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
|
||||
#else
|
||||
#define REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
|
||||
#define VEC_NSG 4
|
||||
#define VEC_WG_SIZE (Q1_WG_SIZE * VEC_NSG)
|
||||
#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
|
||||
|
||||
// Dequant one float4 lane (0..7) from a q4_0 block.
|
||||
// Lanes 0..3 → low nibbles of qs[0..15], lanes 4..7 → high nibbles.
|
||||
inline float4 dequant_q4_0_lane(const global char * block_ptr, int lane) {
|
||||
const float d = vload_half(0, (const global half *)block_ptr);
|
||||
const global uchar * qs = (const global uchar *)(block_ptr + 2);
|
||||
const int g = lane & 3;
|
||||
const int shift = (lane < 4) ? 0 : 4;
|
||||
return d * (float4)((float)((qs[g*4+0] >> shift) & 0x0F) - 8.0f,
|
||||
(float)((qs[g*4+1] >> shift) & 0x0F) - 8.0f,
|
||||
(float)((qs[g*4+2] >> shift) & 0x0F) - 8.0f,
|
||||
(float)((qs[g*4+3] >> shift) & 0x0F) - 8.0f);
|
||||
}
|
||||
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
__kernel void flash_attn_f32_q4_0_q1_vec(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
const global void * v_void, ulong v_offset,
|
||||
global void * o_void, ulong o_offset,
|
||||
const float scale,
|
||||
const int n_q,
|
||||
const int n_kv,
|
||||
const int is_causal,
|
||||
const int n_head,
|
||||
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||
const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
|
||||
const float max_bias,
|
||||
const float m0,
|
||||
const float m1,
|
||||
const int n_head_log2,
|
||||
const float logit_softcap,
|
||||
const int n_head_kv,
|
||||
const global void* mask_void,
|
||||
const ulong mask_offset,
|
||||
const ulong mask_nb1,
|
||||
const ulong mask_nb2,
|
||||
const ulong mask_nb3,
|
||||
const int mask_ne2,
|
||||
const int mask_ne3,
|
||||
const global void* sinks_void,
|
||||
const ulong sinks_offset
|
||||
) {
|
||||
const int tid = get_local_id(0);
|
||||
const int sgid = tid / Q1_WG_SIZE;
|
||||
const int tid_sg = tid % Q1_WG_SIZE;
|
||||
const int head_batch_idx = get_global_id(1);
|
||||
|
||||
const int batch_idx = head_batch_idx / n_head;
|
||||
const int head_idx = head_batch_idx % n_head;
|
||||
|
||||
const int gqa_ratio = n_head / n_head_kv;
|
||||
const int head_kv_idx = head_idx / gqa_ratio;
|
||||
|
||||
const global char * q_base = (const global char *) q_void + q_offset;
|
||||
const global char * k_base = (const global char *) k_void + k_offset;
|
||||
const global char * v_base = (const global char *) v_void + v_offset;
|
||||
global char * o_base = (global char *) o_void + o_offset;
|
||||
|
||||
const global char * mask_base = NULL;
|
||||
if (mask_void != NULL) {
|
||||
const int mask_head_idx = head_idx % mask_ne2;
|
||||
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||
mask_base = (const global char *) mask_void + mask_offset +
|
||||
mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
|
||||
}
|
||||
|
||||
__local ACC_TYPE4 q_shared[DK_VEC];
|
||||
{
|
||||
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
|
||||
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||
for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
|
||||
q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#ifdef FA_HAVE_INT_DOT
|
||||
// quantize Q to int8-packed uints + per-block (qd, q_sum) once per WG for dp4a
|
||||
// one thread per Q block, remaining threads idle this step
|
||||
__local uint q_packed_shared[DK_Q4_BLOCKS * 8];
|
||||
__local float q_d_shared[DK_Q4_BLOCKS];
|
||||
__local int q_sum_shared[DK_Q4_BLOCKS];
|
||||
if (tid < DK_Q4_BLOCKS) {
|
||||
ACC_TYPE4 q_block[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) q_block[i] = q_shared[tid * 8 + i];
|
||||
uint packed[8];
|
||||
q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) q_packed_shared[tid * 8 + i] = packed[i];
|
||||
q_d_shared[tid] = info.qd;
|
||||
q_sum_shared[tid] = info.q_sum;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
#endif
|
||||
|
||||
const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
|
||||
|
||||
const global ACC_TYPE * sinks_ptr = NULL;
|
||||
if (sinks_void != NULL) {
|
||||
sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
|
||||
}
|
||||
|
||||
ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
|
||||
|
||||
ACC_TYPE m_i = FA_M_INIT;
|
||||
ACC_TYPE l_i = 0.0f;
|
||||
|
||||
const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
|
||||
const int kv_start = sgid * kv_per_sg;
|
||||
const int kv_end = min(n_kv, kv_start + kv_per_sg);
|
||||
|
||||
for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
|
||||
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||
|
||||
#ifdef FA_HAVE_INT_DOT
|
||||
// per-lane dp4a: each lane packs 4 raw q4_0 nibbles into a uint,
|
||||
// then dot_acc_sat_4x8packed_ss_int against the matching uint.
|
||||
ACC_TYPE lane_contrib = 0.0f;
|
||||
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||
const int block_idx = qk / 8;
|
||||
const int lane_in_block = qk % 8;
|
||||
const int g = lane_in_block & 3;
|
||||
const int shift = (lane_in_block < 4) ? 0 : 4;
|
||||
const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
|
||||
const float kd = vload_half(0, (const global half *)k_block);
|
||||
const global uchar * k_qs = (const global uchar *)(k_block + 2);
|
||||
const uchar b0 = k_qs[g*4 + 0];
|
||||
const uchar b1 = k_qs[g*4 + 1];
|
||||
const uchar b2 = k_qs[g*4 + 2];
|
||||
const uchar b3 = k_qs[g*4 + 3];
|
||||
const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
|
||||
((uint)((b1 >> shift) & 0x0F)) << 8 |
|
||||
((uint)((b2 >> shift) & 0x0F)) << 16 |
|
||||
((uint)((b3 >> shift) & 0x0F)) << 24;
|
||||
const uint q_packed_lane = q_packed_shared[block_idx * 8 + lane_in_block];
|
||||
const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
|
||||
const float qd = q_d_shared[block_idx];
|
||||
const float block_scale = qd * kd;
|
||||
float contrib = (float)raw_dot * block_scale;
|
||||
if (lane_in_block == 0) {
|
||||
// block bias correction is per-block
|
||||
const int q_sum_b = q_sum_shared[block_idx];
|
||||
contrib -= 8.0f * block_scale * (float)q_sum_b;
|
||||
}
|
||||
lane_contrib += contrib;
|
||||
}
|
||||
ACC_TYPE score = sub_group_reduce_add(lane_contrib) * scale;
|
||||
#else
|
||||
ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
|
||||
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||
const int block_idx = qk / 8;
|
||||
const int lane = qk % 8;
|
||||
const float4 k_v = dequant_q4_0_lane(k_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||
dot4 = mad(q_shared[qk], k_v, dot4);
|
||||
}
|
||||
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
|
||||
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
|
||||
#endif
|
||||
|
||||
if (mask_base != NULL) {
|
||||
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
|
||||
score += slope * (ACC_TYPE) mask_ptr[k_idx];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
score = logit_softcap * tanh(score / logit_softcap);
|
||||
}
|
||||
|
||||
const ACC_TYPE m_new = max(m_i, score);
|
||||
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
|
||||
const ACC_TYPE p = native_exp(score - m_new);
|
||||
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
const int block_idx = dv / 8;
|
||||
const int lane = dv % 8;
|
||||
const float4 v_v = dequant_q4_0_lane(v_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||
o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
|
||||
}
|
||||
l_i = l_i * scale_prev + p;
|
||||
m_i = m_new;
|
||||
}
|
||||
|
||||
__local ACC_TYPE sg_m[VEC_NSG];
|
||||
__local ACC_TYPE sg_l[VEC_NSG];
|
||||
__local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
|
||||
|
||||
if (tid_sg == 0) {
|
||||
sg_m[sgid] = m_i;
|
||||
sg_l[sgid] = l_i;
|
||||
}
|
||||
{
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
sg_o[sgid][dv] = o_acc[idx];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
if (sgid == 0) {
|
||||
ACC_TYPE m_final = sg_m[0];
|
||||
#pragma unroll
|
||||
for (int s = 1; s < VEC_NSG; ++s) {
|
||||
m_final = max(m_final, sg_m[s]);
|
||||
}
|
||||
if (sinks_ptr != NULL) {
|
||||
m_final = max(m_final, sinks_ptr[head_idx]);
|
||||
}
|
||||
|
||||
ACC_TYPE l_final = 0.0f;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < VEC_NSG; ++s) {
|
||||
l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
|
||||
}
|
||||
if (sinks_ptr != NULL) {
|
||||
l_final += native_exp(sinks_ptr[head_idx] - m_final);
|
||||
}
|
||||
const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
|
||||
|
||||
const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
|
||||
global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
|
||||
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < VEC_NSG; ++s) {
|
||||
const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
|
||||
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
|
||||
}
|
||||
o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flash-decoding split pass for q4_0 KV. Merge kernel is type-agnostic and
|
||||
// shared with the f16/q8_0 FA kernels.
|
||||
#define FA_PARTIAL_FLOATS (2 + DV)
|
||||
@@ -583,6 +844,319 @@ __kernel void flash_attn_f32_q4_0_q1_split(
|
||||
#define WG_SIZE BLOCK_M
|
||||
#endif
|
||||
|
||||
#ifndef MQ_GQA
|
||||
#define MQ_GQA 4
|
||||
#endif
|
||||
#ifndef MQ_NSG_SPLIT
|
||||
#define MQ_NSG_SPLIT 4
|
||||
#endif
|
||||
#define MQ_SPLIT_WG_SIZE_Q4 (Q1_WG_SIZE * MQ_NSG_SPLIT)
|
||||
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
__kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
const global void * v_void, ulong v_offset,
|
||||
const float scale,
|
||||
const int n_q,
|
||||
const int n_kv,
|
||||
const int n_head,
|
||||
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||
const float max_bias,
|
||||
const float m0,
|
||||
const float m1,
|
||||
const int n_head_log2,
|
||||
const float logit_softcap,
|
||||
const int n_head_kv,
|
||||
const global void * mask_void,
|
||||
const ulong mask_offset,
|
||||
const ulong mask_nb1,
|
||||
const ulong mask_nb2,
|
||||
const ulong mask_nb3,
|
||||
const int mask_ne2,
|
||||
const int mask_ne3,
|
||||
global float * partial_void,
|
||||
const int n_splits,
|
||||
const int kv_per_split
|
||||
) {
|
||||
const int tid = get_local_id(0);
|
||||
const int sgid = tid / Q1_WG_SIZE;
|
||||
const int tid_sg = tid % Q1_WG_SIZE;
|
||||
const int kvhead_batch_idx = get_global_id(1);
|
||||
const int split_q_idx = get_global_id(2);
|
||||
const int split_idx = split_q_idx % n_splits;
|
||||
const int q_idx = split_q_idx / n_splits;
|
||||
|
||||
const int batch_idx = kvhead_batch_idx / n_head_kv;
|
||||
const int head_kv_idx = kvhead_batch_idx % n_head_kv;
|
||||
|
||||
const int kv_start = split_idx * kv_per_split;
|
||||
const int kv_end = min(kv_start + kv_per_split, n_kv);
|
||||
|
||||
const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
|
||||
|
||||
if (kv_start >= kv_end) {
|
||||
if (tid == 0) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||
* n_splits + split_idx);
|
||||
global float * rec = partial_void + rec_idx * record_stride;
|
||||
rec[0] = FA_M_INIT;
|
||||
rec[1] = 0.0f;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const global char * q_base = (const global char *) q_void + q_offset;
|
||||
const global char * k_base = (const global char *) k_void + k_offset;
|
||||
const global char * v_base = (const global char *) v_void + v_offset;
|
||||
|
||||
__local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
|
||||
for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q4) {
|
||||
const int h = i / DK_VEC;
|
||||
const int k = i % DK_VEC;
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
|
||||
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||
q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#ifdef FA_HAVE_INT_DOT
|
||||
__local uint q_packed_shared[MQ_GQA * DK_Q4_BLOCKS * 8];
|
||||
__local float q_d_shared[MQ_GQA * DK_Q4_BLOCKS];
|
||||
__local int q_sum_shared[MQ_GQA * DK_Q4_BLOCKS];
|
||||
{
|
||||
const int active = MQ_GQA * DK_Q4_BLOCKS;
|
||||
if (tid < active) {
|
||||
const int h = tid / DK_Q4_BLOCKS;
|
||||
const int block_id = tid % DK_Q4_BLOCKS;
|
||||
ACC_TYPE4 q_block[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) q_block[i] = q_shared[h * DK_VEC + block_id * 8 + i];
|
||||
uint packed[8];
|
||||
q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) q_packed_shared[(h * DK_Q4_BLOCKS + block_id) * 8 + i] = packed[i];
|
||||
q_d_shared[h * DK_Q4_BLOCKS + block_id] = info.qd;
|
||||
q_sum_shared[h * DK_Q4_BLOCKS + block_id] = info.q_sum;
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
#endif
|
||||
|
||||
float slope[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
|
||||
}
|
||||
|
||||
const global char * mask_base[MQ_GQA];
|
||||
if (mask_void != NULL) {
|
||||
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||
const global char * mask_base_b = (const global char *) mask_void + mask_offset +
|
||||
mask_batch_idx * mask_nb3 +
|
||||
(ulong) q_idx * mask_nb1;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
const int mask_head_idx = head_idx % mask_ne2;
|
||||
mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
|
||||
}
|
||||
|
||||
ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
|
||||
ACC_TYPE m_i[MQ_GQA];
|
||||
ACC_TYPE l_i[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
m_i[h] = FA_M_INIT;
|
||||
l_i[h] = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
|
||||
}
|
||||
|
||||
const int kv_len = kv_end - kv_start;
|
||||
const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
|
||||
const int kv_lo = kv_start + sgid * kv_per_sg;
|
||||
const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
|
||||
|
||||
for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
|
||||
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||
|
||||
#ifdef FA_HAVE_INT_DOT
|
||||
ACC_TYPE lane_contrib[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
|
||||
|
||||
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||
const int block_idx = qk / 8;
|
||||
const int lane_in_block = qk % 8;
|
||||
const int g = lane_in_block & 3;
|
||||
const int shift = (lane_in_block < 4) ? 0 : 4;
|
||||
const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
|
||||
const float kd = vload_half(0, (const global half *)k_block);
|
||||
const global uchar * k_qs = (const global uchar *)(k_block + 2);
|
||||
const uchar b0 = k_qs[g*4 + 0];
|
||||
const uchar b1 = k_qs[g*4 + 1];
|
||||
const uchar b2 = k_qs[g*4 + 2];
|
||||
const uchar b3 = k_qs[g*4 + 3];
|
||||
const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
|
||||
((uint)((b1 >> shift) & 0x0F)) << 8 |
|
||||
((uint)((b2 >> shift) & 0x0F)) << 16 |
|
||||
((uint)((b3 >> shift) & 0x0F)) << 24;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const uint q_packed_lane = q_packed_shared[(h * DK_Q4_BLOCKS + block_idx) * 8 + lane_in_block];
|
||||
const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
|
||||
const float qd = q_d_shared[h * DK_Q4_BLOCKS + block_idx];
|
||||
const float block_scale = qd * kd;
|
||||
float contrib = (float) raw_dot * block_scale;
|
||||
if (lane_in_block == 0) {
|
||||
const int q_sum_b = q_sum_shared[h * DK_Q4_BLOCKS + block_idx];
|
||||
contrib -= 8.0f * block_scale * (float) q_sum_b;
|
||||
}
|
||||
lane_contrib[h] += contrib;
|
||||
}
|
||||
}
|
||||
|
||||
ACC_TYPE score[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * scale;
|
||||
if (mask_base[h] != NULL) {
|
||||
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
|
||||
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
s = logit_softcap * tanh(s / logit_softcap);
|
||||
}
|
||||
score[h] = s;
|
||||
}
|
||||
#else
|
||||
// fallback float-dequant K dot
|
||||
ACC_TYPE4 dot4[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
|
||||
|
||||
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||
const int block_idx = qk / 8;
|
||||
const int lane = qk % 8;
|
||||
const float4 k_v = dequant_q4_0_lane(k_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
|
||||
}
|
||||
}
|
||||
|
||||
ACC_TYPE score[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
|
||||
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
|
||||
if (mask_base[h] != NULL) {
|
||||
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
|
||||
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
s = logit_softcap * tanh(s / logit_softcap);
|
||||
}
|
||||
score[h] = s;
|
||||
}
|
||||
#endif
|
||||
|
||||
ACC_TYPE p_h[MQ_GQA];
|
||||
ACC_TYPE sp_h[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const ACC_TYPE m_new = max(m_i[h], score[h]);
|
||||
sp_h[h] = native_exp(m_i[h] - m_new);
|
||||
p_h[h] = native_exp(score[h] - m_new);
|
||||
l_i[h] = l_i[h] * sp_h[h] + p_h[h];
|
||||
m_i[h] = m_new;
|
||||
}
|
||||
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
const int block_idx = dv / 8;
|
||||
const int lane = dv % 8;
|
||||
const float4 v_v = dequant_q4_0_lane(v_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// per-h cross-subgroup merge
|
||||
__local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
|
||||
__local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
|
||||
__local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
|
||||
|
||||
if (tid_sg == 0) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
sg_m[h][sgid] = m_i[h];
|
||||
sg_l[h][sgid] = l_i[h];
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
{
|
||||
int idx = 0;
|
||||
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
|
||||
sg_o[sgid][dv_idx] = o_acc[h][idx];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
if (sgid == 0) {
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
|
||||
ACC_TYPE m_c = sg_m[h][0];
|
||||
#pragma unroll
|
||||
for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
|
||||
m_c = max(m_c, sg_m[h][s]);
|
||||
}
|
||||
ACC_TYPE l_c = 0.0f;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||
l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
|
||||
}
|
||||
|
||||
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||
* n_splits + split_idx);
|
||||
global float * rec = partial_void + rec_idx * record_stride;
|
||||
global float4 * rec_o = (global float4 *) (rec + 2);
|
||||
|
||||
if (tid_sg == 0) {
|
||||
rec[0] = (float) m_c;
|
||||
rec[1] = (float) l_c;
|
||||
}
|
||||
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
|
||||
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||
const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
|
||||
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
|
||||
}
|
||||
rec_o[dv_idx] = o_merged;
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
}
|
||||
|
||||
__kernel void flash_attn_f32_q4_0(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
|
||||
@@ -24,7 +24,11 @@
|
||||
|
||||
#define DK_VEC (DK/4)
|
||||
#define DV_VEC (DV/4)
|
||||
#define Q1_WG_SIZE 64
|
||||
|
||||
#ifndef FA_SG
|
||||
#define FA_SG 64
|
||||
#endif
|
||||
#define Q1_WG_SIZE FA_SG
|
||||
|
||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||
@@ -310,6 +314,201 @@ __kernel void flash_attn_f32_q8_0_q1(
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef cl_intel_subgroups
|
||||
#pragma OPENCL EXTENSION cl_intel_subgroups : enable
|
||||
#else
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#endif
|
||||
|
||||
#ifdef cl_qcom_reqd_sub_group_size
|
||||
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
|
||||
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
|
||||
#else
|
||||
#define REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
|
||||
#define VEC_NSG 4
|
||||
#define VEC_WG_SIZE (Q1_WG_SIZE * VEC_NSG)
|
||||
#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
|
||||
|
||||
inline float4 dequant_q8_0_lane(const global char * block_ptr, int lane) {
|
||||
const float d = vload_half(0, (const global half *)block_ptr);
|
||||
const global char * qs = block_ptr + 2 + lane * 4;
|
||||
return d * (float4)((float)qs[0], (float)qs[1], (float)qs[2], (float)qs[3]);
|
||||
}
|
||||
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
__kernel void flash_attn_f32_q8_0_q1_vec(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
const global void * v_void, ulong v_offset,
|
||||
global void * o_void, ulong o_offset,
|
||||
const float scale,
|
||||
const int n_q,
|
||||
const int n_kv,
|
||||
const int is_causal,
|
||||
const int n_head,
|
||||
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||
const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
|
||||
const float max_bias,
|
||||
const float m0,
|
||||
const float m1,
|
||||
const int n_head_log2,
|
||||
const float logit_softcap,
|
||||
const int n_head_kv,
|
||||
const global void* mask_void,
|
||||
const ulong mask_offset,
|
||||
const ulong mask_nb1,
|
||||
const ulong mask_nb2,
|
||||
const ulong mask_nb3,
|
||||
const int mask_ne2,
|
||||
const int mask_ne3,
|
||||
const global void* sinks_void,
|
||||
const ulong sinks_offset
|
||||
) {
|
||||
const int tid = get_local_id(0);
|
||||
const int sgid = tid / Q1_WG_SIZE;
|
||||
const int tid_sg = tid % Q1_WG_SIZE;
|
||||
const int head_batch_idx = get_global_id(1);
|
||||
|
||||
const int batch_idx = head_batch_idx / n_head;
|
||||
const int head_idx = head_batch_idx % n_head;
|
||||
|
||||
const int gqa_ratio = n_head / n_head_kv;
|
||||
const int head_kv_idx = head_idx / gqa_ratio;
|
||||
|
||||
const global char * q_base = (const global char *) q_void + q_offset;
|
||||
const global char * k_base = (const global char *) k_void + k_offset;
|
||||
const global char * v_base = (const global char *) v_void + v_offset;
|
||||
global char * o_base = (global char *) o_void + o_offset;
|
||||
|
||||
const global char * mask_base = NULL;
|
||||
if (mask_void != NULL) {
|
||||
const int mask_head_idx = head_idx % mask_ne2;
|
||||
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||
mask_base = (const global char *) mask_void + mask_offset +
|
||||
mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
|
||||
}
|
||||
|
||||
__local ACC_TYPE4 q_shared[DK_VEC];
|
||||
{
|
||||
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
|
||||
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||
for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
|
||||
q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
|
||||
|
||||
const global ACC_TYPE * sinks_ptr = NULL;
|
||||
if (sinks_void != NULL) {
|
||||
sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
|
||||
}
|
||||
|
||||
ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
|
||||
|
||||
ACC_TYPE m_i = FA_M_INIT;
|
||||
ACC_TYPE l_i = 0.0f;
|
||||
|
||||
const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
|
||||
const int kv_start = sgid * kv_per_sg;
|
||||
const int kv_end = min(n_kv, kv_start + kv_per_sg);
|
||||
|
||||
for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
|
||||
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||
|
||||
ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
|
||||
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||
const int block_idx = qk / 8;
|
||||
const int lane = qk % 8;
|
||||
const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||
dot4 = mad(q_shared[qk], k_v, dot4);
|
||||
}
|
||||
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
|
||||
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
|
||||
|
||||
if (mask_base != NULL) {
|
||||
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
|
||||
score += slope * (ACC_TYPE) mask_ptr[k_idx];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
score = logit_softcap * tanh(score / logit_softcap);
|
||||
}
|
||||
|
||||
const ACC_TYPE m_new = max(m_i, score);
|
||||
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
|
||||
const ACC_TYPE p = native_exp(score - m_new);
|
||||
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
const int block_idx = dv / 8;
|
||||
const int lane = dv % 8;
|
||||
const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||
o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
|
||||
}
|
||||
l_i = l_i * scale_prev + p;
|
||||
m_i = m_new;
|
||||
}
|
||||
|
||||
__local ACC_TYPE sg_m[VEC_NSG];
|
||||
__local ACC_TYPE sg_l[VEC_NSG];
|
||||
__local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
|
||||
|
||||
if (tid_sg == 0) {
|
||||
sg_m[sgid] = m_i;
|
||||
sg_l[sgid] = l_i;
|
||||
}
|
||||
{
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
sg_o[sgid][dv] = o_acc[idx];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
if (sgid == 0) {
|
||||
ACC_TYPE m_final = sg_m[0];
|
||||
#pragma unroll
|
||||
for (int s = 1; s < VEC_NSG; ++s) {
|
||||
m_final = max(m_final, sg_m[s]);
|
||||
}
|
||||
if (sinks_ptr != NULL) {
|
||||
m_final = max(m_final, sinks_ptr[head_idx]);
|
||||
}
|
||||
|
||||
ACC_TYPE l_final = 0.0f;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < VEC_NSG; ++s) {
|
||||
l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
|
||||
}
|
||||
if (sinks_ptr != NULL) {
|
||||
l_final += native_exp(sinks_ptr[head_idx] - m_final);
|
||||
}
|
||||
const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
|
||||
|
||||
const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
|
||||
global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
|
||||
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < VEC_NSG; ++s) {
|
||||
const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
|
||||
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
|
||||
}
|
||||
o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flash-decoding split pass for q8_0 KV. Partial record: [m, l, O[DV]].
|
||||
// Merge kernel from flash_attn_f32_f16.cl is type-agnostic and reused.
|
||||
#define FA_PARTIAL_FLOATS (2 + DV)
|
||||
@@ -533,6 +732,244 @@ __kernel void flash_attn_f32_q8_0_q1_split(
|
||||
#define FA_V_STRATEGY 0
|
||||
#endif
|
||||
|
||||
#ifndef MQ_GQA
|
||||
#define MQ_GQA 4
|
||||
#endif
|
||||
#ifndef MQ_NSG_SPLIT
|
||||
#define MQ_NSG_SPLIT 4
|
||||
#endif
|
||||
#define MQ_SPLIT_WG_SIZE_Q8 (Q1_WG_SIZE * MQ_NSG_SPLIT)
|
||||
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
__kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
const global void * v_void, ulong v_offset,
|
||||
const float scale,
|
||||
const int n_q,
|
||||
const int n_kv,
|
||||
const int n_head,
|
||||
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||
const float max_bias,
|
||||
const float m0,
|
||||
const float m1,
|
||||
const int n_head_log2,
|
||||
const float logit_softcap,
|
||||
const int n_head_kv,
|
||||
const global void * mask_void,
|
||||
const ulong mask_offset,
|
||||
const ulong mask_nb1,
|
||||
const ulong mask_nb2,
|
||||
const ulong mask_nb3,
|
||||
const int mask_ne2,
|
||||
const int mask_ne3,
|
||||
global float * partial_void,
|
||||
const int n_splits,
|
||||
const int kv_per_split
|
||||
) {
|
||||
const int tid = get_local_id(0);
|
||||
const int sgid = tid / Q1_WG_SIZE;
|
||||
const int tid_sg = tid % Q1_WG_SIZE;
|
||||
const int kvhead_batch_idx = get_global_id(1);
|
||||
const int split_q_idx = get_global_id(2);
|
||||
const int split_idx = split_q_idx % n_splits;
|
||||
const int q_idx = split_q_idx / n_splits;
|
||||
|
||||
const int batch_idx = kvhead_batch_idx / n_head_kv;
|
||||
const int head_kv_idx = kvhead_batch_idx % n_head_kv;
|
||||
|
||||
const int kv_start = split_idx * kv_per_split;
|
||||
const int kv_end = min(kv_start + kv_per_split, n_kv);
|
||||
|
||||
const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
|
||||
|
||||
if (kv_start >= kv_end) {
|
||||
// Empty split — write sentinel for each of the MQ_GQA Q-heads.
|
||||
if (tid == 0) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||
* n_splits + split_idx);
|
||||
global float * rec = partial_void + rec_idx * record_stride;
|
||||
rec[0] = FA_M_INIT;
|
||||
rec[1] = 0.0f;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const global char * q_base = (const global char *) q_void + q_offset;
|
||||
const global char * k_base = (const global char *) k_void + k_offset;
|
||||
const global char * v_base = (const global char *) v_void + v_offset;
|
||||
|
||||
__local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
|
||||
for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q8) {
|
||||
const int h = i / DK_VEC;
|
||||
const int k = i % DK_VEC;
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
|
||||
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||
q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
float slope[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
|
||||
}
|
||||
|
||||
const global char * mask_base[MQ_GQA];
|
||||
if (mask_void != NULL) {
|
||||
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||
const global char * mask_base_b = (const global char *) mask_void + mask_offset +
|
||||
mask_batch_idx * mask_nb3 +
|
||||
(ulong) q_idx * mask_nb1;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
const int mask_head_idx = head_idx % mask_ne2;
|
||||
mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
|
||||
}
|
||||
|
||||
ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
|
||||
ACC_TYPE m_i[MQ_GQA];
|
||||
ACC_TYPE l_i[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
m_i[h] = FA_M_INIT;
|
||||
l_i[h] = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
|
||||
}
|
||||
|
||||
const int kv_len = kv_end - kv_start;
|
||||
const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
|
||||
const int kv_lo = kv_start + sgid * kv_per_sg;
|
||||
const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
|
||||
|
||||
for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
|
||||
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||
|
||||
ACC_TYPE4 dot4[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
|
||||
|
||||
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||
const int block_idx = qk / 8;
|
||||
const int lane = qk % 8;
|
||||
const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
|
||||
}
|
||||
}
|
||||
|
||||
ACC_TYPE score[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
|
||||
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
|
||||
if (mask_base[h] != NULL) {
|
||||
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
|
||||
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
s = logit_softcap * tanh(s / logit_softcap);
|
||||
}
|
||||
score[h] = s;
|
||||
}
|
||||
|
||||
ACC_TYPE p_h[MQ_GQA];
|
||||
ACC_TYPE sp_h[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
const ACC_TYPE m_new = max(m_i[h], score[h]);
|
||||
sp_h[h] = native_exp(m_i[h] - m_new);
|
||||
p_h[h] = native_exp(score[h] - m_new);
|
||||
l_i[h] = l_i[h] * sp_h[h] + p_h[h];
|
||||
m_i[h] = m_new;
|
||||
}
|
||||
|
||||
int idx = 0;
|
||||
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||
const int block_idx = dv / 8;
|
||||
const int lane = dv % 8;
|
||||
const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
|
||||
__local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
|
||||
__local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
|
||||
|
||||
if (tid_sg == 0) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
sg_m[h][sgid] = m_i[h];
|
||||
sg_l[h][sgid] = l_i[h];
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
{
|
||||
int idx = 0;
|
||||
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
|
||||
sg_o[sgid][dv_idx] = o_acc[h][idx];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
if (sgid == 0) {
|
||||
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||
|
||||
ACC_TYPE m_c = sg_m[h][0];
|
||||
#pragma unroll
|
||||
for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
|
||||
m_c = max(m_c, sg_m[h][s]);
|
||||
}
|
||||
ACC_TYPE l_c = 0.0f;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||
l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
|
||||
}
|
||||
|
||||
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||
* n_splits + split_idx);
|
||||
global float * rec = partial_void + rec_idx * record_stride;
|
||||
global float4 * rec_o = (global float4 *) (rec + 2);
|
||||
|
||||
if (tid_sg == 0) {
|
||||
rec[0] = (float) m_c;
|
||||
rec[1] = (float) l_c;
|
||||
}
|
||||
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
|
||||
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||
const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
|
||||
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
|
||||
}
|
||||
rec_o[dv_idx] = o_merged;
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
}
|
||||
|
||||
__kernel void flash_attn_f32_q8_0(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
|
||||
@@ -18,6 +18,14 @@
|
||||
#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
|
||||
#endif
|
||||
|
||||
#ifdef cl_khr_subgroup_shuffle
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
|
||||
#define HAS_SUBGROUP_SHUFFLE 1
|
||||
#elif defined(cl_qcom_subgroup_shuffle)
|
||||
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
|
||||
#define HAS_SUBGROUP_SHUFFLE 1
|
||||
#endif
|
||||
|
||||
// Assumes row size (ne00) is a multiple of 4
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
@@ -378,3 +386,848 @@ kernel void kernel_mul_mat_f16_f32_l4_dr_lq(
|
||||
}
|
||||
}
|
||||
#endif // ADRENO_GPU
|
||||
|
||||
#define N_ROWS_PER_WG 8
|
||||
#define N_OUTS_PER_WG 8
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_ROWS_PER_WG;
|
||||
const int im = get_group_id(2);
|
||||
|
||||
const int i12 = im % ne12;
|
||||
const int i13 = im / ne12;
|
||||
|
||||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||||
|
||||
__local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
|
||||
if (sgs_lid < ne00 / 4) {
|
||||
q_loc[sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_ROWS_PER_WG; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
if (r0 >= ne01) return;
|
||||
|
||||
const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
const half4 k4 = x4[i];
|
||||
const float4 q = q_loc[i];
|
||||
sumf += convert_float(k4.s0) * q.s0
|
||||
+ convert_float(k4.s1) * q.s1
|
||||
+ convert_float(k4.s2) * q.s2
|
||||
+ convert_float(k4.s3) * q.s3;
|
||||
}
|
||||
|
||||
const float all_sum = sub_group_reduce_add(sumf);
|
||||
if (sgs_lid == 0) {
|
||||
dst[im * ne1 * ne0 + r0] = all_sum; // ne11 == 1, so r1==0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_y8(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_OUTS_PER_WG;
|
||||
const int im = get_group_id(2);
|
||||
|
||||
const int i12 = im % ne12;
|
||||
const int i13 = im / ne12;
|
||||
|
||||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||||
|
||||
global half4 * x4_o[N_OUTS_PER_WG];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||||
const ulong off = r0c * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||||
x4_o[o] = (global half4 *)(src0 + off);
|
||||
}
|
||||
|
||||
float sum[N_OUTS_PER_WG] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
|
||||
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
const float4 q4 = y4[i];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||||
const half4 v4 = x4_o[o][i];
|
||||
sum[o] += convert_float(v4.s0) * q4.s0
|
||||
+ convert_float(v4.s1) * q4.s1
|
||||
+ convert_float(v4.s2) * q4.s2
|
||||
+ convert_float(v4.s3) * q4.s3;
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const float s = sub_group_reduce_add(sum[o]);
|
||||
if (sgs_lid == 0 && r0 < ne01) {
|
||||
dst[im * ne1 * ne0 + r0] = s;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_OUTS_PAIR 8
|
||||
#define N_PAIRS_PAIR (N_OUTS_PAIR / 2)
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_pair(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int half_id = sgs_lid >> 5; // 0 = lower half, 1 = upper half
|
||||
const int lane_h = sgs_lid & 31; // lane 0..31 within half
|
||||
|
||||
const int r0_base = get_group_id(0) * N_OUTS_PAIR;
|
||||
const int im = get_group_id(2);
|
||||
|
||||
const int i12 = im % ne12;
|
||||
const int i13 = im / ne12;
|
||||
|
||||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||||
|
||||
__local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
|
||||
if (sgs_lid < ne00 / 4) {
|
||||
q_loc[sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int dk_vec = ne00 / 4;
|
||||
|
||||
#pragma unroll
|
||||
for (int p = 0; p < N_PAIRS_PAIR; ++p) {
|
||||
const int r0 = r0_base + 2 * p + half_id;
|
||||
|
||||
const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = lane_h; i < dk_vec; i += 32) {
|
||||
const half4 k4 = x4[i];
|
||||
const float4 q = q_loc[i];
|
||||
sumf += convert_float(k4.s0) * q.s0
|
||||
+ convert_float(k4.s1) * q.s1
|
||||
+ convert_float(k4.s2) * q.s2
|
||||
+ convert_float(k4.s3) * q.s3;
|
||||
}
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 16);
|
||||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_h == 0) {
|
||||
dst[im * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_K_ROWS_GQA 16
|
||||
#define GQA_RATIO_GQA 8
|
||||
#define LANES_PER_QH 8 // 64 / GQA_RATIO_GQA
|
||||
#define DK_VEC_GQA 32 // DK / 4 for DK=128
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
|
||||
const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02; // K-head index (also K2 batch)
|
||||
const int i03 = im_kv / ne02; // n13 batch index
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_GQA;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA]; // 4 × 32 = 128 float4
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
|
||||
if (sgs_lid < DK_VEC_GQA) {
|
||||
q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
// K base offset for this WG. All 8 K-rows × 4 Q-heads share this K-head.
|
||||
const ulong offset_src0_base = (i02) * nb02 + (i03 / r3) * nb03;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
|
||||
const ulong offset_src0 = r0 * nb01 + offset_src0_base;
|
||||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||||
|
||||
float sumf = 0.0f;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < 4; ++t) {
|
||||
const int i = lane_q + t * LANES_PER_QH; // 8, 16, 24-step
|
||||
const half4 k4 = x4[i];
|
||||
const float4 q = q_loc[q_id * DK_VEC_GQA + i];
|
||||
sumf += convert_float(k4.s0) * q.s0
|
||||
+ convert_float(k4.s1) * q.s1
|
||||
+ convert_float(k4.s2) * q.s2
|
||||
+ convert_float(k4.s3) * q.s3;
|
||||
}
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_DV_ROWS_Y8GQA 8
|
||||
#define GQA_RATIO_Y8GQA 8
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_y8_gqa(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02; // K-head index
|
||||
const int i03 = im_kv / ne02; // n13 batch index
|
||||
|
||||
// GQA Q-heads sharing this K-head.
|
||||
const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
|
||||
|
||||
global float4 * y4_q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
}
|
||||
|
||||
global half4 * x4_o[N_DV_ROWS_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||||
const ulong off = r0c * nb01 + (i02) * nb02 + (i03 / r3) * nb03;
|
||||
x4_o[o] = (global half4 *)(src0 + off);
|
||||
}
|
||||
|
||||
float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
|
||||
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
// load 8 V values (one per DV row), same K-head, K-pos = i.
|
||||
half4 v[N_DV_ROWS_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
v[o] = x4_o[o][i];
|
||||
}
|
||||
|
||||
// load 8 softmax values (one per Q-head).
|
||||
float4 q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
q[qh] = y4_q[qh][i];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const float4 vf = (float4)(convert_float(v[o].s0),
|
||||
convert_float(v[o].s1),
|
||||
convert_float(v[o].s2),
|
||||
convert_float(v[o].s3));
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
sum[o][qh] += vf.s0 * q[qh].s0
|
||||
+ vf.s1 * q[qh].s1
|
||||
+ vf.s2 * q[qh].s2
|
||||
+ vf.s3 * q[qh].s3;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const float s = sub_group_reduce_add(sum[o][qh]);
|
||||
if (sgs_lid == 0 && r0 < ne01) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + qh);
|
||||
dst[im_out * ne1 * ne0 + r0] = s;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
|
||||
const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_GQA;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
if (sgs_lid < DK_VEC_GQA) {
|
||||
q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 4);
|
||||
const int pitch_px_head = (int)(nb02 >> 4);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||||
|
||||
float sumf = 0.0f;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < 2; ++t) {
|
||||
const int p = lane_q + t * LANES_PER_QH; // pixel idx in row, 0..15
|
||||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||||
const int i0 = 2 * p; // first half4 idx
|
||||
const float4 qa = q_loc[q_id * DK_VEC_GQA + i0 ];
|
||||
const float4 qb = q_loc[q_id * DK_VEC_GQA + i0 + 1];
|
||||
sumf += convert_float(k8.s0) * qa.s0
|
||||
+ convert_float(k8.s1) * qa.s1
|
||||
+ convert_float(k8.s2) * qa.s2
|
||||
+ convert_float(k8.s3) * qa.s3
|
||||
+ convert_float(k8.s4) * qb.s0
|
||||
+ convert_float(k8.s5) * qb.s1
|
||||
+ convert_float(k8.s6) * qb.s2
|
||||
+ convert_float(k8.s7) * qb.s3;
|
||||
}
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_y8_gqa_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
|
||||
|
||||
// Q (= softmax(KQ)) base pointers per Q-head
|
||||
global float4 * y4_q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
}
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 3);
|
||||
const int pitch_px_head = (int)(nb02 >> 3);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 3);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
// per-DV-row pixel base
|
||||
int row_px_base[N_DV_ROWS_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||||
row_px_base[o] = r0c * pitch_px_row + head_px_base;
|
||||
}
|
||||
|
||||
float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
|
||||
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
half4 v[N_DV_ROWS_Y8GQA];
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
v[o] = read_imageh(src0_img, row_px_base[o] + i);
|
||||
}
|
||||
|
||||
float4 q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
q[qh] = y4_q[qh][i];
|
||||
}
|
||||
// 64 mads.
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const float4 vf = (float4)(convert_float(v[o].s0),
|
||||
convert_float(v[o].s1),
|
||||
convert_float(v[o].s2),
|
||||
convert_float(v[o].s3));
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
sum[o][qh] += vf.s0 * q[qh].s0
|
||||
+ vf.s1 * q[qh].s1
|
||||
+ vf.s2 * q[qh].s2
|
||||
+ vf.s3 * q[qh].s3;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const float s = sub_group_reduce_add(sum[o][qh]);
|
||||
if (sgs_lid == 0 && r0 < ne01) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + qh);
|
||||
dst[im_out * ne1 * ne0 + r0] = s;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_K_ROWS_GQA_R4 16
|
||||
#define GQA_RATIO_R4 4
|
||||
#define LANES_PER_QH_R4 16 // = 64 / GQA_RATIO_R4
|
||||
#define DK_VEC_R4 32 // DK / 4 for DK=128
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 4; // 0..3
|
||||
const int lane_q = sgs_lid & 15; // 0..15
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R4;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_R4;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_R4 * DK_VEC_R4];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_R4; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
if (sgs_lid < DK_VEC_R4) {
|
||||
q_loc[qh * DK_VEC_R4 + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 4);
|
||||
const int pitch_px_head = (int)(nb02 >> 4);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA_R4; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||||
|
||||
const int p = lane_q;
|
||||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||||
const int i0 = 2 * p;
|
||||
const float4 qa = q_loc[q_id * DK_VEC_R4 + i0 ];
|
||||
const float4 qb = q_loc[q_id * DK_VEC_R4 + i0 + 1];
|
||||
|
||||
float sumf =
|
||||
convert_float(k8.s0) * qa.s0
|
||||
+ convert_float(k8.s1) * qa.s1
|
||||
+ convert_float(k8.s2) * qa.s2
|
||||
+ convert_float(k8.s3) * qa.s3
|
||||
+ convert_float(k8.s4) * qb.s0
|
||||
+ convert_float(k8.s5) * qb.s1
|
||||
+ convert_float(k8.s6) * qb.s2
|
||||
+ convert_float(k8.s7) * qb.s3;
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_K_ROWS_GQA_R2_DK256 16
|
||||
#define GQA_RATIO_R2 2
|
||||
#define LANES_PER_QH_R2 32 // = 64 / GQA_RATIO_R2
|
||||
#define DK_VEC_DK256 64 // DK / 4 for DK=256
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 5; // 0..1
|
||||
const int lane_q = sgs_lid & 31; // 0..31
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R2_DK256;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_R2;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_R2 * DK_VEC_DK256];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_R2; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
q_loc[qh * DK_VEC_DK256 + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 4);
|
||||
const int pitch_px_head = (int)(nb02 >> 4);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA_R2_DK256; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||||
|
||||
const int p = lane_q;
|
||||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||||
const int i0 = 2 * p;
|
||||
const float4 qa = q_loc[q_id * DK_VEC_DK256 + i0 ];
|
||||
const float4 qb = q_loc[q_id * DK_VEC_DK256 + i0 + 1];
|
||||
|
||||
float sumf =
|
||||
convert_float(k8.s0) * qa.s0
|
||||
+ convert_float(k8.s1) * qa.s1
|
||||
+ convert_float(k8.s2) * qa.s2
|
||||
+ convert_float(k8.s3) * qa.s3
|
||||
+ convert_float(k8.s4) * qb.s0
|
||||
+ convert_float(k8.s5) * qb.s1
|
||||
+ convert_float(k8.s6) * qb.s2
|
||||
+ convert_float(k8.s7) * qb.s3;
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 16);
|
||||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -129,7 +129,7 @@ typedef struct VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE {
|
||||
#endif
|
||||
|
||||
#define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
|
||||
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
|
||||
#define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
|
||||
static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
|
||||
|
||||
#define VK_VENDOR_ID_AMD 0x1002
|
||||
@@ -17370,21 +17370,24 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
|
||||
return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32;
|
||||
case GGML_OP_SET_ROWS:
|
||||
{
|
||||
switch (op->type) {
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_F16:
|
||||
case GGML_TYPE_BF16:
|
||||
case GGML_TYPE_Q1_0:
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_IQ4_NL:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
if (op->src[0]->type == GGML_TYPE_F32) {
|
||||
switch (op->type) {
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_F16:
|
||||
case GGML_TYPE_BF16:
|
||||
case GGML_TYPE_Q1_0:
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_IQ4_NL:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
case GGML_OP_CONT:
|
||||
case GGML_OP_CPY:
|
||||
|
||||
@@ -186,6 +186,12 @@ function(hf_download version out_var out_resolved)
|
||||
|
||||
set(archive "${UI_BINARY_DIR}/dist.tar.gz")
|
||||
|
||||
# Use HF_TOKEN to benefit from higher rate limits
|
||||
set(auth_headers "")
|
||||
if(DEFINED ENV{HF_TOKEN} AND NOT "$ENV{HF_TOKEN}" STREQUAL "")
|
||||
list(APPEND auth_headers "HTTPHEADER" "Authorization: Bearer $ENV{HF_TOKEN}")
|
||||
endif()
|
||||
|
||||
set(candidates "")
|
||||
if(NOT "${version}" STREQUAL "")
|
||||
list(APPEND candidates "${version}")
|
||||
@@ -198,7 +204,7 @@ function(hf_download version out_var out_resolved)
|
||||
message(STATUS "UI: downloading from ${resolved}: ${base}/dist.tar.gz")
|
||||
|
||||
file(DOWNLOAD "${base}/dist.tar.gz?download=true" "${archive}"
|
||||
STATUS status TIMEOUT 300
|
||||
STATUS status TIMEOUT 300 ${auth_headers}
|
||||
)
|
||||
list(GET status 0 rc)
|
||||
if(NOT rc EQUAL 0)
|
||||
@@ -208,7 +214,7 @@ function(hf_download version out_var out_resolved)
|
||||
endif()
|
||||
|
||||
file(DOWNLOAD "${base}/dist.tar.gz.sha256?download=true" "${archive}.sha256"
|
||||
STATUS status TIMEOUT 30
|
||||
STATUS status TIMEOUT 30 ${auth_headers}
|
||||
)
|
||||
list(GET status 0 rc)
|
||||
if(NOT rc EQUAL 0)
|
||||
|
||||
@@ -953,6 +953,8 @@ static buft_list_t make_gpu_buft_list(ggml_backend_dev_t dev, llama_split_mode s
|
||||
if (buft != nullptr) {
|
||||
buft_list.emplace_back(dev, buft);
|
||||
}
|
||||
} else {
|
||||
throw std::runtime_error(format("device %s does not support split buffers", ggml_backend_dev_name(dev)));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -523,6 +523,7 @@ void server_models::load_models() {
|
||||
|
||||
// collect all threads to join in one pass while the lock is held:
|
||||
// - monitoring threads from just-unloaded models (to_unload)
|
||||
// - threads of finished downloads (DOWNLOADED), they acquire the mutex on exit
|
||||
// - threads of already-UNLOADED models that are being removed from source
|
||||
std::vector<std::thread> threads_to_join;
|
||||
for (const auto & name : to_unload) {
|
||||
@@ -535,6 +536,13 @@ void server_models::load_models() {
|
||||
if (inst.meta.status == SERVER_MODEL_STATUS_DOWNLOADING) {
|
||||
continue; // downloading models are not from config sources, leave them alone
|
||||
}
|
||||
if (inst.meta.status == SERVER_MODEL_STATUS_DOWNLOADED) {
|
||||
// joining this thread under the lock deadlocks: it locks the mutex on its way out
|
||||
if (inst.th.joinable()) {
|
||||
threads_to_join.push_back(std::move(inst.th));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (final_presets.find(name) == final_presets.end() && !inst.meta.is_running() && inst.th.joinable()) {
|
||||
threads_to_join.push_back(std::move(inst.th));
|
||||
}
|
||||
@@ -550,10 +558,8 @@ void server_models::load_models() {
|
||||
if (it->second.meta.status == SERVER_MODEL_STATUS_DOWNLOADING) {
|
||||
++it; // download thread is still busy, skip
|
||||
} else if (it->second.meta.status == SERVER_MODEL_STATUS_DOWNLOADED) {
|
||||
// download finished, safe to erase
|
||||
if (it->second.th.joinable()) {
|
||||
it->second.th.join();
|
||||
}
|
||||
// download finished, thread is joined above, safe to erase
|
||||
GGML_ASSERT(!it->second.th.joinable());
|
||||
it = mapping.erase(it);
|
||||
} else if (final_presets.find(it->first) == final_presets.end()) {
|
||||
SRV_INF("(reload) removing model name=%s (no longer in source)\n", it->first.c_str());
|
||||
|
||||
@@ -31,6 +31,9 @@ import wget
|
||||
|
||||
DEFAULT_HTTP_TIMEOUT = 60
|
||||
|
||||
# per-request timeout, a hung server fails the test instead of stalling the CI for hours
|
||||
DEFAULT_REQUEST_TIMEOUT = 600
|
||||
|
||||
|
||||
class ServerResponse:
|
||||
headers: dict
|
||||
@@ -330,7 +333,7 @@ class ServerProcess:
|
||||
path: str,
|
||||
data: dict | Any | None = None,
|
||||
headers: dict | None = None,
|
||||
timeout: float | None = None,
|
||||
timeout: float | None = DEFAULT_REQUEST_TIMEOUT,
|
||||
) -> ServerResponse:
|
||||
url = f"http://{self.server_host}:{self.server_port}{path}"
|
||||
parse_body = False
|
||||
@@ -389,7 +392,7 @@ class ServerProcess:
|
||||
path: str,
|
||||
data: dict | None = None,
|
||||
headers: dict | None = None,
|
||||
timeout: float | None = None,
|
||||
timeout: float | None = DEFAULT_REQUEST_TIMEOUT,
|
||||
) -> dict:
|
||||
stream = data.get('stream', False)
|
||||
if stream:
|
||||
|
||||
+4
-1
@@ -27,7 +27,10 @@
|
||||
|
||||
let { onSearchClick = () => {} }: Props = $props();
|
||||
|
||||
const { handleKeydown } = useKeyboardShortcuts({ activateSearchMode: () => onSearchClick() });
|
||||
const { handleKeydown } = useKeyboardShortcuts({
|
||||
activateSearchMode: () => onSearchClick(),
|
||||
toggleSidebar: () => toggleExpandedMode()
|
||||
});
|
||||
|
||||
let isExpandedMode = $state(false);
|
||||
let hoveredTooltip = $state<string | null>(null);
|
||||
|
||||
@@ -9,6 +9,7 @@ export enum KeyboardKey {
|
||||
ARROW_LEFT = 'ArrowLeft',
|
||||
ARROW_RIGHT = 'ArrowRight',
|
||||
TAB = 'Tab',
|
||||
B_LOWER = 'b',
|
||||
D_LOWER = 'd',
|
||||
D_UPPER = 'D',
|
||||
E_UPPER = 'E',
|
||||
|
||||
@@ -9,6 +9,7 @@ interface KeyboardShortcutsCallbacks {
|
||||
deleteActiveConversation?: () => void;
|
||||
navigateToPrevConversation?: () => void;
|
||||
navigateToNextConversation?: () => void;
|
||||
toggleSidebar?: () => void;
|
||||
}
|
||||
|
||||
export function useKeyboardShortcuts(callbacks: KeyboardShortcutsCallbacks) {
|
||||
@@ -21,6 +22,11 @@ export function useKeyboardShortcuts(callbacks: KeyboardShortcutsCallbacks) {
|
||||
callbacks.onSearchActivated?.();
|
||||
}
|
||||
|
||||
if (isCmdOrCtrl && event.key === KeyboardKey.B_LOWER) {
|
||||
event.preventDefault();
|
||||
callbacks.toggleSidebar?.();
|
||||
}
|
||||
|
||||
if (
|
||||
isCmdOrCtrl &&
|
||||
event.shiftKey &&
|
||||
|
||||
@@ -314,6 +314,30 @@ export class MCPService {
|
||||
)
|
||||
);
|
||||
|
||||
if (method === 'DELETE' && url.includes(CORS_PROXY_ENDPOINT)) {
|
||||
const response = new Response(null, { status: 200, statusText: 'OK' });
|
||||
|
||||
logIfEnabled(
|
||||
this.createLog(
|
||||
MCPConnectionPhase.INITIALIZING,
|
||||
`HTTP 200 ${method} ${url} (fake response)`,
|
||||
MCPLogLevel.INFO,
|
||||
{
|
||||
response: {
|
||||
url,
|
||||
status: response.status,
|
||||
statusText: response.statusText,
|
||||
durationMs: 0,
|
||||
isFake: true
|
||||
}
|
||||
}
|
||||
)
|
||||
);
|
||||
|
||||
// fake response, bypass real fetch()
|
||||
return response;
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await fetch(input, {
|
||||
...baseInit,
|
||||
|
||||
@@ -154,6 +154,32 @@ describe('MCPService', () => {
|
||||
});
|
||||
});
|
||||
|
||||
it('DELETE request with CORS proxy should return a fake 200 response', async () => {
|
||||
const logs: MCPConnectionLog[] = [];
|
||||
const fetchMock = vi.fn();
|
||||
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const config: MCPServerConfig = {
|
||||
url: 'https://example.com/mcp',
|
||||
transport: MCPTransportType.STREAMABLE_HTTP,
|
||||
useProxy: true
|
||||
};
|
||||
|
||||
const controller = createDiagnosticFetch(config, (log) => logs.push(log), {}, true);
|
||||
|
||||
const response = await controller.fetch(
|
||||
'http://localhost:8080/cors-proxy?url=https%3A%2F%2Fexample.com%2Fmcp',
|
||||
{ method: 'DELETE' }
|
||||
);
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(response.status).toBe(200);
|
||||
expect(logs.at(-1)?.details).toMatchObject({
|
||||
response: { status: 200, isFake: true }
|
||||
});
|
||||
});
|
||||
|
||||
it('partially redacts mcp-session-id in diagnostic request and response logs', async () => {
|
||||
const logs: MCPConnectionLog[] = [];
|
||||
const response = new Response('{}', {
|
||||
|
||||
Reference in New Issue
Block a user