mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-10 14:35:53 +02:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 049326a000 | |||
| 074944998d | |||
| 3de7dd4c8f | |||
| fb30ba9a6c |
+1
-1
@@ -2849,7 +2849,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
params.out_file = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_CVECTOR_GENERATOR, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_FINETUNE,
|
||||
LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS}));
|
||||
LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, LLAMA_EXAMPLE_CLI}));
|
||||
add_opt(common_arg(
|
||||
{"-ofreq", "--output-frequency"}, "N",
|
||||
string_format("output the imatrix every N iterations (default: %d)", params.n_out_freq),
|
||||
|
||||
@@ -28,6 +28,20 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr
|
||||
#endif // STRIDED_ITERATOR_AVAILABLE
|
||||
|
||||
#ifdef GGML_CUDA_USE_CUB
|
||||
|
||||
// returns the suggested maximum number of rows to process during one argsort_f32_i32_cuda_cub() call
|
||||
int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows) {
|
||||
// perform argsort in chunks up to approximately this size (currently 64MB)
|
||||
// to avoid excessive temporary buffers memory usage
|
||||
const int chunk_bytes = 1 << 26;
|
||||
|
||||
// calculate how many rows will fit in one chunk (must be at least one)
|
||||
const int chunk_nrows = std::max((int) (chunk_bytes / nb01), 1);
|
||||
|
||||
// limit the resulting amount to total nrows
|
||||
return std::min((int64_t) chunk_nrows, nrows);
|
||||
}
|
||||
|
||||
void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
|
||||
const float * x,
|
||||
int * dst,
|
||||
@@ -254,11 +268,23 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const size_t shared_mem = ncols_pad * sizeof(int);
|
||||
const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;
|
||||
|
||||
if (shared_mem > max_shared_mem || ncols > 1024) {
|
||||
ggml_cuda_pool & pool = ctx.pool();
|
||||
argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream);
|
||||
} else {
|
||||
// early return if we can use bitonic argsort
|
||||
if (shared_mem <= max_shared_mem && ncols <= 1024) {
|
||||
argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
|
||||
return;
|
||||
}
|
||||
|
||||
const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows);
|
||||
|
||||
ggml_cuda_pool & pool = ctx.pool();
|
||||
|
||||
for (int64_t i = 0; i < nrows; i += chunk_nrows) {
|
||||
int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i);
|
||||
|
||||
argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, iter_nrows, order, stream);
|
||||
|
||||
src0_d += ncols * iter_nrows;
|
||||
dst_d += ncols * iter_nrows;
|
||||
}
|
||||
#else
|
||||
argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
|
||||
|
||||
#ifdef GGML_CUDA_USE_CUB
|
||||
int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows);
|
||||
void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
|
||||
const float * x,
|
||||
int * dst,
|
||||
|
||||
@@ -75,17 +75,26 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const int ncols_pad = next_power_of_2(ncols);
|
||||
const size_t shared_mem = ncols_pad * sizeof(int);
|
||||
const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;
|
||||
const bool use_bitonic = shared_mem <= max_shared_mem && ncols <= 1024;
|
||||
const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows);
|
||||
|
||||
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
|
||||
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * chunk_nrows);
|
||||
int * tmp_dst = temp_dst_alloc.get();
|
||||
|
||||
if (shared_mem > max_shared_mem || ncols > 1024) {
|
||||
argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
|
||||
} else {
|
||||
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
|
||||
for (int64_t i = 0; i < nrows; i += chunk_nrows) {
|
||||
int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i);
|
||||
|
||||
if (use_bitonic) {
|
||||
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream);
|
||||
} else {
|
||||
argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream);
|
||||
}
|
||||
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), iter_nrows,
|
||||
cudaMemcpyDeviceToDevice, stream));
|
||||
|
||||
src0_d += ncols * iter_nrows;
|
||||
dst_d += k * iter_nrows;
|
||||
}
|
||||
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
|
||||
cudaMemcpyDeviceToDevice, stream));
|
||||
#else // GGML_CUDA_USE_CUB
|
||||
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
|
||||
int * tmp_dst = temp_dst_alloc.get();
|
||||
|
||||
@@ -44,6 +44,7 @@
|
||||
#include "htp-ops.h"
|
||||
#include "htp/matmul-ops.h"
|
||||
#include "htp/flash-attn-ops.h"
|
||||
#include "htp/unary-ops.h"
|
||||
#include "htp_iface.h"
|
||||
#include "htp-drv.h"
|
||||
|
||||
@@ -170,8 +171,8 @@ static inline bool ggml_hexagon_is_hmx_weight_type(enum ggml_type type) {
|
||||
return type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || ggml_hexagon_is_repack_type(type);
|
||||
}
|
||||
|
||||
struct htp_mm_kernel_params;
|
||||
struct ggml_hexagon_session;
|
||||
|
||||
static void ggml_hexagon_precompute_matmul_params(
|
||||
const struct ggml_hexagon_session * sess,
|
||||
const struct ggml_tensor * src0,
|
||||
@@ -180,6 +181,15 @@ static void ggml_hexagon_precompute_matmul_params(
|
||||
struct htp_mm_kernel_params * kparams
|
||||
);
|
||||
|
||||
static void ggml_hexagon_precompute_unary_params(
|
||||
const struct ggml_hexagon_session * sess,
|
||||
uint32_t op,
|
||||
const struct ggml_tensor * src0,
|
||||
const struct ggml_tensor * src1,
|
||||
const struct ggml_tensor * dst,
|
||||
struct htp_unary_kernel_params * kparams
|
||||
);
|
||||
|
||||
static void ggml_hexagon_precompute_fused_qkv_params(
|
||||
const struct ggml_hexagon_session * sess,
|
||||
const struct ggml_tensor * src0,
|
||||
@@ -2591,6 +2601,74 @@ finalize:
|
||||
kparams->div_ne11 = init_fastdiv_values(ne11);
|
||||
}
|
||||
|
||||
static void ggml_hexagon_precompute_unary_params(
|
||||
const struct ggml_hexagon_session * sess,
|
||||
uint32_t op,
|
||||
const struct ggml_tensor * src0,
|
||||
const struct ggml_tensor * src1,
|
||||
const struct ggml_tensor * dst,
|
||||
struct htp_unary_kernel_params * kparams
|
||||
) {
|
||||
memset(kparams, 0, sizeof(*kparams));
|
||||
|
||||
const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
|
||||
const uint32_t n_threads = (std::min)((uint32_t)sess->n_threads, src0_nrows);
|
||||
|
||||
kparams->n_threads = n_threads;
|
||||
|
||||
const size_t src0_data_row_size = src0->ne[0] * sizeof(float);
|
||||
const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
|
||||
|
||||
const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
|
||||
const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
|
||||
|
||||
kparams->src0_row_size_aligned = src0_row_size_aligned;
|
||||
kparams->dst_row_size_aligned = dst_row_size_aligned;
|
||||
|
||||
size_t src1_data_row_size = 0;
|
||||
size_t src1_row_size_aligned = 0;
|
||||
bool broadcast_weight = false;
|
||||
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
GGML_ASSERT(src1 != nullptr);
|
||||
src1_data_row_size = src1->ne[0] * sizeof(float);
|
||||
src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
|
||||
broadcast_weight = (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
|
||||
}
|
||||
|
||||
kparams->src1_row_size_aligned = src1_row_size_aligned;
|
||||
kparams->broadcast_weight = broadcast_weight;
|
||||
|
||||
struct htp_unary_vtcm_layout L;
|
||||
uint32_t col_tile = 0;
|
||||
uint32_t vtcm_row_per_thread = 0;
|
||||
|
||||
htp_unary_vtcm_layout_build(&L, op, src0->ne[0], dst->ne[0],
|
||||
op == HTP_OP_RMS_NORM_MUL ? src1->ne[0] : 0,
|
||||
broadcast_weight, n_threads, sess->vtcm_size,
|
||||
&col_tile, &vtcm_row_per_thread);
|
||||
|
||||
kparams->col_tile = col_tile;
|
||||
kparams->vtcm_row_per_thread = vtcm_row_per_thread;
|
||||
kparams->vtcm_size = L.total_bytes;
|
||||
|
||||
kparams->vtcm_src0_size_per_thread = L.src0_bytes;
|
||||
kparams->vtcm_src1_size_per_thread = L.src1_bytes;
|
||||
kparams->vtcm_dst_size_per_thread = L.dst_bytes;
|
||||
|
||||
kparams->vtcm_src0_size = L.src0_bytes * n_threads;
|
||||
kparams->vtcm_src1_size = L.src1_bytes * n_threads;
|
||||
kparams->vtcm_dst_size = L.dst_bytes * n_threads;
|
||||
|
||||
kparams->block = col_tile ? 0 : ((L.src0_bytes / 2) / src0_row_size_aligned);
|
||||
|
||||
const uint32_t tiles_per_row = col_tile > 0 ? (src0->ne[0] + col_tile - 1) / col_tile : 1;
|
||||
kparams->div_ne01 = init_fastdiv_values(src0->ne[1]);
|
||||
kparams->div_ne02 = init_fastdiv_values(src0->ne[2]);
|
||||
kparams->div_ne012 = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
|
||||
kparams->div_tpr = init_fastdiv_values(tiles_per_row);
|
||||
}
|
||||
|
||||
static void ggml_hexagon_precompute_fused_qkv_params(
|
||||
const struct ggml_hexagon_session * sess,
|
||||
const struct ggml_tensor * src0, // Wk
|
||||
@@ -2866,6 +2944,9 @@ static bool ggml_hexagon_supported_binary(const struct ggml_hexagon_session * se
|
||||
return false;
|
||||
}
|
||||
|
||||
if (ggml_is_permuted(src0) || ggml_is_permuted(dst)) {
|
||||
return false;
|
||||
}
|
||||
if (!ggml_are_same_shape(src0, dst)) {
|
||||
return false;
|
||||
}
|
||||
@@ -2912,6 +2993,9 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
|
||||
if (dst->type != GGML_TYPE_F32) {
|
||||
return false;
|
||||
}
|
||||
if (ggml_is_permuted(src0)) {
|
||||
return false;
|
||||
}
|
||||
if (!ggml_are_same_shape(src0, dst)) {
|
||||
return false;
|
||||
}
|
||||
@@ -3451,6 +3535,15 @@ static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph *
|
||||
if (next_node->op == GGML_OP_MUL && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) {
|
||||
htp_opnode node(n, {}, HTP_OP_RMS_NORM_MUL);
|
||||
node.add_fused(next_node);
|
||||
|
||||
auto inputs = node.get_inputs();
|
||||
const struct ggml_tensor * src0 = inputs[0];
|
||||
const struct ggml_tensor * src1 = inputs.size() > 1 ? inputs[1] : nullptr;
|
||||
ggml_hexagon_precompute_unary_params(sess,
|
||||
node.opcode, src0, src1, node.dst(),
|
||||
(struct htp_unary_kernel_params *)node.kernel_params
|
||||
);
|
||||
|
||||
nodes.push_back(std::move(node));
|
||||
i++; // skip the fused MUL node
|
||||
return true;
|
||||
@@ -3555,6 +3648,14 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
|
||||
node.node,
|
||||
(struct htp_fa_kernel_params *)node.kernel_params
|
||||
);
|
||||
} else if (htp_op_is_unary(node.opcode)) {
|
||||
auto inputs = node.get_inputs();
|
||||
const struct ggml_tensor * src0 = inputs[0];
|
||||
const struct ggml_tensor * src1 = inputs.size() > 1 ? inputs[1] : nullptr;
|
||||
ggml_hexagon_precompute_unary_params(sess,
|
||||
node.opcode, src0, src1, node.dst(),
|
||||
(struct htp_unary_kernel_params *)node.kernel_params
|
||||
);
|
||||
}
|
||||
computed_nodes.push_back(std::move(node));
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
#include "htp-ops.h"
|
||||
#include "htp/matmul-ops.h"
|
||||
#include "htp/flash-attn-ops.h"
|
||||
#include "htp/unary-ops.h"
|
||||
|
||||
struct htp_opnode {
|
||||
ggml_tensor * node = nullptr;
|
||||
@@ -362,6 +363,9 @@ struct htp_opformat {
|
||||
path = "hvx";
|
||||
}
|
||||
snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
|
||||
} else if (htp_op_is_unary(node.opcode)) {
|
||||
const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
|
||||
snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
|
||||
} else {
|
||||
snprintf(str, max_size, "----");
|
||||
}
|
||||
|
||||
@@ -39,8 +39,8 @@ add_library(${HTP_LIB} SHARED
|
||||
diag-ops.c
|
||||
solve-tri-ops.c
|
||||
pad-ops.c
|
||||
flash-attn-ops.c
|
||||
matmul-ops.c
|
||||
flash-attn-ops.c
|
||||
)
|
||||
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE
|
||||
|
||||
@@ -120,7 +120,6 @@ int op_concat(struct htp_ops_context * octx);
|
||||
int op_diag(struct htp_ops_context * octx);
|
||||
int op_solve_tri(struct htp_ops_context * octx);
|
||||
int op_gated_delta_net(struct htp_ops_context * octx);
|
||||
int op_tri(struct htp_ops_context * octx);
|
||||
int op_pad(struct htp_ops_context * octx);
|
||||
|
||||
#endif /* HTP_CTX_H */
|
||||
|
||||
@@ -0,0 +1,257 @@
|
||||
#ifndef HVX_NORM_H
|
||||
#define HVX_NORM_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include "hvx-base.h"
|
||||
#include "hvx-reduce.h"
|
||||
#include "hvx-inverse.h"
|
||||
#include "hvx-sqrt.h"
|
||||
#include "hvx-repl.h"
|
||||
|
||||
static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
|
||||
const HVX_Vector * restrict v_src = (HVX_Vector *) src;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32; // number of full vectors
|
||||
const int nloe = num_elems % VLEN_FP32; // leftover elements
|
||||
|
||||
// Compute sum of squares for full vectors
|
||||
HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Reduce HVX sum
|
||||
sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
|
||||
|
||||
HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
|
||||
HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
|
||||
HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
|
||||
HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
|
||||
|
||||
// Scale full vectors
|
||||
HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(v2);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
HVX_Vector result = Q6_Vsf_equals_Vqf32(v2);
|
||||
|
||||
// Store with masking to avoid overwriting memory beyond the tensor
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
|
||||
const uint8_t * restrict weight,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
|
||||
const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32; // number of full vectors
|
||||
const int nloe = num_elems % VLEN_FP32; // leftover elements
|
||||
|
||||
// Compute sum of squares for full vectors
|
||||
HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Reduce HVX sum
|
||||
sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
|
||||
|
||||
HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
|
||||
HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
|
||||
HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
|
||||
HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
|
||||
|
||||
// Scale and multiply
|
||||
HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
|
||||
HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[i]);
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(result);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
|
||||
HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[nvec]);
|
||||
HVX_Vector res_v = Q6_Vsf_equals_Vqf32(result);
|
||||
|
||||
// Store with masking to avoid overwriting memory beyond the tensor
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, res_v);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_fast_norm_f32(const uint8_t * restrict src,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
|
||||
const HVX_Vector * restrict v_src = (HVX_Vector *) src;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32; // number of full vectors
|
||||
const int nloe = num_elems % VLEN_FP32; // leftover elements
|
||||
|
||||
// Compute sum of squares and sum of values for full vectors
|
||||
HVX_Vector sum_sq_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector sum_x_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
|
||||
sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
|
||||
sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
|
||||
}
|
||||
|
||||
// Reduce HVX sums
|
||||
sum_sq_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_sq_v));
|
||||
sum_x_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_x_v));
|
||||
|
||||
HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
|
||||
HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
|
||||
HVX_Vector mean_sq_v = Q6_Vqf32_vmpy_VsfVsf(sum_sq_v, denom_v);
|
||||
HVX_Vector mean_x_v = Q6_Vqf32_vmpy_VsfVsf(sum_x_v, denom_v);
|
||||
HVX_Vector mean_x_sq_v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(mean_x_v), Q6_Vsf_equals_Vqf32(mean_x_v));
|
||||
HVX_Vector var_v = Q6_Vqf32_vsub_Vqf32Vqf32(mean_sq_v, mean_x_sq_v);
|
||||
HVX_Vector var_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(var_v, epsilon_v);
|
||||
|
||||
// scale = rsqrt(variance + epsilon), mean_x broadcast for subtraction
|
||||
HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(var_epsilon_v));
|
||||
HVX_Vector mean_x_b = hvx_vec_repl_f32(Q6_Vsf_equals_Vqf32(mean_x_v));
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
|
||||
HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(v3);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
|
||||
HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
|
||||
HVX_Vector result = Q6_Vsf_equals_Vqf32(v3);
|
||||
|
||||
// Store with masking to avoid overwriting memory beyond the tensor
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_fast_l2_norm_f32(const uint8_t * restrict src,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
|
||||
const HVX_Vector * restrict v_src = (HVX_Vector *) src;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
HVX_Vector sum_v = hvx_vec_splat_f32(0.0f);
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32;
|
||||
const int nloe = num_elems % VLEN_FP32;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
|
||||
}
|
||||
|
||||
// Include tail elements in the sum-of-squares using a predicate mask
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
|
||||
}
|
||||
|
||||
// Compute scale = 1/fmax(sqrt(sum), epsilon) entirely in HVX registers.
|
||||
// hvx_vec_rsqrt_f32 + hvx_vec_inverse_f32 avoids scalar extraction.
|
||||
HVX_Vector sum_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
|
||||
HVX_Vector rsqrt_v = hvx_vec_rsqrt_f32(sum_sf); // 1/sqrt(sum)
|
||||
HVX_Vector sqrt_v = hvx_vec_inverse_f32(rsqrt_v); // sqrt(sum)
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
HVX_Vector denom_v = Q6_Vsf_vmax_VsfVsf(sqrt_v, epsilon_v); // fmax(sqrt(sum), epsilon)
|
||||
HVX_Vector scale_v = hvx_vec_inverse_f32(denom_v); // 1/fmax(sqrt(sum), epsilon)
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
|
||||
}
|
||||
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector result = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // HVX_NORM_H
|
||||
@@ -19,5 +19,6 @@
|
||||
#include "hvx-base.h"
|
||||
#include "hvx-pow.h"
|
||||
#include "hvx-log.h"
|
||||
#include "hvx-norm.h"
|
||||
|
||||
#endif /* HVX_UTILS_H */
|
||||
|
||||
@@ -667,7 +667,7 @@ static int execute_op(struct htp_ops_context * octx) {
|
||||
return op_gated_delta_net(octx);
|
||||
|
||||
case HTP_OP_TRI:
|
||||
return op_tri(octx);
|
||||
return op_unary(octx);
|
||||
|
||||
case HTP_OP_INVALID:
|
||||
break;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,162 @@
|
||||
#ifndef HTP_UNARY_OPS_H
|
||||
#define HTP_UNARY_OPS_H
|
||||
|
||||
#include "hex-common.h"
|
||||
#include "htp-ops.h"
|
||||
|
||||
// Op-specific struct for precomputed unary params
|
||||
struct htp_unary_kernel_params {
|
||||
uint32_t n_threads;
|
||||
uint32_t col_tile;
|
||||
uint32_t vtcm_row_per_thread;
|
||||
uint32_t block;
|
||||
uint32_t broadcast_weight;
|
||||
|
||||
uint32_t vtcm_src0_size_per_thread;
|
||||
uint32_t vtcm_src1_size_per_thread;
|
||||
uint32_t vtcm_dst_size_per_thread;
|
||||
|
||||
uint32_t vtcm_src0_size;
|
||||
uint32_t vtcm_src1_size;
|
||||
uint32_t vtcm_dst_size;
|
||||
|
||||
uint32_t src0_row_size_aligned;
|
||||
uint32_t src1_row_size_aligned;
|
||||
uint32_t dst_row_size_aligned;
|
||||
|
||||
uint32_t vtcm_size;
|
||||
|
||||
// Fastdiv helpers
|
||||
struct fastdiv_values div_ne01;
|
||||
struct fastdiv_values div_ne02;
|
||||
struct fastdiv_values div_ne012;
|
||||
struct fastdiv_values div_tpr;
|
||||
};
|
||||
|
||||
#if defined(__cplusplus)
|
||||
static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
|
||||
#else
|
||||
_Static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
|
||||
#endif
|
||||
|
||||
static inline bool htp_op_is_unary(uint32_t opcode) {
|
||||
switch (opcode) {
|
||||
case HTP_OP_NORM:
|
||||
case HTP_OP_RMS_NORM:
|
||||
case HTP_OP_RMS_NORM_MUL:
|
||||
case HTP_OP_SCALE:
|
||||
case HTP_OP_SQR:
|
||||
case HTP_OP_SQRT:
|
||||
case HTP_OP_UNARY_NEG:
|
||||
case HTP_OP_UNARY_EXP:
|
||||
case HTP_OP_UNARY_SIGMOID:
|
||||
case HTP_OP_UNARY_SOFTPLUS:
|
||||
case HTP_OP_UNARY_TANH:
|
||||
case HTP_OP_L2_NORM:
|
||||
case HTP_OP_TRI:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
struct htp_unary_vtcm_layout {
|
||||
size_t total_bytes;
|
||||
size_t off_src0;
|
||||
size_t off_src1;
|
||||
size_t off_dst;
|
||||
|
||||
size_t src0_bytes;
|
||||
size_t src1_bytes;
|
||||
size_t dst_bytes;
|
||||
};
|
||||
|
||||
static inline void htp_unary_vtcm_layout_build(
|
||||
struct htp_unary_vtcm_layout * L,
|
||||
uint32_t op,
|
||||
uint32_t ne00,
|
||||
uint32_t ne10,
|
||||
uint32_t ne11,
|
||||
bool broadcast_weight,
|
||||
uint32_t n_threads,
|
||||
size_t vtcm_size,
|
||||
uint32_t * out_col_tile,
|
||||
uint32_t * out_vtcm_row_per_thread
|
||||
) {
|
||||
const size_t src0_data_row_size = ne00 * sizeof(float);
|
||||
const size_t dst_data_row_size = ne10 * sizeof(float);
|
||||
|
||||
const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
|
||||
const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
|
||||
|
||||
size_t src1_row_size_aligned = 0;
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
const size_t src1_data_row_size = ne11 * sizeof(float);
|
||||
src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
|
||||
}
|
||||
|
||||
size_t vtcm_size_per_row = 0;
|
||||
size_t vtcm_row_per_thread = 0;
|
||||
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
if (broadcast_weight) {
|
||||
size_t available_vtcm = vtcm_size;
|
||||
size_t src1_vtcm_total = n_threads * src1_row_size_aligned;
|
||||
if (available_vtcm > src1_vtcm_total) {
|
||||
available_vtcm -= src1_vtcm_total;
|
||||
} else {
|
||||
available_vtcm = 0;
|
||||
}
|
||||
vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
|
||||
vtcm_row_per_thread = available_vtcm / (n_threads * vtcm_size_per_row);
|
||||
} else {
|
||||
vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned + src1_row_size_aligned);
|
||||
vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
|
||||
}
|
||||
} else {
|
||||
vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
|
||||
vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
|
||||
}
|
||||
|
||||
const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
|
||||
op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM);
|
||||
uint32_t col_tile = 0;
|
||||
|
||||
if (vtcm_row_per_thread == 0 && !is_reduction) {
|
||||
const size_t per_thread_budget = vtcm_size / n_threads;
|
||||
const size_t col_tile_bytes = hex_align_down(per_thread_budget / 4, 128);
|
||||
col_tile = (uint32_t) (col_tile_bytes / sizeof(float));
|
||||
|
||||
L->src0_bytes = col_tile_bytes * 2;
|
||||
L->dst_bytes = col_tile_bytes * 2;
|
||||
L->src1_bytes = 0;
|
||||
} else {
|
||||
L->src0_bytes = src0_row_size_aligned * vtcm_row_per_thread * 2;
|
||||
L->dst_bytes = dst_row_size_aligned * vtcm_row_per_thread * 2;
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
if (broadcast_weight) {
|
||||
L->src1_bytes = src1_row_size_aligned;
|
||||
} else {
|
||||
L->src1_bytes = src1_row_size_aligned * vtcm_row_per_thread * 2;
|
||||
}
|
||||
} else {
|
||||
L->src1_bytes = 0;
|
||||
}
|
||||
}
|
||||
|
||||
L->off_src0 = 0;
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
L->off_src1 = L->off_src0 + L->src0_bytes * n_threads;
|
||||
L->off_dst = L->off_src1 + L->src1_bytes * n_threads;
|
||||
} else {
|
||||
L->off_src1 = 0;
|
||||
L->off_dst = L->off_src0 + L->src0_bytes * n_threads;
|
||||
}
|
||||
|
||||
L->total_bytes = L->off_dst + L->dst_bytes * n_threads;
|
||||
|
||||
*out_col_tile = col_tile;
|
||||
*out_vtcm_row_per_thread = vtcm_row_per_thread;
|
||||
}
|
||||
|
||||
#endif /* HTP_UNARY_OPS_H */
|
||||
@@ -439,6 +439,19 @@ struct ggml_opencl_fa_kernels {
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_k_img;
|
||||
// k-image variant of MQ_GQA=4 vec_mq_split
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_k_img;
|
||||
// Cluster-parallel decode
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_c8;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c8;
|
||||
// NSG_SPLIT=2 specializations (WG=128): the c8 kernel's register footprint
|
||||
// caps its per-kernel WG at 128 on X2, below the stock 256/192 requirement.
|
||||
// 2 subgroups × FA_CL_NCL streams still gives 16 in-flight rows per WG.
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_c8_ns2;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c8_ns2;
|
||||
// FA_CL_C=32 / MQ_GQA=8 / NSG_SPLIT=2 specialization for the DK=DV=256
|
||||
// GQA=8 class (Qwen3.5/3.6-35B-A3B: 16 Q heads, 2 KV heads). o_acc =
|
||||
// DV_VEC/32 × 8 = 128B/lane (in budget); the baseline fa1 path for this
|
||||
// shape has NO MQ/FD at all and pays an 8× KV re-read per Q head.
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c32;
|
||||
// alternative decode
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_local_tile;
|
||||
// hybrid local-tile + MQ + FD-split kernel for DK=DV=128 only
|
||||
@@ -456,6 +469,8 @@ struct ggml_opencl_fa_kernels {
|
||||
// KV-head-coalesced + flash-decoding split for q8_0 KV
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split_g8;
|
||||
// Cluster-parallel q8_0 decode
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split_c8;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q8_0; // prefill (baseline)
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_split; // N_SPLIT>1 variant
|
||||
std::map<std::pair<int, int>, int> f32_q8_0_split_wg_size; // wg_size = bm*n_split
|
||||
@@ -468,6 +483,9 @@ struct ggml_opencl_fa_kernels {
|
||||
// kv-head-coalesced + flash-decoding split for q4_0 kv (dp4a K dot)
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_g8;
|
||||
// Cluster-parallel q4_0 decode
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_g8_c8;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_c8;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q4_0;
|
||||
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_split;
|
||||
std::map<std::pair<int, int>, int> f32_q4_0_split_wg_size;
|
||||
@@ -4195,11 +4213,20 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac
|
||||
variant == FA_VARIANT_Q4_0_SPLIT;
|
||||
if (is_split) {
|
||||
opts += " -D N_SPLIT=" + std::to_string(cfg->n_split);
|
||||
if (backend_ctx->has_subgroup_shuffle) {
|
||||
opts += backend_ctx->has_qcom_subgroup_shuffle
|
||||
? " -D cl_qcom_subgroup_shuffle=1"
|
||||
: " -D cl_khr_subgroup_shuffle=1";
|
||||
}
|
||||
}
|
||||
// Shuffle define for the split tile paths AND the cluster-parallel decode
|
||||
// kernel (q1_vec_mq_split_c8) in the plain F32_F16 program. Without it the
|
||||
// c8 kernel is compiled out (HAS_SUBGROUP_SHUFFLE guard) and dispatch
|
||||
// falls back to the baseline mq_split.
|
||||
if ((is_split || variant == FA_VARIANT_F32_F16) && backend_ctx->has_subgroup_shuffle) {
|
||||
opts += backend_ctx->has_qcom_subgroup_shuffle
|
||||
? " -D cl_qcom_subgroup_shuffle=1"
|
||||
: " -D cl_khr_subgroup_shuffle=1";
|
||||
}
|
||||
// X1E drops the explicit sub-group size pin on the c8 kernels, compiler
|
||||
// routes the fp16-heavy kernel to a slow variant with explicit subgroup size
|
||||
if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
|
||||
opts += " -D FA_C8_NO_SG_PIN";
|
||||
}
|
||||
return opts;
|
||||
}
|
||||
@@ -4474,6 +4501,21 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
opts += " -D FA_DECODE_ONLY -D FA_DECODE_MINIMAL";
|
||||
}
|
||||
|
||||
// c8 cluster width (GGML_OPENCL_FA_CL_C overrides): value = GQA4 cluster
|
||||
// width (kernel default 8); the g8 programs use 2x the value (default 16).
|
||||
// Wider clusters halve per-lane o_acc at the cost of position streams per
|
||||
// subgroup
|
||||
static const int fa_cl_c_env = []{
|
||||
const char * e = std::getenv("GGML_OPENCL_FA_CL_C");
|
||||
const int x = (e && e[0]) ? atoi(e) : 0;
|
||||
return (x == 8 || x == 16 || x == 32) ? x : 0; // 0 = per-gen default
|
||||
}();
|
||||
const int fa_cl_c_gqa4 = fa_cl_c_env ? fa_cl_c_env
|
||||
: (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 16 : 0);
|
||||
const std::string opts_cl_c_gqa4 = fa_cl_c_gqa4
|
||||
? " -D FA_CL_C=" + std::to_string(fa_cl_c_gqa4) : std::string();
|
||||
const std::string fa_cl_c_g8_val = std::to_string(fa_cl_c_gqa4 ? fa_cl_c_gqa4 * 2 : 16);
|
||||
|
||||
const char * tag = nullptr;
|
||||
switch (variant) {
|
||||
case FA_VARIANT_F16: tag = "fa f16"; break;
|
||||
@@ -4487,7 +4529,7 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
default: break;
|
||||
}
|
||||
cl_program prog = build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts,
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts + opts_cl_c_gqa4,
|
||||
/*fatal=*/false, tag, backend_ctx->queue);
|
||||
if (!prog) { return false; }
|
||||
|
||||
@@ -4570,6 +4612,17 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
clReleaseKernel(k_q1_vec_mq_split_k_img);
|
||||
}
|
||||
}
|
||||
// Cluster-parallel decode variant
|
||||
cl_kernel k_q1_vec_mq_split_c8 = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_c8, 256,
|
||||
"flash_attn_f32_f16_q1_vec_mq_split_c8", dk, dv)) {
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_c8[{dk, dv}] = k_q1_vec_mq_split_c8;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_q1_vec_mq_split_c8);
|
||||
}
|
||||
}
|
||||
cl_kernel k_merge = clCreateKernel(prog, "flash_attn_f32_merge", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
backend_ctx->fa.f32_merge[{dk, dv}] = k_merge;
|
||||
@@ -4602,7 +4655,11 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
// second compile of the same source with -DMQ_GQA=8.
|
||||
// FA_MQ_ONLY keeps only the vec_mq kernels so that the program
|
||||
// compiles within the Adreno compiler's memory budget at DK>=256.
|
||||
const std::string opts_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3 -D FA_MQ_ONLY";
|
||||
// FA_CL_C for the g8 program: MQ_GQA=8 doubles the c8 kernel's
|
||||
// per-lane o_acc, so widen the cluster to keep the register
|
||||
// footprint inside the 192-thread WG cap (see fa_cl_c_gqa4 above
|
||||
// for the per-gen default).
|
||||
const std::string opts_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3 -D FA_MQ_ONLY -D FA_CL_C=" + fa_cl_c_g8_val;
|
||||
cl_program prog_g8 = fa_decode_only ? nullptr : build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8,
|
||||
/*fatal=*/false, "fa f32_f16 MQ_GQA=8", backend_ctx->queue);
|
||||
@@ -4639,6 +4696,17 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
clReleaseKernel(k_q1_vec_mq_split_g8_k_img);
|
||||
}
|
||||
}
|
||||
// Cluster-parallel decode, MQ_GQA=8 / FA_CL_C=16 specialization
|
||||
cl_kernel k_q1_vec_mq_split_g8_c8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_g8_c8, mq_g8_required_wg,
|
||||
"flash_attn_f32_f16_q1_vec_mq_split_c8 (g8)", dk, dv)) {
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8[{dk, dv}] = k_q1_vec_mq_split_g8_c8;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_g8_c8", dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_q1_vec_mq_split_g8_c8);
|
||||
}
|
||||
}
|
||||
// hybrid local-tile + MQ_GQA=8
|
||||
if (dk == 128 && dv == 128) {
|
||||
cl_kernel k_lmq_g8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_local_mq_split", &err);
|
||||
@@ -4654,6 +4722,76 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
}
|
||||
clReleaseProgram(prog_g8);
|
||||
}
|
||||
// NSG_SPLIT=2 programs for the cluster-parallel kernel: its register
|
||||
// footprint caps the per-kernel WG at 128 on X2 (< the stock 256/192
|
||||
// requirement), so it can never register from the stock programs.
|
||||
// With FA_CL_NCL position streams per subgroup, 2 subgroups still
|
||||
// carry 16 in-flight rows per WG (baseline WG has 4). FA_MQ_ONLY
|
||||
// keeps these compiles minimal; skipped when the stock program c8
|
||||
// registered (some other device) or shuffles are absent.
|
||||
if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count({dk, dv}) == 0) {
|
||||
const std::string opts_c8_ns2 = opts + " -D FA_MQ_ONLY -D MQ_GQA=4 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2" + opts_cl_c_gqa4;
|
||||
cl_program prog_c8 = build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8_ns2,
|
||||
/*fatal=*/false, "fa f32_f16 c8 NSG2", backend_ctx->queue);
|
||||
if (prog_c8) {
|
||||
cl_kernel k_c8 = clCreateKernel(prog_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
// WG = MQ_NSG(2) × Q1_WG_SIZE(=FA_SG): 128 Adreno (64), 64 Intel (32).
|
||||
const size_t c8_ns2_wg = backend_ctx->gpu_family == INTEL ? 64 : 128;
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8, c8_ns2_wg,
|
||||
"flash_attn_f32_f16_q1_vec_mq_split_c8 (ns2)", dk, dv)) {
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2[{dk, dv}] = k_c8;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8_ns2", dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_c8);
|
||||
}
|
||||
}
|
||||
clReleaseProgram(prog_c8);
|
||||
}
|
||||
}
|
||||
// FA_CL_C=32 g8 program for the DK=DV=256 GQA=8
|
||||
if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
|
||||
dk == 256 && dv == 256) {
|
||||
const std::string opts_g8_c32 = opts + " -D FA_MQ_ONLY -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2 -D FA_CL_C=32";
|
||||
cl_program prog_g8_c32 = build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8_c32,
|
||||
/*fatal=*/false, "fa f32_f16 c32 g8 d256 NSG2", backend_ctx->queue);
|
||||
if (prog_g8_c32) {
|
||||
cl_kernel k_g8_c32 = clCreateKernel(prog_g8_c32, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8_c32, 128,
|
||||
"flash_attn_f32_f16_q1_vec_mq_split_c8 (g8 c32 d256)", dk, dv)) {
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32[{dk, dv}] = k_g8_c32;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8_c32, "flash_attn_f32_f16_q1_vec_mq_split_g8_c32", dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_g8_c32);
|
||||
}
|
||||
}
|
||||
clReleaseProgram(prog_g8_c32);
|
||||
}
|
||||
}
|
||||
if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count({dk, dv}) == 0) {
|
||||
const std::string opts_g8_c8_ns2 = opts + " -D FA_MQ_ONLY -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2 -D FA_CL_C=" + fa_cl_c_g8_val;
|
||||
cl_program prog_g8_c8 = build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8_c8_ns2,
|
||||
/*fatal=*/false, "fa f32_f16 c8 g8 NSG2", backend_ctx->queue);
|
||||
if (prog_g8_c8) {
|
||||
cl_kernel k_g8_c8 = clCreateKernel(prog_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8_c8, 128,
|
||||
"flash_attn_f32_f16_q1_vec_mq_split_c8 (g8 ns2)", dk, dv)) {
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2[{dk, dv}] = k_g8_c8;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_g8_c8_ns2", dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_g8_c8);
|
||||
}
|
||||
}
|
||||
clReleaseProgram(prog_g8_c8);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case FA_VARIANT_Q8_0:
|
||||
@@ -4735,6 +4873,50 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
|
||||
}
|
||||
clReleaseProgram(prog_mq_g8);
|
||||
}
|
||||
// GQA=4 cluster-parallel program (NSG_SPLIT=2 / WG=128)
|
||||
if (backend_ctx->has_subgroup_shuffle) {
|
||||
auto & m_c8_gqa4 = is_q8 ? backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8
|
||||
: backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8;
|
||||
const std::string name_c8_gqa4 = name_q1 + "_vec_mq_split_c8";
|
||||
const std::string opts_c8_gqa4 = opts + " -D MQ_GQA=4 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2" + opts_cl_c_gqa4;
|
||||
cl_program prog_c8_gqa4 = build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8_gqa4,
|
||||
/*fatal=*/false, is_q8 ? "fa q8_0 c8 GQA4 NSG2" : "fa q4_0 c8 GQA4 NSG2",
|
||||
backend_ctx->queue);
|
||||
if (prog_c8_gqa4) {
|
||||
cl_kernel k_c8_gqa4 = clCreateKernel(prog_c8_gqa4, name_c8_gqa4.c_str(), &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8_gqa4, 128,
|
||||
name_c8_gqa4.c_str(), dk, dv)) {
|
||||
m_c8_gqa4[{dk, dv}] = k_c8_gqa4;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8_gqa4, name_c8_gqa4.c_str(), dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_c8_gqa4);
|
||||
}
|
||||
}
|
||||
clReleaseProgram(prog_c8_gqa4);
|
||||
}
|
||||
}
|
||||
// Cluster-parallel q4_0 decode kernel
|
||||
if (!is_q8 && backend_ctx->has_subgroup_shuffle) {
|
||||
const std::string opts_c8 = opts + " -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2";
|
||||
cl_program prog_c8 = build_program_from_source_ex(
|
||||
backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8,
|
||||
/*fatal=*/false, "fa q4_0 c8 NSG2", backend_ctx->queue);
|
||||
if (prog_c8) {
|
||||
cl_kernel k_c8 = clCreateKernel(prog_c8, "flash_attn_f32_q4_0_q1_vec_mq_split_c8", &err);
|
||||
if (err == CL_SUCCESS) {
|
||||
if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8, 128,
|
||||
"flash_attn_f32_q4_0_q1_vec_mq_split_c8 (g8 ns2)", dk, dv)) {
|
||||
backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8[{dk, dv}] = k_c8;
|
||||
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8, "flash_attn_f32_q4_0_q1_vec_mq_split_g8_c8", dk, dv);
|
||||
} else {
|
||||
clReleaseKernel(k_c8);
|
||||
}
|
||||
}
|
||||
clReleaseProgram(prog_c8);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case FA_VARIANT_F32_F16_SPLIT: {
|
||||
@@ -13794,6 +13976,18 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co
|
||||
|
||||
const bool nq_in_vec_range = (n_q >= 1) && (n_q <= N_MAX_VEC_NQ);
|
||||
const bool nq1_only = (n_q == 1);
|
||||
|
||||
// Cluster-parallel decode default on for Adreno X2E/X1E
|
||||
static const int c8_env_state = []{
|
||||
const char * e = getenv("GGML_OPENCL_FA_C8");
|
||||
if (e == NULL || e[0] == '\0') { return -1; }
|
||||
return (e[0] != '0') ? 1 : 0;
|
||||
}();
|
||||
const bool c8_default_on = backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ||
|
||||
backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E;
|
||||
const bool c8_f16_on = (c8_env_state >= 0) ? (c8_env_state == 1) : c8_default_on;
|
||||
// Quant-KV (q4_0/q8_0) GQA4 c8: default-on X2E + X1E
|
||||
const bool c8_quant_on = (c8_env_state >= 0) ? (c8_env_state == 1) : c8_default_on;
|
||||
if (mq_enabled && mq_kv_ok && nq_in_vec_range && !is_causal &&
|
||||
backend_ctx->gpu_family != INTEL &&
|
||||
!use_local_tile &&
|
||||
@@ -13819,7 +14013,16 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co
|
||||
getenv("GGML_OPENCL_FA_K_IMG") != NULL &&
|
||||
getenv("GGML_OPENCL_FA_K_IMG")[0] != '0' &&
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img.count(dk_dv) > 0;
|
||||
if (k_img_on) {
|
||||
// Cluster-parallel decode
|
||||
const bool c8_env = d_head_q == 128 && d_head_v == 128 && c8_f16_on;
|
||||
if (c8_env && backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
} else if (c8_env && backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
fd_mq_wg = 128;
|
||||
} else if (k_img_on) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
use_fa_k_img = true;
|
||||
@@ -13827,6 +14030,28 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
}
|
||||
// Cluster-parallel decode, DK=DV=256 GQA=8
|
||||
} else if (nq1_only && is_mixed && gqa_ratio_dispatch == 8 &&
|
||||
d_head_q == 256 && d_head_v == 256 &&
|
||||
c8_env_state == 1 &&
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
fd_mq_wg = 128;
|
||||
// Cluster-parallel decode for the g8
|
||||
} else if (is_mixed && gqa_ratio_dispatch == 8 &&
|
||||
d_head_q == 128 && d_head_v == 128 &&
|
||||
c8_f16_on &&
|
||||
(backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count(dk_dv) > 0 ||
|
||||
backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2.count(dk_dv) > 0)) {
|
||||
if (backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.at(dk_dv);
|
||||
fd_mq_wg = 192;
|
||||
} else {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2.at(dk_dv);
|
||||
fd_mq_wg = 128;
|
||||
}
|
||||
use_fd_mq = true;
|
||||
} else if (is_mixed && gqa_ratio_dispatch == 8 &&
|
||||
d_head_q == 128 && d_head_v == 128 &&
|
||||
getenv("GGML_OPENCL_FA_K_IMG") != NULL &&
|
||||
@@ -13851,12 +14076,27 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co
|
||||
} else if (nq1_only && is_q8_0 && gqa_ratio_dispatch == 4 &&
|
||||
d_head_q == 128 && d_head_v == 128 &&
|
||||
backend_ctx->fa.f32_q8_0_q1_vec_mq_split.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split.at(dk_dv);
|
||||
// Cluster-parallel q8_0 GQA4
|
||||
if (c8_quant_on &&
|
||||
backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8.at(dk_dv);
|
||||
fd_mq_wg = 128;
|
||||
} else {
|
||||
fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split.at(dk_dv);
|
||||
}
|
||||
use_fd_mq = true;
|
||||
} else if (nq1_only && is_q4_0) {
|
||||
const char * q4_mq_env = getenv("GGML_OPENCL_FA_Q4_MQ");
|
||||
const bool q4_mq_on = (q4_mq_env != NULL) && (q4_mq_env[0] != '0');
|
||||
if (q4_mq_on && gqa_ratio_dispatch == 8 &&
|
||||
// Cluster-parallel q4_0
|
||||
const bool q4_c8_on = c8_env_state == 1 &&
|
||||
backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8.count(dk_dv) > 0;
|
||||
if (q4_c8_on && gqa_ratio_dispatch == 8 &&
|
||||
d_head_q == 64 && d_head_v == 64) {
|
||||
fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
fd_mq_wg = 128;
|
||||
} else if (q4_mq_on && gqa_ratio_dispatch == 8 &&
|
||||
d_head_q == 128 && d_head_v == 128 &&
|
||||
backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8.at(dk_dv);
|
||||
@@ -13865,12 +14105,35 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co
|
||||
} else if (gqa_ratio_dispatch == 4 &&
|
||||
d_head_q == 128 && d_head_v == 128 &&
|
||||
backend_ctx->fa.f32_q4_0_q1_vec_mq_split.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split.at(dk_dv);
|
||||
// Cluster-parallel q4_0 GQA4
|
||||
if (c8_quant_on &&
|
||||
backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8.at(dk_dv);
|
||||
fd_mq_wg = 128;
|
||||
} else {
|
||||
fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split.at(dk_dv);
|
||||
}
|
||||
use_fd_mq = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Intel cluster-parallel decode FA
|
||||
if (fd_k_split == NULL && backend_ctx->gpu_family == INTEL && n_q == 1 && !is_causal &&
|
||||
is_mixed && gqa_ratio_dispatch == 4 && d_head_q == 128 && d_head_v == 128 &&
|
||||
n_kv >= FD_MIN_N_KV &&
|
||||
getenv("GGML_OPENCL_FA_C8") != NULL && getenv("GGML_OPENCL_FA_C8")[0] != '0' &&
|
||||
backend_ctx->fa.f32_merge.count(dk_dv) > 0) {
|
||||
if (backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
fd_mq_wg = 128;
|
||||
} else if (backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.count(dk_dv) > 0) {
|
||||
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.at(dk_dv);
|
||||
use_fd_mq = true;
|
||||
fd_mq_wg = 64;
|
||||
}
|
||||
}
|
||||
if (fd_k_split == NULL &&
|
||||
n_q >= 1 && n_q <= fd_max_n_q && n_kv >= FD_MIN_N_KV && !is_causal &&
|
||||
d_head_q <= FD_MAX_DK &&
|
||||
|
||||
@@ -1803,6 +1803,345 @@ __kernel void flash_attn_f32_f16_q1_vec_mq_split(
|
||||
}
|
||||
}
|
||||
|
||||
// Cluster-parallel variant of _q1_vec_mq_split
|
||||
//
|
||||
// Tthe baseline keeps one 256B K row in flight per subgroup (32 lanes cooperate
|
||||
// on one position, serialized by the reduce+exp chain). This kernel
|
||||
// takes q1_split's memory-level parallelism at MQ's read-once traffic:
|
||||
// - the 64-lane subgroup is split into FA_CL_NCL clusters of FA_CL_C lanes;
|
||||
// - each cluster owns its own KV position stream (positions strided by
|
||||
// FA_CL_NCL) with private per-cluster online-softmax state, hence FA_CL_NCL
|
||||
// independent K rows in flight per subgroup, no cross-cluster serial chain;
|
||||
// - within a cluster, lanes split DK for the dot (cluster-reduce via
|
||||
// sub_group_shuffle_xor, steps < FA_CL_C stay inside the cluster) and
|
||||
// split DV for o_acc (each lane owns dv indices {lic + FA_CL_C*i} — the
|
||||
// same slice for every position, so accumulation is lane-local);
|
||||
// - merge stage 1 folds the FA_CL_NCL cluster partials with cross-cluster
|
||||
// shuffles (distances >= FA_CL_C); stage 2 is the baseline cross-subgroup
|
||||
// LDS merge (o published by cluster 0's lanes, layout-identical to the
|
||||
// baseline's sg_o).
|
||||
// The KV sweep runs a UNIFORM trip count (max over clusters) with a clamped
|
||||
// row address + FA_M_INIT score on the tail — keeps every shuffle convergent
|
||||
// (p = exp(FA_M_INIT - m) underflows to 0, so clamped-row reads are inert).
|
||||
// Register cost vs baseline: o_acc grows from DV_VEC/64 to DV_VEC/FA_CL_C
|
||||
// float4 per lane per head — FA_CL_C=8 / MQ_GQA=4 => 16 float4 (256B).
|
||||
|
||||
#ifdef HAS_SUBGROUP_SHUFFLE // cluster reduce/merge needs shuffles; absent -> kernel dropped, dispatch falls back
|
||||
|
||||
#ifndef FA_CL_C
|
||||
#define FA_CL_C 8
|
||||
#endif
|
||||
|
||||
// The lane striping requires DK/DV to divide evenly across the cluster;
|
||||
// otherwise (e.g. DK=40 with FA_CL_C=16 -> zero-size arrays) compile the
|
||||
// kernel out — host soft-create falls back silently.
|
||||
#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
|
||||
#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
|
||||
#define FA_CL_DK (DK_VEC / FA_CL_C) // half4s of K per lane per row
|
||||
#define FA_CL_DV (DV_VEC / FA_CL_C) // float4s of o_acc per lane per head
|
||||
|
||||
// explicit "half" sub-group attribute routes this fp16-heavy kernel to a slow
|
||||
// codegen path on the X1 compiler. X2 keeps the pin: its driver miscompile
|
||||
// without it.
|
||||
#ifdef FA_C8_NO_SG_PIN
|
||||
#define FA_C8_SG_ATTR
|
||||
#else
|
||||
// REQD_FA_SG pins the HW subgroup on Intel (intel_reqd_sub_group_size(FA_SG),
|
||||
// host passes -D FA_SG=32); empty on Adreno. REQD_SUBGROUP_SIZE_64 pins 64 on
|
||||
// Adreno; empty on Intel.
|
||||
#define FA_C8_SG_ATTR REQD_FA_SG REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
|
||||
FA_C8_SG_ATTR
|
||||
__kernel void flash_attn_f32_f16_q1_vec_mq_split_c8(
|
||||
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 cl = tid_sg / FA_CL_C; // cluster id
|
||||
const int lic = tid_sg % FA_CL_C; // lane in cluster
|
||||
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;
|
||||
|
||||
// Stage MQ_GQA Q rows in __local once (uniform across WG).
|
||||
__local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
|
||||
for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE) {
|
||||
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;
|
||||
}
|
||||
|
||||
// Per-CLUSTER online-softmax state (uniform across the cluster's lanes);
|
||||
// o_acc holds this lane's DV slice {lic + FA_CL_C*i}.
|
||||
ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DV];
|
||||
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 < FA_CL_DV; ++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);
|
||||
|
||||
// Uniform trip count across the subgroup: every cluster runs n_iter
|
||||
// iterations; tail positions clamp the row address and drop the score to
|
||||
// FA_M_INIT so shuffles stay convergent and the contribution is exactly 0.
|
||||
const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
|
||||
const ulong kv_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
|
||||
const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
|
||||
|
||||
for (int it = 0; it < n_iter; ++it) {
|
||||
const int k_idx = kv_lo + cl + it * FA_CL_NCL;
|
||||
const int valid = k_idx < kv_hi;
|
||||
const int k_safe = valid ? k_idx : (kv_hi - 1);
|
||||
|
||||
const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + kv_row_base + (ulong) k_safe * k_nb1);
|
||||
const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_base + (ulong) k_safe * v_nb1);
|
||||
|
||||
// Dot: this lane covers DK elements {lic + FA_CL_C*i} of the cluster's row.
|
||||
ACC_TYPE4 dot4[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DK; ++i) {
|
||||
const int kk = lic + FA_CL_C * i;
|
||||
const ACC_TYPE4 k_vec = CONVERT_KV_ACC4(k_ptr[kk]);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
dot4[h] = mad(q_shared[h * DK_VEC + kk], k_vec, dot4[h]);
|
||||
}
|
||||
}
|
||||
|
||||
// Cluster-reduce (xor steps < FA_CL_C stay inside the cluster) + score.
|
||||
ACC_TYPE score[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE s = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
|
||||
#pragma unroll
|
||||
for (int step = 1; step < FA_CL_C; step <<= 1) {
|
||||
s += sub_group_shuffle_xor(s, step);
|
||||
}
|
||||
s *= 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_safe];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
s = logit_softcap * tanh(s / logit_softcap);
|
||||
}
|
||||
score[h] = valid ? s : FA_M_INIT;
|
||||
}
|
||||
|
||||
// Per-cluster online update — identical math to the baseline, but the
|
||||
// serial chain is per cluster (depth n_iter, not kv_per_sg).
|
||||
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;
|
||||
}
|
||||
|
||||
// V accumulate on this lane's DV slice (p = 0 on tail -> inert).
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DV; ++i) {
|
||||
const ACC_TYPE4 v_vec = CONVERT_KV_ACC4(v_ptr[lic + FA_CL_C * i]);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
o_acc[h][i] = mad(p_h[h], v_vec, o_acc[h][i] * sp_h[h]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Merge stage 1: fold the FA_CL_NCL cluster partials inside the subgroup.
|
||||
// Lanes with equal lic across clusters hold the SAME dv slice, so a
|
||||
// cross-cluster xor-reduce (distances FA_CL_C..Q1_WG_SIZE/2) sums o
|
||||
// slice-wise; m/l fold the same way. All shuffles are subgroup-convergent.
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE m_c = m_i[h];
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
|
||||
}
|
||||
const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
|
||||
ACC_TYPE l_c = l_i[h] * alpha;
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
l_c += sub_group_shuffle_xor(l_c, step);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DV; ++i) {
|
||||
ACC_TYPE4 o = o_acc[h][i] * alpha;
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
o.s0 += sub_group_shuffle_xor(o.s0, step);
|
||||
o.s1 += sub_group_shuffle_xor(o.s1, step);
|
||||
o.s2 += sub_group_shuffle_xor(o.s2, step);
|
||||
o.s3 += sub_group_shuffle_xor(o.s3, step);
|
||||
}
|
||||
o_acc[h][i] = o;
|
||||
}
|
||||
m_i[h] = m_c;
|
||||
l_i[h] = l_c;
|
||||
}
|
||||
|
||||
// Merge stage 2: baseline cross-subgroup LDS merge. Cluster 0's lanes hold
|
||||
// the subgroup's merged o (dv indices {lic + FA_CL_C*i}) — same sg_o layout
|
||||
// and fold loop as q1_vec_mq_split.
|
||||
__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) {
|
||||
if (cl == 0) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DV; ++i) {
|
||||
sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
|
||||
}
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // DK_VEC/DV_VEC divisible by FA_CL_C
|
||||
#endif // HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
|
||||
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
__kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img(
|
||||
const global void * q_void, ulong q_offset,
|
||||
|
||||
@@ -1157,6 +1157,359 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
|
||||
}
|
||||
}
|
||||
|
||||
// flash_attn_f32_q4_0_q1_vec_mq_split_c8 — cluster-parallel variant of the MQ
|
||||
// split, port of flash_attn_f32_f16_q1_vec_mq_split_c8
|
||||
// Requires dp4a + subgroup shuffles
|
||||
|
||||
#if defined(FA_HAVE_INT_DOT) && defined(HAS_SUBGROUP_SHUFFLE)
|
||||
|
||||
#ifndef FA_CL_C
|
||||
#define FA_CL_C 8
|
||||
#endif
|
||||
|
||||
// Lane striping requires DK/DV to divide across the cluster (see f16 c8).
|
||||
#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
|
||||
#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
|
||||
#define FA_CL_DKQ (DK_VEC / FA_CL_C) // K quartets per lane per row
|
||||
#define FA_CL_DVQ (DV_VEC / FA_CL_C) // V quartets (o_acc float4s) per lane per head
|
||||
|
||||
#ifdef FA_C8_NO_SG_PIN
|
||||
#define FA_C8_SG_ATTR_Q4
|
||||
#else
|
||||
#define FA_C8_SG_ATTR_Q4 REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
|
||||
FA_C8_SG_ATTR_Q4
|
||||
__kernel void flash_attn_f32_q4_0_q1_vec_mq_split_c8(
|
||||
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 cl = tid_sg / FA_CL_C; // cluster id
|
||||
const int lic = tid_sg % FA_CL_C; // lane in cluster
|
||||
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;
|
||||
|
||||
// Stage MQ_GQA Q rows in __local as float4 (source for the quantize pass).
|
||||
__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);
|
||||
|
||||
// Per-(h, block) int8-packed Q + (qd, q_sum), quantized once per WG.
|
||||
__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);
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
// Per-CLUSTER online state; o_acc holds this lane's V quartets {lic + FA_CL_C*i}.
|
||||
ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DVQ];
|
||||
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 < FA_CL_DVQ; ++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);
|
||||
|
||||
// Uniform trip count; tail clamps the row address and drops the score to
|
||||
// FA_M_INIT (p underflows to 0) so shuffles stay convergent.
|
||||
const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
|
||||
const ulong k_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
|
||||
const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
|
||||
|
||||
for (int it = 0; it < n_iter; ++it) {
|
||||
const int k_idx = kv_lo + cl + it * FA_CL_NCL;
|
||||
const int valid = k_idx < kv_hi;
|
||||
const int k_safe = valid ? k_idx : (kv_hi - 1);
|
||||
|
||||
const global char * k_row = k_base + k_row_base + (ulong) k_safe * k_nb1;
|
||||
const global char * v_row = v_base + v_row_base + (ulong) k_safe * v_nb1;
|
||||
|
||||
// dp4a K dot over this lane's quartets of the cluster's row.
|
||||
ACC_TYPE lane_contrib[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DKQ; ++i) {
|
||||
const int qk = lic + FA_CL_C * i;
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// Cluster-reduce + score.
|
||||
ACC_TYPE score[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE s = lane_contrib[h];
|
||||
#pragma unroll
|
||||
for (int step = 1; step < FA_CL_C; step <<= 1) {
|
||||
s += sub_group_shuffle_xor(s, step);
|
||||
}
|
||||
s *= 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_safe];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
s = logit_softcap * tanh(s / logit_softcap);
|
||||
}
|
||||
score[h] = valid ? s : FA_M_INIT;
|
||||
}
|
||||
|
||||
// Per-cluster online update (serial chain depth n_iter, not kv_per_sg).
|
||||
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;
|
||||
}
|
||||
|
||||
// V accumulate on this lane's quartets (p = 0 on tail -> inert).
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DVQ; ++i) {
|
||||
const int dv = lic + FA_CL_C * i;
|
||||
const float4 v_v = dequant_q4_0_lane(v_row + (dv / 8) * Q4_0_BLOCK_SIZE, dv % 8);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
o_acc[h][i] = mad(p_h[h], v_v, o_acc[h][i] * sp_h[h]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Merge stage 1: fold cluster partials inside the subgroup via shuffles.
|
||||
// Lanes with equal lic across clusters hold the SAME dv slice.
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE m_c = m_i[h];
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
|
||||
}
|
||||
const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
|
||||
ACC_TYPE l_c = l_i[h] * alpha;
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
l_c += sub_group_shuffle_xor(l_c, step);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DVQ; ++i) {
|
||||
ACC_TYPE4 o = o_acc[h][i] * alpha;
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
o.s0 += sub_group_shuffle_xor(o.s0, step);
|
||||
o.s1 += sub_group_shuffle_xor(o.s1, step);
|
||||
o.s2 += sub_group_shuffle_xor(o.s2, step);
|
||||
o.s3 += sub_group_shuffle_xor(o.s3, step);
|
||||
}
|
||||
o_acc[h][i] = o;
|
||||
}
|
||||
m_i[h] = m_c;
|
||||
l_i[h] = l_c;
|
||||
}
|
||||
|
||||
// Merge stage 2: baseline cross-subgroup LDS merge (o published by
|
||||
// cluster 0's lanes; layout identical to the baseline sg_o).
|
||||
__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) {
|
||||
if (cl == 0) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DVQ; ++i) {
|
||||
sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
|
||||
}
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // DK_VEC/DV_VEC divisible by FA_CL_C
|
||||
#endif // FA_HAVE_INT_DOT && HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
|
||||
|
||||
__kernel void flash_attn_f32_q4_0(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
|
||||
@@ -970,6 +970,311 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
|
||||
}
|
||||
}
|
||||
|
||||
// flash_attn_f32_q8_0_q1_vec_mq_split_c8 — cluster-parallel variant of the MQ
|
||||
// split above, port of the f16/q4_0 c8 kernels
|
||||
|
||||
#ifdef HAS_SUBGROUP_SHUFFLE
|
||||
|
||||
#ifndef FA_CL_C
|
||||
#define FA_CL_C 8
|
||||
#endif
|
||||
|
||||
// Lane striping requires DK/DV to divide across the cluster (see f16 c8).
|
||||
#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
|
||||
#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
|
||||
#define FA_CL_DKQ (DK_VEC / FA_CL_C) // K quartets per lane per row
|
||||
#define FA_CL_DVQ (DV_VEC / FA_CL_C) // V quartets (o_acc float4s) per lane per head
|
||||
|
||||
#ifdef FA_C8_NO_SG_PIN
|
||||
#define FA_C8_SG_ATTR_Q8
|
||||
#else
|
||||
#define FA_C8_SG_ATTR_Q8 REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
|
||||
FA_C8_SG_ATTR_Q8
|
||||
__kernel void flash_attn_f32_q8_0_q1_vec_mq_split_c8(
|
||||
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 cl = tid_sg / FA_CL_C; // cluster id
|
||||
const int lic = tid_sg % FA_CL_C; // lane in cluster
|
||||
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;
|
||||
|
||||
// Stage MQ_GQA Q rows in __local once (uniform across WG).
|
||||
__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;
|
||||
}
|
||||
|
||||
// Per-CLUSTER online state; o_acc holds this lane's V quartets {lic + FA_CL_C*i}.
|
||||
ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DVQ];
|
||||
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 < FA_CL_DVQ; ++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);
|
||||
|
||||
// Uniform trip count; tail clamps the row address and drops the score to
|
||||
// FA_M_INIT (p underflows to 0) so shuffles stay convergent.
|
||||
const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
|
||||
const ulong k_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
|
||||
const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
|
||||
|
||||
for (int it = 0; it < n_iter; ++it) {
|
||||
const int k_idx = kv_lo + cl + it * FA_CL_NCL;
|
||||
const int valid = k_idx < kv_hi;
|
||||
const int k_safe = valid ? k_idx : (kv_hi - 1);
|
||||
|
||||
const global char * k_row = k_base + k_row_base + (ulong) k_safe * k_nb1;
|
||||
const global char * v_row = v_base + v_row_base + (ulong) k_safe * v_nb1;
|
||||
|
||||
// Float-dequant K dot over this lane's quartets of the cluster's row.
|
||||
ACC_TYPE4 dot4[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DKQ; ++i) {
|
||||
const int qk = lic + FA_CL_C * i;
|
||||
const float4 k_v = dequant_q8_0_lane(k_row + (qk / 8) * Q8_0_BLOCK_SIZE, qk % 8);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
|
||||
}
|
||||
}
|
||||
|
||||
// Cluster-reduce (xor steps < FA_CL_C stay inside the cluster) + score.
|
||||
ACC_TYPE score[MQ_GQA];
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE s = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
|
||||
#pragma unroll
|
||||
for (int step = 1; step < FA_CL_C; step <<= 1) {
|
||||
s += sub_group_shuffle_xor(s, step);
|
||||
}
|
||||
s *= 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_safe];
|
||||
}
|
||||
if (logit_softcap > 0.0f) {
|
||||
s = logit_softcap * tanh(s / logit_softcap);
|
||||
}
|
||||
score[h] = valid ? s : FA_M_INIT;
|
||||
}
|
||||
|
||||
// Per-cluster online update (serial chain depth n_iter, not kv_per_sg).
|
||||
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;
|
||||
}
|
||||
|
||||
// V accumulate on this lane's quartets (p = 0 on tail -> inert).
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DVQ; ++i) {
|
||||
const int dv = lic + FA_CL_C * i;
|
||||
const float4 v_v = dequant_q8_0_lane(v_row + (dv / 8) * Q8_0_BLOCK_SIZE, dv % 8);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
o_acc[h][i] = mad(p_h[h], v_v, o_acc[h][i] * sp_h[h]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Merge stage 1: fold cluster partials inside the subgroup via shuffles.
|
||||
#pragma unroll
|
||||
for (int h = 0; h < MQ_GQA; ++h) {
|
||||
ACC_TYPE m_c = m_i[h];
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
|
||||
}
|
||||
const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
|
||||
ACC_TYPE l_c = l_i[h] * alpha;
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
l_c += sub_group_shuffle_xor(l_c, step);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DVQ; ++i) {
|
||||
ACC_TYPE4 o = o_acc[h][i] * alpha;
|
||||
#pragma unroll
|
||||
for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
|
||||
o.s0 += sub_group_shuffle_xor(o.s0, step);
|
||||
o.s1 += sub_group_shuffle_xor(o.s1, step);
|
||||
o.s2 += sub_group_shuffle_xor(o.s2, step);
|
||||
o.s3 += sub_group_shuffle_xor(o.s3, step);
|
||||
}
|
||||
o_acc[h][i] = o;
|
||||
}
|
||||
m_i[h] = m_c;
|
||||
l_i[h] = l_c;
|
||||
}
|
||||
|
||||
// Merge stage 2: baseline cross-subgroup LDS merge (o published by
|
||||
// cluster 0's lanes; layout identical to the baseline sg_o).
|
||||
__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) {
|
||||
if (cl == 0) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FA_CL_DVQ; ++i) {
|
||||
sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
|
||||
}
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // DK_VEC/DV_VEC divisible by FA_CL_C
|
||||
#endif // HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
|
||||
|
||||
__kernel void flash_attn_f32_q8_0(
|
||||
const global void * q_void, ulong q_offset,
|
||||
const global void * k_void, ulong k_offset,
|
||||
|
||||
@@ -162,6 +162,14 @@ bool cli_context::init() {
|
||||
|
||||
fetch_server_props();
|
||||
|
||||
if (!params.out_file.empty()) {
|
||||
output_file.emplace(params.out_file);
|
||||
if (!output_file->is_open()) {
|
||||
ui::show_error(string_format("failed to open output file '%s'", params.out_file.c_str()));
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -323,7 +331,14 @@ bool cli_context::stage_media_file(const std::string & fname, const std::string
|
||||
return true;
|
||||
}
|
||||
|
||||
bool cli_context::generate_completion(std::string & assistant_content, cli_timings & timings) {
|
||||
void cli_context::write_output_file(const std::string & content) {
|
||||
if (output_file) {
|
||||
(*output_file) << content;
|
||||
output_file->flush();
|
||||
}
|
||||
}
|
||||
|
||||
bool cli_context::generate_completion(generated_content & content_out, cli_timings & timings) {
|
||||
json body = {
|
||||
{"messages", impl->messages},
|
||||
{"stream", true},
|
||||
@@ -364,13 +379,14 @@ bool cli_context::generate_completion(std::string & assistant_content, cli_timin
|
||||
if (delta.contains("reasoning_content") && delta.at("reasoning_content").is_string()) {
|
||||
const std::string text = delta.at("reasoning_content").get<std::string>();
|
||||
if (!text.empty()) {
|
||||
content_out.reasoning += text;
|
||||
a.push(ui::ASSISTANT_DISPLAY_MODE_REASONING, text);
|
||||
}
|
||||
}
|
||||
if (delta.contains("content") && delta.at("content").is_string()) {
|
||||
const std::string text = delta.at("content").get<std::string>();
|
||||
if (!text.empty()) {
|
||||
assistant_content += text;
|
||||
content_out.content += text;
|
||||
a.push(ui::ASSISTANT_DISPLAY_MODE_CONTENT, text);
|
||||
}
|
||||
}
|
||||
@@ -520,10 +536,12 @@ int cli_context::run() {
|
||||
continue;
|
||||
}
|
||||
ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
|
||||
write_output_file(string_format("User: Added media: %s\n", fname.c_str()));
|
||||
continue;
|
||||
} else if (string_starts_with(buffer, "/read ")) {
|
||||
std::string fname = string_strip(buffer.substr(6));
|
||||
add_text_file(fname);
|
||||
write_output_file(string_format("User: Added text file: %s\n", fname.c_str()));
|
||||
continue;
|
||||
} else if (string_starts_with(buffer, "/glob ")) {
|
||||
std::error_code ec;
|
||||
@@ -568,9 +586,11 @@ int cli_context::run() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!add_text_file((rel_path / rel).string())) {
|
||||
const std::string full_path = (curdir / rel).string();
|
||||
if (!add_text_file(full_path)) {
|
||||
continue;
|
||||
}
|
||||
write_output_file(string_format("User: Added text file: %s\n", full_path.c_str()));
|
||||
|
||||
if (++count >= FILE_GLOB_MAX_RESULTS) {
|
||||
ui::show_error(string_format("Maximum number of globbed files allowed (%zu) reached.", FILE_GLOB_MAX_RESULTS));
|
||||
@@ -586,16 +606,34 @@ int cli_context::run() {
|
||||
// generate response
|
||||
if (add_user_msg) {
|
||||
push_user_message(cur_msg);
|
||||
write_output_file(string_format("User:\n%s\n\n", cur_msg.c_str()));
|
||||
cur_msg.clear();
|
||||
}
|
||||
|
||||
cli_timings timings;
|
||||
std::string assistant_content;
|
||||
generate_completion(assistant_content, timings);
|
||||
generated_content content;
|
||||
generate_completion(content, timings);
|
||||
|
||||
impl->messages.push_back({
|
||||
{"role", "assistant"},
|
||||
{"content", assistant_content}
|
||||
{"content", content.content}
|
||||
});
|
||||
|
||||
if (output_file) {
|
||||
std::string out_content = "Assistant:\n";
|
||||
if (!content.reasoning.empty()) {
|
||||
out_content += "[Start thinking]\n\n";
|
||||
out_content += content.reasoning;
|
||||
out_content += "[End thinking]\n\n";
|
||||
}
|
||||
out_content += content.content;
|
||||
if (!out_content.empty() && out_content.back() != '\n') {
|
||||
out_content += "\n";
|
||||
}
|
||||
out_content += "\n";
|
||||
write_output_file(out_content);
|
||||
}
|
||||
|
||||
if (params.show_timings) {
|
||||
ui::show_info(string_format(
|
||||
"\n[ Prompt: %.1f t/s | Generation: %.1f t/s ]",
|
||||
@@ -619,4 +657,8 @@ void cli_context::shutdown() {
|
||||
server->stop();
|
||||
server.reset();
|
||||
}
|
||||
if (output_file) {
|
||||
output_file->close();
|
||||
output_file.reset();
|
||||
}
|
||||
}
|
||||
|
||||
+11
-1
@@ -9,6 +9,7 @@
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <fstream>
|
||||
|
||||
struct cli_timings {
|
||||
double prompt_per_second = 0.0;
|
||||
@@ -32,6 +33,8 @@ struct cli_context {
|
||||
bool has_audio = false;
|
||||
bool has_video = false;
|
||||
|
||||
std::optional<std::ofstream> output_file;
|
||||
|
||||
cli_context(const common_params & params);
|
||||
~cli_context();
|
||||
|
||||
@@ -49,7 +52,11 @@ struct cli_context {
|
||||
static std::atomic<bool> & interrupted();
|
||||
|
||||
private:
|
||||
bool generate_completion(std::string & assistant_content, cli_timings & timings);
|
||||
struct generated_content {
|
||||
std::string reasoning;
|
||||
std::string content;
|
||||
};
|
||||
bool generate_completion(generated_content & content_out, cli_timings & timings);
|
||||
void fetch_server_props();
|
||||
void add_system_prompt();
|
||||
void push_user_message(const std::string & text);
|
||||
@@ -62,5 +69,8 @@ private:
|
||||
// "image", "audio", "video"; returns false if the file cannot be read
|
||||
bool stage_media_file(const std::string & fname, const std::string & type);
|
||||
|
||||
// no-op if output file is not set
|
||||
void write_output_file(const std::string & content);
|
||||
|
||||
std::unique_ptr<cli_context_impl> impl;
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user