mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-13 16:05:54 +02:00
Compare commits
20 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e474bba7af | |||
| 38fd5c9993 | |||
| 99f3dc3229 | |||
| 34558825a2 | |||
| 8014d2cf97 | |||
| 4114ba18b2 | |||
| 0c4fa7a989 | |||
| 6b4dc2116a | |||
| e3546c7948 | |||
| d72bfa38f7 | |||
| 3cec3bcd16 | |||
| 13f2b28b09 | |||
| c92e806d1c | |||
| ea1f7bbb5d | |||
| 00f5442cc4 | |||
| 76f2798059 | |||
| 1d1d9a9ed7 | |||
| 4f37f51972 | |||
| c749cb0417 | |||
| 67776eaee5 |
@@ -1081,6 +1081,9 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *);
|
||||
struct common_prompt_checkpoint {
|
||||
int64_t n_tokens;
|
||||
|
||||
// (optional) id of the task that created the checkpoint
|
||||
int id_task = -1;
|
||||
|
||||
llama_pos pos_min;
|
||||
llama_pos pos_max;
|
||||
|
||||
|
||||
@@ -8,10 +8,10 @@ extern "C" {
|
||||
|
||||
#define RPC_PROTO_MAJOR_VERSION 4
|
||||
#define RPC_PROTO_MINOR_VERSION 0
|
||||
#define RPC_PROTO_PATCH_VERSION 1
|
||||
#define RPC_PROTO_PATCH_VERSION 2
|
||||
|
||||
#ifdef __cplusplus
|
||||
static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
|
||||
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
|
||||
#endif
|
||||
|
||||
#define GGML_RPC_MAX_SERVERS 16
|
||||
|
||||
@@ -570,6 +570,7 @@ extern "C" {
|
||||
GGML_OP_RWKV_WKV7,
|
||||
GGML_OP_SOLVE_TRI,
|
||||
GGML_OP_GATED_DELTA_NET,
|
||||
GGML_OP_LIGHTNING_INDEXER,
|
||||
|
||||
GGML_OP_UNARY,
|
||||
|
||||
@@ -2575,6 +2576,24 @@ extern "C" {
|
||||
struct ggml_tensor * state,
|
||||
int64_t K);
|
||||
|
||||
// DSA lightning indexer
|
||||
//
|
||||
// q: [n_embd_idx, n_head_idx, n_batch, ne3 ]
|
||||
// k: [n_embd_idx, 1, n_kv, ne3 ]
|
||||
// weights: [n_head_idx, n_batch, 1, ne3 ] !! prescaled !!
|
||||
// mask: [n_kv, n_batch, 1, ne33] !! f16 !!
|
||||
// res: [n_kv, n_batch, 1, ne3 ]
|
||||
//
|
||||
// broadcast:
|
||||
// ne3 % ne33 == 0
|
||||
//
|
||||
GGML_API struct ggml_tensor * ggml_lightning_indexer(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * q,
|
||||
struct ggml_tensor * k,
|
||||
struct ggml_tensor * weights,
|
||||
struct ggml_tensor * mask);
|
||||
|
||||
// custom operators
|
||||
|
||||
typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);
|
||||
|
||||
@@ -2060,6 +2060,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
|
||||
{
|
||||
ggml_compute_forward_gated_delta_net(params, tensor);
|
||||
} break;
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
{
|
||||
ggml_compute_forward_lightning_indexer(params, tensor);
|
||||
} break;
|
||||
case GGML_OP_MAP_CUSTOM1:
|
||||
{
|
||||
ggml_compute_forward_map_custom1(params, tensor);
|
||||
@@ -2380,6 +2384,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
|
||||
case GGML_OP_FLASH_ATTN_BACK:
|
||||
case GGML_OP_SSM_CONV:
|
||||
case GGML_OP_SSM_SCAN:
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
{
|
||||
n_tasks = n_threads;
|
||||
} break;
|
||||
@@ -2965,6 +2970,12 @@ struct ggml_cplan ggml_graph_plan(
|
||||
{
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
{
|
||||
// temp buffer for dequantizing lightning indexer keys
|
||||
const int64_t ne10 = node->src[1]->ne[0];
|
||||
cur += sizeof(float)*ne10*n_tasks;
|
||||
} break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -11568,3 +11568,87 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor *
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ggml_compute_forward_lightning_indexer
|
||||
|
||||
void ggml_compute_forward_lightning_indexer(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
|
||||
const ggml_tensor * q = dst->src[0];
|
||||
const ggml_tensor * k = dst->src[1];
|
||||
const ggml_tensor * w = dst->src[2]; // weights
|
||||
const ggml_tensor * m = dst->src[3]; // mask
|
||||
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( q->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( w->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( m->type == GGML_TYPE_F16);
|
||||
|
||||
GGML_TENSOR_LOCALS(int64_t, neq, q, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbq, q, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, nek, k, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbk, k, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, new, w, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbw, w, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, nem, m, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbm, m, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nb, dst, nb)
|
||||
|
||||
GGML_ASSERT( nb0 == ggml_type_size(dst->type));
|
||||
GGML_ASSERT(nbq0 == ggml_type_size( q->type));
|
||||
GGML_ASSERT(nbk0 == ggml_type_size( k->type));
|
||||
GGML_ASSERT(nbw0 == ggml_type_size( w->type));
|
||||
GGML_ASSERT(nbm0 == ggml_type_size( m->type));
|
||||
|
||||
const int n_embd = q->ne[0];
|
||||
const int n_head = q->ne[1];
|
||||
const int n_tokens = q->ne[2];
|
||||
const int n_stream = q->ne[3];
|
||||
const int n_kv = k->ne[2];
|
||||
|
||||
ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float;
|
||||
GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type");
|
||||
|
||||
const int nr = n_kv;
|
||||
const int ith = params->ith;
|
||||
const int nth = params->nth;
|
||||
|
||||
// (temporary) buffer for K converted to float
|
||||
float * k_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32);
|
||||
|
||||
// rows per thread
|
||||
const int dr = (nr + nth - 1)/nth;
|
||||
|
||||
// row range for this thread
|
||||
const int ir0 = dr*ith;
|
||||
const int ir1 = MIN(ir0 + dr, nr);
|
||||
|
||||
for (int s = 0; s < n_stream; ++s) {
|
||||
for (int t = 0; t < n_tokens; ++t) {
|
||||
const float * w_row = (float *) ((char *) w->data + t*nbw1 + s*nbw3);
|
||||
const ggml_fp16_t * m_row = (ggml_fp16_t *) ((char *) m->data + t*nbm1 + (s%nem3)*nbm3);
|
||||
float * dst_row = (float *) ((char *) dst->data + t*nb1 + s*nb3 );
|
||||
for (int ik = ir0; ik < ir1; ++ik) {
|
||||
char * k_row = (char *) k->data + ik*nbk2 + s*nbk3;
|
||||
if (k_to_float) {
|
||||
k_to_float(k_row, k_row_f32, n_embd);
|
||||
} else {
|
||||
k_row_f32 = (float *) k_row;
|
||||
}
|
||||
float score = 0.0f;
|
||||
for (int h = 0; h < n_head; ++h) {
|
||||
// dot product of q and k for head h
|
||||
float qk = 0.0f;
|
||||
const float * q_row = (float *) ((char *) q->data + h*nbq1 + t*nbq2 + s*nbq3);
|
||||
ggml_vec_dot_f32(n_embd, &qk, 0, q_row, 0, k_row_f32, 0, 1);
|
||||
// ReLU and weights (prescaled)
|
||||
score += MAX(qk, 0.0f) * w_row[h];
|
||||
}
|
||||
// apply mask
|
||||
dst_row[ik] = score + GGML_CPU_FP16_TO_FP32(m_row[ik]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,6 +105,7 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s
|
||||
void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
|
||||
@@ -4493,7 +4493,14 @@ static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_k
|
||||
static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
|
||||
ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context;
|
||||
ggml_cuda_set_device(ctx->device);
|
||||
CUDA_CHECK(cudaMemGetInfo(free, total));
|
||||
cudaError_t err = cudaMemGetInfo(free, total);
|
||||
if (err != cudaSuccess) {
|
||||
(void)cudaGetLastError();
|
||||
GGML_LOG_WARN("%s: cudaMemGetInfo failed (%s), returning 0/0\n", __func__, cudaGetErrorString(err));
|
||||
*free = 0;
|
||||
*total = 0;
|
||||
return;
|
||||
}
|
||||
|
||||
// ref: https://github.com/ggml-org/llama.cpp/pull/17368
|
||||
#if defined(__linux__)
|
||||
|
||||
@@ -31,7 +31,6 @@ add_library(${HTP_LIB} SHARED
|
||||
get-rows-ops.c
|
||||
cpy-ops.c
|
||||
repeat-ops.c
|
||||
argsort-ops.c
|
||||
ssm-conv.c
|
||||
cumsum-ops.c
|
||||
fill-ops.c
|
||||
@@ -39,8 +38,9 @@ add_library(${HTP_LIB} SHARED
|
||||
diag-ops.c
|
||||
solve-tri-ops.c
|
||||
pad-ops.c
|
||||
matmul-ops.c
|
||||
flash-attn-ops.c
|
||||
matmul-ops.c
|
||||
argsort-ops.c
|
||||
)
|
||||
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE
|
||||
|
||||
@@ -22,6 +22,8 @@
|
||||
struct htp_argsort_context {
|
||||
struct htp_ops_context * octx;
|
||||
uint32_t nrows_per_thread;
|
||||
uint8_t * vtcm_base;
|
||||
size_t vtcm_per_thread;
|
||||
};
|
||||
|
||||
static inline bool all_greater_f32(HVX_Vector x, HVX_Vector y)
|
||||
@@ -170,7 +172,208 @@ int32_t argosrt_ramp_lut[32] __attribute__((aligned(VLEN))) = {
|
||||
16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
|
||||
};
|
||||
|
||||
static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
|
||||
__attribute__((always_inline))
|
||||
static inline void vec_cas(HVX_Vector * X_val, HVX_Vector * X_idx, HVX_Vector * Y_val, HVX_Vector * Y_idx, bool asc) {
|
||||
HVX_VectorPred pred = asc ? Q6_Q_vcmp_gt_VsfVsf(*X_val, *Y_val)
|
||||
: Q6_Q_vcmp_gt_VsfVsf(*Y_val, *X_val);
|
||||
HVX_Vector next_X_val = Q6_V_vmux_QVV(pred, *Y_val, *X_val);
|
||||
HVX_Vector next_Y_val = Q6_V_vmux_QVV(pred, *X_val, *Y_val);
|
||||
HVX_Vector next_X_idx = Q6_V_vmux_QVV(pred, *Y_idx, *X_idx);
|
||||
HVX_Vector Y_tmp_idx = Q6_V_vmux_QVV(pred, *X_idx, *Y_idx);
|
||||
*X_val = next_X_val;
|
||||
*Y_val = next_Y_val;
|
||||
*X_idx = next_X_idx;
|
||||
*Y_idx = Y_tmp_idx;
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void bitonic_cas_32(HVX_Vector * V, HVX_Vector * I, int d, HVX_VectorPred dir_mask, HVX_Vector idx_vec, HVX_Vector zero_vec) {
|
||||
HVX_VectorPred mask_left;
|
||||
HVX_Vector V_rot_left, V_rot_right;
|
||||
HVX_Vector I_rot_left, I_rot_right;
|
||||
|
||||
if (d == 1) {
|
||||
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1)), zero_vec);
|
||||
V_rot_left = Q6_V_vror_VR(*V, 4);
|
||||
V_rot_right = Q6_V_vror_VR(*V, 124);
|
||||
I_rot_left = Q6_V_vror_VR(*I, 4);
|
||||
I_rot_right = Q6_V_vror_VR(*I, 124);
|
||||
} else if (d == 2) {
|
||||
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(2)), zero_vec);
|
||||
V_rot_left = Q6_V_vror_VR(*V, 8);
|
||||
V_rot_right = Q6_V_vror_VR(*V, 120);
|
||||
I_rot_left = Q6_V_vror_VR(*I, 8);
|
||||
I_rot_right = Q6_V_vror_VR(*I, 120);
|
||||
} else if (d == 4) {
|
||||
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(4)), zero_vec);
|
||||
V_rot_left = Q6_V_vror_VR(*V, 16);
|
||||
V_rot_right = Q6_V_vror_VR(*V, 112);
|
||||
I_rot_left = Q6_V_vror_VR(*I, 16);
|
||||
I_rot_right = Q6_V_vror_VR(*I, 112);
|
||||
} else if (d == 8) {
|
||||
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(8)), zero_vec);
|
||||
V_rot_left = Q6_V_vror_VR(*V, 32);
|
||||
V_rot_right = Q6_V_vror_VR(*V, 96);
|
||||
I_rot_left = Q6_V_vror_VR(*I, 32);
|
||||
I_rot_right = Q6_V_vror_VR(*I, 96);
|
||||
} else { // d == 16
|
||||
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(16)), zero_vec);
|
||||
V_rot_left = Q6_V_vror_VR(*V, 64);
|
||||
V_rot_right = Q6_V_vror_VR(*V, 64);
|
||||
I_rot_left = Q6_V_vror_VR(*I, 64);
|
||||
I_rot_right = Q6_V_vror_VR(*I, 64);
|
||||
}
|
||||
|
||||
HVX_Vector V_paired = Q6_V_vmux_QVV(mask_left, V_rot_left, V_rot_right);
|
||||
HVX_Vector I_paired = Q6_V_vmux_QVV(mask_left, I_rot_left, I_rot_right);
|
||||
|
||||
HVX_VectorPred V_gt_Vpaired = Q6_Q_vcmp_gt_VsfVsf(*V, V_paired);
|
||||
HVX_VectorPred Vpaired_gt_V = Q6_Q_vcmp_gt_VsfVsf(V_paired, *V);
|
||||
HVX_VectorPred mask_right = Q6_Q_not_Q(mask_left);
|
||||
HVX_VectorPred Q_asc = Q6_Q_or_QQ(
|
||||
Q6_Q_and_QQ(mask_left, V_gt_Vpaired),
|
||||
Q6_Q_and_QQ(Vpaired_gt_V, mask_right)
|
||||
);
|
||||
HVX_VectorPred Q_swap = Q6_Q_or_QQ(
|
||||
Q6_Q_and_QQ(dir_mask, Q_asc),
|
||||
Q6_Q_and_QQ(Q6_Q_not_Q(dir_mask), Q6_Q_not_Q(Q_asc))
|
||||
);
|
||||
|
||||
*V = Q6_V_vmux_QVV(Q_swap, V_paired, *V);
|
||||
*I = Q6_V_vmux_QVV(Q_swap, I_paired, *I);
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void bitonic_sort_generic_hvx(uint8_t * values, uint8_t * indices, int K, bool asc_order) {
|
||||
HVX_Vector V[32];
|
||||
HVX_Vector I[32];
|
||||
|
||||
HVX_Vector zero_vec = Q6_V_vzero();
|
||||
HVX_Vector idx_vec = *(HVX_Vector *)argosrt_ramp_lut;
|
||||
|
||||
// Load values and initialize indices
|
||||
for (int v = 0; v < K; v++) {
|
||||
V[v] = *(HVX_Vector *)(values + v * 128);
|
||||
I[v] = Q6_Vw_vadd_VwVw(idx_vec, Q6_V_vsplat_R(v * 32));
|
||||
}
|
||||
|
||||
HVX_VectorPred pred_all_1s = Q6_Q_vcmp_eq_VwVw(zero_vec, zero_vec);
|
||||
HVX_VectorPred pred_all_0s = Q6_Q_not_Q(pred_all_1s);
|
||||
|
||||
int M = 5;
|
||||
while ((1 << (M - 5)) < K) M++;
|
||||
|
||||
for (int s = 1; s <= M; s++) {
|
||||
for (int stage_d = s - 1; stage_d >= 0; stage_d--) {
|
||||
int d = 1 << stage_d;
|
||||
if (d >= 32) {
|
||||
int v_dist = d / 32;
|
||||
for (int v1 = 0; v1 < K; v1++) {
|
||||
if ((v1 & v_dist) == 0) {
|
||||
int v2 = v1 + v_dist;
|
||||
bool asc = (s < M) ? ((((v1 * 32) >> s) % 2) == 0) : asc_order;
|
||||
vec_cas(&V[v1], &I[v1], &V[v2], &I[v2], asc);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (s < 5) {
|
||||
HVX_VectorPred dir_mask = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1 << s)), zero_vec);
|
||||
for (int v = 0; v < K; v++) {
|
||||
bitonic_cas_32(&V[v], &I[v], d, dir_mask, idx_vec, zero_vec);
|
||||
}
|
||||
} else {
|
||||
for (int v = 0; v < K; v++) {
|
||||
bool asc = (s < M) ? ((((v * 32) >> s) % 2) == 0) : asc_order;
|
||||
HVX_VectorPred dir_mask = asc ? pred_all_1s : pred_all_0s;
|
||||
bitonic_cas_32(&V[v], &I[v], d, dir_mask, idx_vec, zero_vec);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Write back sorted values and indices
|
||||
for (int v = 0; v < K; v++) {
|
||||
*(HVX_Vector *)(values + v * 128) = V[v];
|
||||
*(HVX_Vector *)(indices + v * 128) = I[v];
|
||||
}
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void sort32_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
|
||||
bitonic_sort_generic_hvx(values, indices, 1, order == GGML_SORT_ORDER_ASC);
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void sort64_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
|
||||
bitonic_sort_generic_hvx(values, indices, 2, order == GGML_SORT_ORDER_ASC);
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void sort128_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
|
||||
bitonic_sort_generic_hvx(values, indices, 4, order == GGML_SORT_ORDER_ASC);
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void sort256_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
|
||||
bitonic_sort_generic_hvx(values, indices, 8, order == GGML_SORT_ORDER_ASC);
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void sort512_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
|
||||
bitonic_sort_generic_hvx(values, indices, 16, order == GGML_SORT_ORDER_ASC);
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
static inline void sort1024_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
|
||||
bitonic_sort_generic_hvx(values, indices, 32, order == GGML_SORT_ORDER_ASC);
|
||||
}
|
||||
|
||||
#define HTP_ARGSORT_FN(ne00, order_name, order_enum, sort_fn) \
|
||||
static void htp_argsort_f32_##ne00##_##order_name(unsigned int n, unsigned int i, void * data) { \
|
||||
struct htp_argsort_context * actx = (struct htp_argsort_context *)data; \
|
||||
struct htp_ops_context * octx = actx->octx; \
|
||||
const struct htp_tensor * src0 = octx->src[0]; \
|
||||
const struct htp_tensor * dst = octx->dst; \
|
||||
uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i; \
|
||||
uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3]; \
|
||||
uint32_t rows_per_thread = actx->nrows_per_thread; \
|
||||
uint32_t start_row = rows_per_thread * i; \
|
||||
uint32_t end_row = MIN(start_row + rows_per_thread, total_rows); \
|
||||
size_t values_size = hex_round_up(ne00 * sizeof(float), 128); \
|
||||
float * values_buf = (float *) spad; \
|
||||
int32_t * indices_buf = (int32_t *) (spad + values_size); \
|
||||
uint32_t nb01 = src0->nb[1]; \
|
||||
uint32_t nb1 = dst->nb[1]; \
|
||||
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL; \
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row); \
|
||||
for (uint32_t r = start_row; r < end_row; r++) { \
|
||||
uint32_t src_offset = r * nb01; \
|
||||
uint32_t dst_offset = r * nb1; \
|
||||
uint8_t * src_ptr = (uint8_t *) src0->data + src_offset; \
|
||||
uint8_t * dst_ptr = (uint8_t *) dst->data + dst_offset; \
|
||||
hex_l2fetch(src_ptr, ne00 * sizeof(float), ne00 * sizeof(float), 1); \
|
||||
hvx_copy_f32_au((uint8_t*)values_buf, src_ptr, ne00); \
|
||||
sort_fn((uint8_t*)values_buf, (uint8_t*)indices_buf, order_enum); \
|
||||
hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00); \
|
||||
} \
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row); \
|
||||
}
|
||||
|
||||
HTP_ARGSORT_FN(32, asc, GGML_SORT_ORDER_ASC, sort32_f32_hvx)
|
||||
HTP_ARGSORT_FN(32, dsc, GGML_SORT_ORDER_DESC, sort32_f32_hvx)
|
||||
HTP_ARGSORT_FN(64, asc, GGML_SORT_ORDER_ASC, sort64_f32_hvx)
|
||||
HTP_ARGSORT_FN(64, dsc, GGML_SORT_ORDER_DESC, sort64_f32_hvx)
|
||||
HTP_ARGSORT_FN(128, asc, GGML_SORT_ORDER_ASC, sort128_f32_hvx)
|
||||
HTP_ARGSORT_FN(128, dsc, GGML_SORT_ORDER_DESC, sort128_f32_hvx)
|
||||
HTP_ARGSORT_FN(256, asc, GGML_SORT_ORDER_ASC, sort256_f32_hvx)
|
||||
HTP_ARGSORT_FN(256, dsc, GGML_SORT_ORDER_DESC, sort256_f32_hvx)
|
||||
HTP_ARGSORT_FN(512, asc, GGML_SORT_ORDER_ASC, sort512_f32_hvx)
|
||||
HTP_ARGSORT_FN(512, dsc, GGML_SORT_ORDER_DESC, sort512_f32_hvx)
|
||||
HTP_ARGSORT_FN(1024, asc, GGML_SORT_ORDER_ASC, sort1024_f32_hvx)
|
||||
HTP_ARGSORT_FN(1024, dsc, GGML_SORT_ORDER_DESC, sort1024_f32_hvx)
|
||||
|
||||
static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data) {
|
||||
struct htp_argsort_context * actx = (struct htp_argsort_context *)data;
|
||||
struct htp_ops_context * octx = actx->octx;
|
||||
|
||||
@@ -179,7 +382,7 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
|
||||
const struct htp_tensor * dst = octx->dst;
|
||||
|
||||
// Scratchpad memory
|
||||
uint8_t * spad = octx->src0_spad.data + octx->src0_spad.size_per_thread * i;
|
||||
uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;
|
||||
|
||||
// Dimensions
|
||||
uint32_t ne00 = src0->ne[0];
|
||||
@@ -188,12 +391,8 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
|
||||
uint32_t ne03 = src0->ne[3];
|
||||
|
||||
uint32_t nb01 = src0->nb[1];
|
||||
//uint32_t nb02 = src0->nb[2];
|
||||
//uint32_t nb03 = src0->nb[3];
|
||||
|
||||
uint32_t nb1 = dst->nb[1];
|
||||
//uint32_t nb2 = dst->nb[2];
|
||||
//uint32_t nb3 = dst->nb[3];
|
||||
|
||||
// Sort order
|
||||
enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
|
||||
@@ -204,20 +403,17 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
|
||||
uint32_t start_row = rows_per_thread * i;
|
||||
uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);
|
||||
|
||||
// Scratchpad layout:
|
||||
// We need space for one row of float data (values) and one row of int32 indices.
|
||||
// values: ne00 * sizeof(float)
|
||||
// indices: ne00 * sizeof(int32_t)
|
||||
// Padded to 128 bytes.
|
||||
|
||||
size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
|
||||
size_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
|
||||
uint32_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
|
||||
float * values_buf = (float *) spad;
|
||||
int32_t * indices_buf = (int32_t *) (spad + values_size);
|
||||
HVX_Vector * indices_buf_vec = (HVX_Vector *) (spad + values_size);
|
||||
const HVX_Vector ind_init_vec = *(HVX_Vector *)argosrt_ramp_lut;
|
||||
const HVX_Vector ind_diff_vec = Q6_V_vsplat_R(32);
|
||||
|
||||
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL;
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
|
||||
|
||||
for (uint32_t r = start_row; r < end_row; r++) {
|
||||
uint32_t src_offset = r * nb01;
|
||||
uint32_t dst_offset = r * nb1;
|
||||
@@ -245,6 +441,8 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
|
||||
// Copy indices back to DDR
|
||||
hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00);
|
||||
}
|
||||
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
|
||||
}
|
||||
|
||||
int op_argsort(struct htp_ops_context * octx) {
|
||||
@@ -273,11 +471,6 @@ int op_argsort(struct htp_ops_context * octx) {
|
||||
return HTP_STATUS_VTCM_TOO_SMALL;
|
||||
}
|
||||
|
||||
octx->src0_spad.data = octx->ctx->vtcm_base;
|
||||
octx->src0_spad.size = total_spad_size;
|
||||
octx->src0_spad.size_per_thread = spad_per_thread;
|
||||
octx->src0_spad.src = NULL;
|
||||
|
||||
FARF(HIGH, "argsort: %ux%ux%ux%u -> %ux%ux%ux%u (0x%x, 0x%x)",
|
||||
octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
|
||||
octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
|
||||
@@ -286,9 +479,36 @@ int op_argsort(struct htp_ops_context * octx) {
|
||||
struct htp_argsort_context actx;
|
||||
actx.octx = octx;
|
||||
actx.nrows_per_thread = (total_rows + n_threads - 1) / n_threads;
|
||||
actx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
|
||||
actx.vtcm_per_thread = spad_per_thread;
|
||||
|
||||
enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
|
||||
worker_callback_t job_func = htp_argsort_f32_fallback;
|
||||
|
||||
if (order == GGML_SORT_ORDER_ASC) {
|
||||
switch (ne00) {
|
||||
case 1024: job_func = htp_argsort_f32_1024_asc; break;
|
||||
case 512: job_func = htp_argsort_f32_512_asc; break;
|
||||
case 256: job_func = htp_argsort_f32_256_asc; break;
|
||||
case 128: job_func = htp_argsort_f32_128_asc; break;
|
||||
case 64: job_func = htp_argsort_f32_64_asc; break;
|
||||
case 32: job_func = htp_argsort_f32_32_asc; break;
|
||||
default: job_func = htp_argsort_f32_fallback; break;
|
||||
}
|
||||
} else {
|
||||
switch (ne00) {
|
||||
case 1024: job_func = htp_argsort_f32_1024_dsc; break;
|
||||
case 512: job_func = htp_argsort_f32_512_dsc; break;
|
||||
case 256: job_func = htp_argsort_f32_256_dsc; break;
|
||||
case 128: job_func = htp_argsort_f32_128_dsc; break;
|
||||
case 64: job_func = htp_argsort_f32_64_dsc; break;
|
||||
case 32: job_func = htp_argsort_f32_32_dsc; break;
|
||||
default: job_func = htp_argsort_f32_fallback; break;
|
||||
}
|
||||
}
|
||||
|
||||
// Run jobs
|
||||
worker_pool_run_func(octx->ctx->worker_pool, htp_argsort_f32, &actx, n_threads);
|
||||
worker_pool_run_func(octx->ctx->worker_pool, job_func, &actx, n_threads);
|
||||
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
|
||||
@@ -114,7 +114,9 @@ set(GGML_OPENCL_KERNELS
|
||||
mul_mv_id_mxfp4_f32
|
||||
mul_mv_id_mxfp4_f32_flat
|
||||
gemm_moe_q4_0_f32_ns
|
||||
gemm_moe_q4_0_q8_1_dp4a
|
||||
gemv_moe_q4_0_f32_ns
|
||||
gemm_moe_q8_0_f32_ns
|
||||
gemm_moe_q4_1_f32_ns
|
||||
gemv_moe_q4_1_f32_ns
|
||||
gemm_moe_q5_0_f32_ns
|
||||
@@ -122,6 +124,18 @@ set(GGML_OPENCL_KERNELS
|
||||
gemm_moe_q5_1_f32_ns
|
||||
gemv_moe_q5_1_f32_ns
|
||||
gemm_moe_q4_k_f32_ns
|
||||
gemm_moe_q4_k_q8_1_dp4a
|
||||
gemm_moe_q6_k_q8_1_dp4a
|
||||
gemm_moe_q8_1_dp4a
|
||||
moe_reorder_quant_a_q8_1
|
||||
gemm_noshuffle_q4_k_q8_1_dp4a
|
||||
gemm_noshuffle_q5_k_q8_1_dp4a
|
||||
gemm_noshuffle_q6_k_q8_1_dp4a
|
||||
gemm_noshuffle_q8_0_q8_1_dp4a
|
||||
gemm_noshuffle_q5_0_q8_1_dp4a
|
||||
gemm_noshuffle_iq4_nl_q8_1_dp4a
|
||||
gemm_noshuffle_q4_0_q8_1_dp4a
|
||||
quant_a_q8_1
|
||||
gemv_moe_q4_k_f32_ns
|
||||
gemm_moe_q5_k_f32_ns
|
||||
gemv_moe_q5_k_f32_ns
|
||||
@@ -130,8 +144,10 @@ set(GGML_OPENCL_KERNELS
|
||||
gemm_moe_mxfp4_f32
|
||||
gemv_moe_mxfp4_f32
|
||||
gemm_moe_mxfp4_f32_ns
|
||||
gemm_moe_mxfp4_q8_1_dp4a
|
||||
gemv_moe_mxfp4_f32_ns
|
||||
moe_reorder_b
|
||||
moe_combine
|
||||
moe_sort_by_expert
|
||||
mul_mm_f32_f32_l4_lm
|
||||
mul_mm_f16_f32_l4_lm
|
||||
|
||||
+1815
-137
File diff suppressed because it is too large
Load Diff
@@ -2372,3 +2372,121 @@ kernel void kernel_restore_block_iq4_nl_noshuffle(
|
||||
b->qs[2*i + 1] = convert_uchar(((x0 & mask_F0) >> 4) | (x1 & mask_F0));
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// kernel_moe_expand_scale_q8_0
|
||||
//
|
||||
// Expand the q8_0 per-32-block scale d (one half/block, [expert][row][block]) into
|
||||
// the UNIFORM scale[16] format the generic dp4a MoE GEMM (kernel_gemm_moe_q8_1_dp4a,
|
||||
// MOE_QT=80) consumes: 16 f16 per 256-superblock (per-16-element segment), where the
|
||||
// two segments of each 32-block share the block's d. q8_0 is symmetric -> no min
|
||||
// buffer (the GEMM runs with has_min=0). The int8 weight codes are reused verbatim
|
||||
// from the existing flat q8_0 weight buffer (extra0_q8_0->q), so only the scale is
|
||||
// rebuilt here. One work-item per (row, superblock, expert).
|
||||
// ---------------------------------------------------------------------------
|
||||
kernel void kernel_moe_expand_scale_q8_0(
|
||||
global const half * src_d, // [expert][row][block], one scale per 32-block
|
||||
global half * dst_scale, // [expert][row][block][2] (FLAT per-32-block)
|
||||
int ne00,
|
||||
int ne01
|
||||
) {
|
||||
int row = get_global_id(0);
|
||||
int blk = get_global_id(1); // 32-block index along K
|
||||
int e = get_global_id(2);
|
||||
if (row >= ne01) { return; }
|
||||
|
||||
long nb = ne00 / 32; // 32-blocks per row (K only needs % 32 == 0)
|
||||
half d = src_d[((long)e*ne01 + row)*nb + blk];
|
||||
long b = (((long)e*ne01 + row)*nb + blk) * 2;
|
||||
dst_scale[b + 0] = d;
|
||||
dst_scale[b + 1] = d;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// kernel_moe_expand_scale_q5_0
|
||||
//
|
||||
// q5_0 = symmetric, value = d*(code-16), code = nibble | (hi<<4) in 0..31. The
|
||||
// generic dp4a MoE GEMM keeps the unsigned code and centers via the min term:
|
||||
// scale*dp4a(code,a) - min*sum(a), scale = d, min = d*16.
|
||||
// Reads the existing q5_0 d ([expert][block][row], one half/32-block, from the
|
||||
// trans4 convert) and writes the FLAT per-32-block uniform scale[2]/min[1] in
|
||||
// [expert][row][block] order (a transpose). One work-item per (row, block, expert).
|
||||
// ---------------------------------------------------------------------------
|
||||
kernel void kernel_moe_expand_scale_q5_0(
|
||||
global const half * src_d, // [expert][block][row]
|
||||
global half * dst_scale, // [expert][row][block][2]
|
||||
global half * dst_min, // [expert][row][block]
|
||||
int ne00,
|
||||
int ne01
|
||||
) {
|
||||
int row = get_global_id(0);
|
||||
int blk = get_global_id(1);
|
||||
int e = get_global_id(2);
|
||||
if (row >= ne01) { return; }
|
||||
|
||||
long nb = ne00 / 32;
|
||||
half d = src_d[(long)e*nb*ne01 + (long)blk*ne01 + row]; // [expert][block][row]
|
||||
long sb = (((long)e*ne01 + row)*nb + blk) * 2;
|
||||
long mb = ((long)e*ne01 + row)*nb + blk;
|
||||
dst_scale[sb + 0] = d;
|
||||
dst_scale[sb + 1] = d;
|
||||
dst_min[mb] = (half)((float)d * 16.0f);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// kernel_moe_expand_scale_q5_K
|
||||
//
|
||||
// q5_K value = d*sv*code + (-dm*mn), with the 6-bit packed per-sub-block scale sv
|
||||
// and min mn (8 sub-blocks of 32 per 256-superblock, decoded by get_scale_min_k4
|
||||
// from the 12-byte s[]). The generic dp4a MoE GEMM (kernel_gemm_moe_q8_1_dp4a,
|
||||
// MOE_QT=5) keeps the unsigned 5-bit code and applies scale/min via the uniform
|
||||
// per-32-block buffers:
|
||||
// acc += sc0*a_d*raw1 + sc1*a_d*raw2 - mn_u*a_s,
|
||||
// sc0 = sc1 = d*sv (both per-16 segments of a 32-block share the sub-block scale),
|
||||
// mn_u = dm*mn (positive; the GEMM subtracts it -> the -dm*mn min term).
|
||||
// q5_K's q_img (low nibbles) + qh (hi-bit plane) are already in the layout the GEMM
|
||||
// reads (same trans4_ns convert that feeds gemm_moe_q5_k_f32_ns), so only the scale
|
||||
// is rebuilt here.
|
||||
//
|
||||
// One work-item per (row, superblock, expert); each emits 8 sub-blocks.
|
||||
// ---------------------------------------------------------------------------
|
||||
kernel void kernel_moe_expand_scale_q5_K(
|
||||
global const uchar * src_s, // [expert][row][superblock][12]
|
||||
global const half * src_d, // [expert][superblock][row]
|
||||
global const half * src_dm, // [expert][superblock][row]
|
||||
global half * dst_scale, // [expert][row][32block][2]
|
||||
global half * dst_min, // [expert][row][32block]
|
||||
int ne00,
|
||||
int ne01
|
||||
) {
|
||||
int row = get_global_id(0);
|
||||
int sb = get_global_id(1); // superblock index along K
|
||||
int e = get_global_id(2);
|
||||
if (row >= ne01) { return; }
|
||||
|
||||
long nsb = ne00 / 256; // superblocks per row
|
||||
long nblk32 = ne00 / 32; // 32-blocks per row
|
||||
|
||||
float d = (float)src_d [((long)e*nsb + sb)*ne01 + row];
|
||||
float dm = (float)src_dm[((long)e*nsb + sb)*ne01 + row];
|
||||
|
||||
__global const uchar * sc = src_s + ((long)e*ne01 + row)*nsb*12 + (long)sb*12;
|
||||
|
||||
for (int j = 0; j < 8; ++j) {
|
||||
uchar sv, mn;
|
||||
// get_scale_min_k4 (6-bit packed scale/min for sub-block j of 8)
|
||||
if (j < 4) {
|
||||
sv = sc[j] & 63;
|
||||
mn = sc[j+4] & 63;
|
||||
} else {
|
||||
sv = (sc[j+4] & 0x0F) | ((sc[j-4] & 0xC0) >> 2);
|
||||
mn = ((sc[j+4] >> 4) & 0x0F) | ((sc[j] & 0xC0) >> 2);
|
||||
}
|
||||
long sub = (long)sb*8 + j;
|
||||
long sbase = (((long)e*ne01 + row)*nblk32 + sub) * 2;
|
||||
half s_val = (half)(d * (float)sv);
|
||||
dst_scale[sbase + 0] = s_val;
|
||||
dst_scale[sbase + 1] = s_val;
|
||||
dst_min[((long)e*ne01 + row)*nblk32 + sub] = (half)(dm * (float)mn);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#define TILESIZE_M 64
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// 2*mxfp4_value as signed int8, packed 4 codes per uint. Divergent nibble
|
||||
// lookups read a __constant *uint* array + shift, never a byte array
|
||||
// (byte-indexed __constant loads serialize on Adreno and are far slower).
|
||||
// idx 0-3: 0, 1, 2, 3 = 0x03020100
|
||||
// idx 4-7: 4, 6, 8, 12 = 0x0C080604
|
||||
// idx 8-11: 0, -1, -2, -3 = 0xFDFEFF00 (-1=0xFF,-2=0xFE,-3=0xFD)
|
||||
// idx 12-15:-4, -6, -8,-12 = 0xF4F8FAFC (-4=0xFC,-6=0xFA,-8=0xF8,-12=0xF4)
|
||||
__constant uint mxfp4_i8x4[4] = {
|
||||
0x03020100u, 0x0C080604u, 0xFDFEFF00u, 0xF4F8FAFCu
|
||||
};
|
||||
inline uint mxfp4_code(uint n) {
|
||||
return (mxfp4_i8x4[n >> 2] >> ((n & 3u) * 8u)) & 0xFFu;
|
||||
}
|
||||
// 4 nibbles in the low 16 bits of u -> 4 codebook int8, packed for dp4a.
|
||||
inline uint mxfp4_pack(ushort u) {
|
||||
return mxfp4_code((uint)( u & 0xF))
|
||||
| (mxfp4_code((uint)((u >> 4) & 0xF)) << 8)
|
||||
| (mxfp4_code((uint)((u >> 8) & 0xF)) << 16)
|
||||
| (mxfp4_code((uint)((u >> 12) & 0xF)) << 24);
|
||||
}
|
||||
|
||||
static inline float e8m0_to_fp32(uchar x) {
|
||||
int bits;
|
||||
bits = (x == 0) ? 0x00400000 : ((uint) x << 23);
|
||||
return as_float(bits);
|
||||
}
|
||||
|
||||
// One token's dp4a dot (8 uints = 32 K elems) + mxfp4 block-scale epilogue.
|
||||
// blk_scale already carries the 0.5 factor (== 0.5 * 2^e).
|
||||
#define MOE_MXFP4_DP4A_T(t) do { \
|
||||
int raw = 0; \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
|
||||
acc[t] += blk_scale * (float)sh_d[t] * (float)raw; \
|
||||
} while (0)
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_moe_mxfp4_q8_1_dp4a(
|
||||
__read_only image1d_buffer_t src0_q, // mxfp4 codes (transposed, packed nibbles)
|
||||
__global uchar * src0_e, // e8m0 per-32-block scale
|
||||
__global uint * src1_qa, // q8_1 activations: int8 quants (as uint, 4/elem)
|
||||
__global half * src1_da, // q8_1 per-block scale [tok_slot * ne00/32]
|
||||
__global uint * src2, // post-router (orig out positions)
|
||||
__global ushort * src2_emap, // tile -> expert id
|
||||
__write_only image1d_buffer_t dst,
|
||||
__global int * total_tiles,
|
||||
uint ne00,
|
||||
uint ne01,
|
||||
int is_ragged // 1: compute only real tokens per tile
|
||||
) {
|
||||
const uint block_id_m = get_global_id(1); // m_tile
|
||||
const uint block_id_n = get_global_id(2); // n_tile
|
||||
|
||||
if (block_id_n >= total_tiles[0]) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63, == this WI's output row in the M-tile
|
||||
|
||||
const ushort expert_id = src2_emap[block_id_n];
|
||||
const uint row = block_id_m * TILESIZE_M;
|
||||
const uint col = block_id_n * TILESIZE_N;
|
||||
|
||||
const uint num_blocks = ne00 >> 5; // blocks-of-32 per token
|
||||
const uint row_idx = row + lid;
|
||||
|
||||
const uint ne00_u = ne00 >> 2; // ne00 in uint (int8x4) units
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8]; // 32 tokens x 8 uints (32 int8) = 1 KiB
|
||||
__local half sh_d[TILESIZE_N];
|
||||
|
||||
// Real token count for this tile.
|
||||
// Real tokens are packed contiguously at the tile start; padded slots hold
|
||||
// 0xFFFFFFFF (only the last tile of each expert is partial). is_ragged skips
|
||||
// the dp4a/staging/scatter for padded slots; is_ragged==0 forces n_real=32.
|
||||
__local uint sh_src2[TILESIZE_N];
|
||||
__local int sh_nreal;
|
||||
if (lid < TILESIZE_N) {
|
||||
sh_src2[lid] = src2[col + lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (lid == 0) {
|
||||
int nr = TILESIZE_N;
|
||||
if (is_ragged) {
|
||||
nr = 0;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) {
|
||||
if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
|
||||
}
|
||||
}
|
||||
sh_nreal = nr;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
const int n_real = sh_nreal;
|
||||
|
||||
float acc[TILESIZE_N];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
|
||||
|
||||
for (uint step = 0; step < ne00; step += 32) {
|
||||
const uint sub = step >> 5; // 32-block index along K
|
||||
|
||||
// e8m0 block scale for this WI's row, this 32-block (folded x0.5)
|
||||
const uint e_offset = row_idx + sub * ne01 + expert_id * num_blocks * ne01;
|
||||
const float blk_scale = 0.5f * e8m0_to_fp32(src0_e[e_offset]);
|
||||
|
||||
// repack this WI's 32 weight nibbles into 8 dp4a uints
|
||||
const uint qoff0 = row + ((ne01 * step) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint r0 = read_imageui(src0_q, qoff0 + lid).x;
|
||||
const uint r1 = read_imageui(src0_q, qoff0 + lid + ne01).x;
|
||||
const uint r2 = read_imageui(src0_q, qoff1 + lid).x;
|
||||
const uint r3 = read_imageui(src0_q, qoff1 + lid + ne01).x;
|
||||
uint qw[8];
|
||||
qw[0] = mxfp4_pack((ushort)(r0)); qw[1] = mxfp4_pack((ushort)(r0 >> 16));
|
||||
qw[2] = mxfp4_pack((ushort)(r1)); qw[3] = mxfp4_pack((ushort)(r1 >> 16));
|
||||
qw[4] = mxfp4_pack((ushort)(r2)); qw[5] = mxfp4_pack((ushort)(r2 >> 16));
|
||||
qw[6] = mxfp4_pack((ushort)(r3)); qw[7] = mxfp4_pack((ushort)(r3 >> 16));
|
||||
|
||||
// cooperatively stage the n_real-token x 32-K int8 activations
|
||||
const uint stage_lim = (uint)n_real * 8;
|
||||
for (uint idx = lid; idx < stage_lim; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
|
||||
}
|
||||
if (lid < (uint)n_real) {
|
||||
sh_d[lid] = src1_da[(col + lid) * num_blocks + sub];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
// Full tiles keep the fully-unrolled 32-wide loop; partial tiles run only n_real
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) { MOE_MXFP4_DP4A_T(t); }
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int t = 0; t < n_real; ++t) { MOE_MXFP4_DP4A_T(t); }
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (row_idx >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
// scatter results to original output rows (reuse sh_src2 from the top)
|
||||
__local uint out_idx[TILESIZE_N];
|
||||
if (lid < TILESIZE_N) {
|
||||
uint idx = sh_src2[lid];
|
||||
if (idx == 0xFFFFFFFF) {
|
||||
idx = sh_src2[0];
|
||||
}
|
||||
out_idx[lid] = idx * ne01;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const uint m_offset = row + lid;
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 1; t < TILESIZE_N; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
barrier(CLK_GLOBAL_MEM_FENCE);
|
||||
write_imagef(dst, out_idx[0] + m_offset, acc[0]);
|
||||
} else {
|
||||
for (int t = 0; t < n_real; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#define TILESIZE_M 64
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// Expand the 4 nibbles held in the low 16 bits of `u` into 4 bytes (one nibble
|
||||
// per byte, value 0..15), packed for the int8 dp4a. The -8 zero-point is applied
|
||||
// in the epilogue via the activation sum term (cheaper than biasing every byte).
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
// One token's dp4a dot (8 uints = 32 K elems) + q4_0 scale/zero-point epilogue.
|
||||
#define MOE_Q40_DP4A_T(t) do { \
|
||||
int raw = 0; \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
|
||||
acc[t] += d_val * ((float)sh_d[t] * (float)raw - 8.0f * (float)sh_s[t]); \
|
||||
} while (0)
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_moe_q4_0_q8_1_dp4a(
|
||||
__read_only image1d_buffer_t src0_q, // q4_0 weights (transposed, packed nibbles)
|
||||
__global half * src0_d, // per-32-block scale
|
||||
__global uint * src1_qa, // q8_1 activations: int8 quants (as uint, 4/elem)
|
||||
__global half * src1_da, // q8_1 per-block scale [tok_slot * ne00/32]
|
||||
__global half * src1_sa, // q8_1 per-block sum*d [tok_slot * ne00/32]
|
||||
__global uint * src2, // post-router (orig out positions)
|
||||
__global ushort * src2_emap,// tile -> expert id
|
||||
__write_only image1d_buffer_t dst,
|
||||
__global int * total_tiles,
|
||||
uint ne00,
|
||||
uint ne01,
|
||||
int is_ragged // 1: compute only real tokens per tile
|
||||
) {
|
||||
const uint block_id_m = get_global_id(1); // m_tile
|
||||
const uint block_id_n = get_global_id(2); // n_tile
|
||||
|
||||
if (block_id_n >= total_tiles[0]) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63, == this WI's output row in the M-tile
|
||||
|
||||
const ushort expert_id = src2_emap[block_id_n];
|
||||
const uint row = block_id_m * TILESIZE_M;
|
||||
const uint col = block_id_n * TILESIZE_N;
|
||||
|
||||
const uint num_blocks = ne00 >> 5; // blocks-of-32 per token
|
||||
const uint row_idx = row + lid;
|
||||
|
||||
const uint ne00_u = ne00 >> 2; // ne00 in uint (int8x4) units
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8]; // 32 tokens x 8 uints (32 int8) = 1 KiB
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
// Real-token count for this tile
|
||||
__local uint sh_src2[TILESIZE_N];
|
||||
__local int sh_nreal;
|
||||
if (lid < TILESIZE_N) {
|
||||
sh_src2[lid] = src2[col + lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (lid == 0) {
|
||||
int nr = TILESIZE_N;
|
||||
if (is_ragged) {
|
||||
nr = 0;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) {
|
||||
if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
|
||||
}
|
||||
}
|
||||
sh_nreal = nr;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
const int n_real = sh_nreal;
|
||||
|
||||
float acc[TILESIZE_N];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
|
||||
|
||||
for (uint step = 0; step < ne00; step += 32) {
|
||||
const uint sub = step >> 5; // 32-block index along K
|
||||
|
||||
// per-32-block scale for this WI's row
|
||||
const uint d_offset = row_idx + sub * ne01 + expert_id * num_blocks * ne01;
|
||||
const float d_val = (float)src0_d[d_offset];
|
||||
|
||||
// repack this WI's 32 weight nibbles into 8 dp4a uints
|
||||
const uint qoff0 = row + ((ne01 * step) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint r0 = read_imageui(src0_q, qoff0 + lid).x;
|
||||
const uint r1 = read_imageui(src0_q, qoff0 + lid + ne01).x;
|
||||
const uint r2 = read_imageui(src0_q, qoff1 + lid).x;
|
||||
const uint r3 = read_imageui(src0_q, qoff1 + lid + ne01).x;
|
||||
uint qw[8];
|
||||
qw[0] = EXP4(r0); qw[1] = EXP4(r0 >> 16);
|
||||
qw[2] = EXP4(r1); qw[3] = EXP4(r1 >> 16);
|
||||
qw[4] = EXP4(r2); qw[5] = EXP4(r2 >> 16);
|
||||
qw[6] = EXP4(r3); qw[7] = EXP4(r3 >> 16);
|
||||
|
||||
// cooperatively stage the n_real-token x 32-K int8 activations
|
||||
const uint stage_lim = (uint)n_real * 8;
|
||||
for (uint idx = lid; idx < stage_lim; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
|
||||
}
|
||||
if (lid < (uint)n_real) {
|
||||
sh_d[lid] = src1_da[(col + lid) * num_blocks + sub];
|
||||
sh_s[lid] = src1_sa[(col + lid) * num_blocks + sub];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) { MOE_Q40_DP4A_T(t); }
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int t = 0; t < n_real; ++t) { MOE_Q40_DP4A_T(t); }
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (row_idx >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
// scatter results to original output rows (reuse sh_src2 from the top)
|
||||
__local uint out_idx[TILESIZE_N];
|
||||
if (lid < TILESIZE_N) {
|
||||
uint idx = sh_src2[lid];
|
||||
if (idx == 0xFFFFFFFF) {
|
||||
idx = sh_src2[0];
|
||||
}
|
||||
out_idx[lid] = idx * ne01;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const uint m_offset = row + lid;
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 1; t < TILESIZE_N; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
barrier(CLK_GLOBAL_MEM_FENCE);
|
||||
write_imagef(dst, out_idx[0] + m_offset, acc[0]);
|
||||
} else {
|
||||
for (int t = 0; t < n_real; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
// q4_K subblock (32 elems): w_i = scale*q_i - minv, q_i in [0,15], scale =
|
||||
// d_super*sv6, minv = dmin_super*mn6. With activation block (a_d, a_s, qa[32]):
|
||||
// Sum_i w_i * a_i = scale * a_d * dp4a(q, qa) - minv * a_s
|
||||
// where a_s = a_d * Sum(qa) (the q8_1 "s" field)
|
||||
|
||||
#define TILESIZE_M 64
|
||||
#define TILESIZE_N 32
|
||||
#define QK_K 256
|
||||
#define K_SCALE_SIZE 12
|
||||
|
||||
inline void get_scale_min_k4(
|
||||
int j,
|
||||
global const uchar * q,
|
||||
uchar * d,
|
||||
uchar * m
|
||||
) {
|
||||
if (j < 4) {
|
||||
*d = q[j] & 63;
|
||||
*m = q[j+4] & 63;
|
||||
} else {
|
||||
*d = (q[j+4] & 0x0F) | ((q[j-4] & 0xC0) >> 2);
|
||||
*m = ((q[j+4] >> 4) & 0x0F) | ((q[j] & 0xC0) >> 2);
|
||||
}
|
||||
}
|
||||
|
||||
// Expand the 4 nibbles held in the low 16 bits of `u` into 4 bytes (one nibble
|
||||
// per byte, value 0..15), packed for the int8 dp4a.
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
// One token's dp4a dot (8 uints = 32 K elems) + q4_K scale/min epilogue into acc[t].
|
||||
#define MOE_Q4K_DP4A_T(t) do { \
|
||||
int raw = 0; \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
|
||||
acc[t] += scale * (float)sh_d[t] * (float)raw - minv * (float)sh_s[t]; \
|
||||
} while (0)
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_moe_q4_k_q8_1_dp4a(
|
||||
__read_only image1d_buffer_t src0_q, // q4_K weights (transposed, packed nibbles)
|
||||
__global half * src0_d, // per-superblock scale
|
||||
__global half * src0_dm, // per-superblock min
|
||||
__global uchar * src0_s, // 6-bit scale/min codes
|
||||
__global uint * src1_qa, // q8_1 activations: int8 quants (as uint, 4/elem)
|
||||
__global half * src1_da, // q8_1 per-block scale [tok_slot * ne00/32]
|
||||
__global half * src1_sa, // q8_1 per-block sum*d [tok_slot * ne00/32]
|
||||
__global uint * src2, // post-router (orig out positions)
|
||||
__global ushort * src2_emap,// tile -> expert id
|
||||
__write_only image1d_buffer_t dst,
|
||||
__global int * total_tiles,
|
||||
uint ne00,
|
||||
uint ne01,
|
||||
int is_ragged // 1: compute only real tokens per tile
|
||||
) {
|
||||
const uint block_id_m = get_global_id(1); // m_tile
|
||||
const uint block_id_n = get_global_id(2); // n_tile
|
||||
|
||||
if (block_id_n >= total_tiles[0]) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63, == this WI's output row in the M-tile
|
||||
|
||||
const ushort expert_id = src2_emap[block_id_n];
|
||||
const uint row = block_id_m * TILESIZE_M;
|
||||
const uint col = block_id_n * TILESIZE_N;
|
||||
|
||||
const uint num_superblocks = ne00 / QK_K;
|
||||
const uint scales_per_row = num_superblocks * K_SCALE_SIZE;
|
||||
const uint row_idx = row + lid;
|
||||
|
||||
const uint ne00_u = ne00 >> 2; // ne00 in uint (int8x4) units
|
||||
const uint ne00_b = ne00 >> 5; // blocks-of-32 per token
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8]; // 32 tokens x 8 uints (32 int8) = 1 KiB
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
// Real token count for this tile
|
||||
__local uint sh_src2[TILESIZE_N];
|
||||
__local int sh_nreal;
|
||||
if (lid < TILESIZE_N) {
|
||||
sh_src2[lid] = src2[col + lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (lid == 0) {
|
||||
int nr = TILESIZE_N;
|
||||
if (is_ragged) {
|
||||
nr = 0;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) {
|
||||
if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
|
||||
}
|
||||
}
|
||||
sh_nreal = nr;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
const int n_real = sh_nreal;
|
||||
|
||||
float acc[TILESIZE_N];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
|
||||
|
||||
for (uint step = 0; step < ne00; step += 32) {
|
||||
const uint sub = step >> 5; // subblock index along K
|
||||
const uint sb = sub >> 3; // superblock index
|
||||
const uint j = sub & 7; // subblock within superblock
|
||||
|
||||
// --- weight scale / min for this WI's row, this subblock ---
|
||||
const uint d_offset = row + sb * ne01 + expert_id * num_superblocks * ne01 + lid;
|
||||
const float d_val = (float)src0_d[d_offset];
|
||||
const float dm_val = (float)src0_dm[d_offset];
|
||||
|
||||
global const uchar * sc = src0_s + (expert_id * ne01 + row_idx) * scales_per_row + sb * K_SCALE_SIZE;
|
||||
uchar sv, mn;
|
||||
get_scale_min_k4(j, sc, &sv, &mn);
|
||||
const float scale = d_val * (float)sv;
|
||||
const float minv = dm_val * (float)mn;
|
||||
|
||||
// --- repack this WI's 32 weight nibbles into 8 dp4a uints ---
|
||||
const uint qoff0 = row + ((ne01 * step) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint r0 = read_imageui(src0_q, qoff0 + lid).x;
|
||||
const uint r1 = read_imageui(src0_q, qoff0 + lid + ne01).x;
|
||||
const uint r2 = read_imageui(src0_q, qoff1 + lid).x;
|
||||
const uint r3 = read_imageui(src0_q, qoff1 + lid + ne01).x;
|
||||
uint qw[8];
|
||||
qw[0] = EXP4(r0); qw[1] = EXP4(r0 >> 16);
|
||||
qw[2] = EXP4(r1); qw[3] = EXP4(r1 >> 16);
|
||||
qw[4] = EXP4(r2); qw[5] = EXP4(r2 >> 16);
|
||||
qw[6] = EXP4(r3); qw[7] = EXP4(r3 >> 16);
|
||||
|
||||
// --- cooperatively stage the n_real-token x 32-K int8 activations to LDS ---
|
||||
const uint stage_lim = (uint)n_real * 8;
|
||||
for (uint idx = lid; idx < stage_lim; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
|
||||
}
|
||||
if (lid < (uint)n_real) {
|
||||
sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];
|
||||
sh_s[lid] = src1_sa[(col + lid) * ne00_b + sub];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
// dp4a - each real token sum over 8 uints (32 K), then scale/min
|
||||
// Full tiles keep the fully-unrolled 32-wide loop;
|
||||
// partial tiles run only n_real (saves the padded-slot dp4a + staging).
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) { MOE_Q4K_DP4A_T(t); }
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int t = 0; t < n_real; ++t) { MOE_Q4K_DP4A_T(t); }
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (row_idx >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
// scatter results to original output rows
|
||||
__local uint out_idx[TILESIZE_N];
|
||||
if (lid < TILESIZE_N) {
|
||||
uint idx = sh_src2[lid];
|
||||
if (idx == 0xFFFFFFFF) {
|
||||
idx = sh_src2[0];
|
||||
}
|
||||
out_idx[lid] = idx * ne01;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const uint m_offset = row + lid;
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 1; t < TILESIZE_N; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
barrier(CLK_GLOBAL_MEM_FENCE);
|
||||
write_imagef(dst, out_idx[0] + m_offset, acc[0]);
|
||||
} else {
|
||||
for (int t = 0; t < n_real; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#define TILESIZE_N 32
|
||||
#define QK_K 256
|
||||
|
||||
// 4 nibbles in the low 16 bits of `u` -> 4 bytes (value 0..15, in bits 0-3).
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
// 4 2-bit highs in byte `b` (8 bits) -> 4 bytes, value 0..3 in bits 4-5
|
||||
// (pre-multiplied by 16 so it ORs with the EXP4 nibble to form q6 in 0..63).
|
||||
#define EXP2(b) ( (((uint)((b) & 0x03u)) << 4) | \
|
||||
(((uint)((b) & 0x0Cu)) << 10) | \
|
||||
(((uint)((b) & 0x30u)) << 16) | \
|
||||
(((uint)((b) & 0xC0u)) << 22) )
|
||||
|
||||
// q6 (0..63, bits 0-5 of each byte) -> (q6-32) as a signed int8 per byte.
|
||||
// Flipping bit5 subtracts 32 in 6-bit two's complement; then replicate bit5
|
||||
// into bits 6-7 to sign-extend to int8. Per-byte, no inter-byte carry.
|
||||
inline uint SIGN6(uint q6p) {
|
||||
uint x = q6p ^ 0x20202020u;
|
||||
uint s = x & 0x20202020u;
|
||||
return x | (s << 1) | (s << 2);
|
||||
}
|
||||
|
||||
inline int dp4a_q6(uint qw0, uint qw1, uint qw2, uint qw3,
|
||||
uint a0, uint a1, uint a2, uint a3) {
|
||||
int raw = 0;
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw0, a0, raw);
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw1, a1, raw);
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw2, a2, raw);
|
||||
raw = dot_acc_sat_4x8packed_ss_int(qw3, a3, raw);
|
||||
return raw;
|
||||
}
|
||||
|
||||
// One token's q6_K dp4a dot (two halves, per-16 scales) + epilogue into acc[t].
|
||||
#define MOE_Q6K_DP4A_T(t) do { \
|
||||
const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], sh_qa[t][0], sh_qa[t][1], sh_qa[t][2], sh_qa[t][3]); \
|
||||
const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], sh_qa[t][4], sh_qa[t][5], sh_qa[t][6], sh_qa[t][7]); \
|
||||
const float a_d = (float)sh_d[t]; \
|
||||
acc[t] += scale0 * a_d * (float)raw1 + scale1 * a_d * (float)raw2; \
|
||||
} while (0)
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_moe_q6_k_q8_1_dp4a(
|
||||
__read_only image1d_buffer_t src0_ql, // q6_K low nibbles (image, q4_K-style layout)
|
||||
__global uint * src0_qh, // q6_K high 2-bit (16 elems/uint)
|
||||
__global char * src0_s, // int8 scales (one per 16 elems)
|
||||
__global half * src0_d, // per-superblock scale
|
||||
__global uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem)
|
||||
__global half * src1_da, // q8_1 per-block scale [tok_slot * ne00/32]
|
||||
__global uint * src2, // post-router (orig out positions)
|
||||
__global ushort * src2_emap, // tile -> expert id
|
||||
__write_only image1d_buffer_t dst,
|
||||
__global int * total_tiles,
|
||||
uint ne00,
|
||||
uint ne01,
|
||||
int is_ragged // 1: compute only real tokens per tile
|
||||
) {
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
if (block_id_n >= total_tiles[0]) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within M-tile
|
||||
|
||||
const ushort expert_id = src2_emap[block_id_n];
|
||||
const uint row = block_id_m * 64;
|
||||
const uint col = block_id_n * TILESIZE_N;
|
||||
|
||||
const uint num_superblocks = ne00 / QK_K;
|
||||
const uint scales_per_row = num_superblocks * 16;
|
||||
const uint row_idx = row + lid;
|
||||
|
||||
const uint ne00_u = ne00 >> 2;
|
||||
const uint ne00_b = ne00 >> 5;
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
|
||||
// Real token count for this tile
|
||||
__local uint sh_src2[TILESIZE_N];
|
||||
__local int sh_nreal;
|
||||
if (lid < TILESIZE_N) {
|
||||
sh_src2[lid] = src2[col + lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (lid == 0) {
|
||||
int nr = TILESIZE_N;
|
||||
if (is_ragged) {
|
||||
nr = 0;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) {
|
||||
if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
|
||||
}
|
||||
}
|
||||
sh_nreal = nr;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
const int n_real = sh_nreal;
|
||||
|
||||
float acc[TILESIZE_N];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
|
||||
|
||||
for (uint step = 0; step < ne00; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
const uint sb = sub >> 3;
|
||||
const uint j = sub & 7;
|
||||
|
||||
const float d_val = (float)src0_d[row + sb * ne01 + expert_id * num_superblocks * ne01 + lid];
|
||||
global const char * sc = src0_s + (expert_id * ne01 + row_idx) * scales_per_row + sb * 16;
|
||||
const float scale0 = d_val * (float)sc[j * 2];
|
||||
const float scale1 = d_val * (float)sc[j * 2 + 1];
|
||||
|
||||
// high bits: one uint covers 16 elems; first/second 16 of this 32-block
|
||||
const uint qh_base = row + (sub * 2) * ne01 + expert_id * (num_superblocks * 16) * ne01 + lid;
|
||||
const uint qh1 = src0_qh[qh_base];
|
||||
const uint qh2 = src0_qh[qh_base + ne01];
|
||||
|
||||
// low nibbles: same image layout as q4_K (8 ushorts over the 32 K)
|
||||
const uint qoff0 = row + ((ne01 * step) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
|
||||
const uint r0 = read_imageui(src0_ql, qoff0 + lid).x;
|
||||
const uint r1 = read_imageui(src0_ql, qoff0 + lid + ne01).x;
|
||||
const uint r2 = read_imageui(src0_ql, qoff1 + lid).x;
|
||||
const uint r3 = read_imageui(src0_ql, qoff1 + lid + ne01).x;
|
||||
|
||||
uint qw[8];
|
||||
qw[0] = SIGN6(EXP4(r0) | EXP2((qh1) & 0xFFu));
|
||||
qw[1] = SIGN6(EXP4(r0 >> 16) | EXP2((qh1 >> 8) & 0xFFu));
|
||||
qw[2] = SIGN6(EXP4(r1) | EXP2((qh1 >> 16) & 0xFFu));
|
||||
qw[3] = SIGN6(EXP4(r1 >> 16) | EXP2((qh1 >> 24) & 0xFFu));
|
||||
qw[4] = SIGN6(EXP4(r2) | EXP2((qh2) & 0xFFu));
|
||||
qw[5] = SIGN6(EXP4(r2 >> 16) | EXP2((qh2 >> 8) & 0xFFu));
|
||||
qw[6] = SIGN6(EXP4(r3) | EXP2((qh2 >> 16) & 0xFFu));
|
||||
qw[7] = SIGN6(EXP4(r3 >> 16) | EXP2((qh2 >> 24) & 0xFFu));
|
||||
|
||||
const uint stage_lim = (uint)n_real * 8;
|
||||
for (uint idx = lid; idx < stage_lim; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
|
||||
}
|
||||
if (lid < (uint)n_real) {
|
||||
sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
// Full tiles keep the fully-unrolled 32-wide loop; partial tiles run n_real.
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) { MOE_Q6K_DP4A_T(t); }
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int t = 0; t < n_real; ++t) { MOE_Q6K_DP4A_T(t); }
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (row_idx >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
__local uint out_idx[TILESIZE_N];
|
||||
if (lid < TILESIZE_N) {
|
||||
uint idx = sh_src2[lid];
|
||||
if (idx == 0xFFFFFFFF) {
|
||||
idx = sh_src2[0];
|
||||
}
|
||||
out_idx[lid] = idx * ne01;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const uint m_offset = row + lid;
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 1; t < TILESIZE_N; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
barrier(CLK_GLOBAL_MEM_FENCE);
|
||||
write_imagef(dst, out_idx[0] + m_offset, acc[0]);
|
||||
} else {
|
||||
for (int t = 0; t < n_real; ++t) {
|
||||
write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#pragma OPENCL EXTENSION cl_qcom_subgroup_uniform_load: enable
|
||||
#pragma OPENCL EXTENSION cl_qcom_subgroup_constant_load: enable
|
||||
#pragma OPENCL EXTENSION cl_qcom_extra_vector_types : enable
|
||||
|
||||
#define TILESIZE_K 16
|
||||
#define TILESIZE_M 64
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// q8_0: 16 signed int8 weights (one uint4 = 16 chars) -> half16, scaled.
|
||||
#define dequantize_q8_0(q4, a_f16, scale) \
|
||||
a_f16 = convert_half16(as_char16(q4)) * scale;
|
||||
|
||||
#define dotx16_reduce8(a_reg, b_lm, c_reg, lm_offset) \
|
||||
acc.s0 = dot(a_reg.s0123, b_lm[lm_offset + 0]); \
|
||||
acc.s1 = dot(a_reg.s0123, b_lm[lm_offset + 1]); \
|
||||
acc.s2 = dot(a_reg.s0123, b_lm[lm_offset + 2]); \
|
||||
acc.s3 = dot(a_reg.s0123, b_lm[lm_offset + 3]); \
|
||||
acc.s4 = dot(a_reg.s0123, b_lm[lm_offset + 4]); \
|
||||
acc.s5 = dot(a_reg.s0123, b_lm[lm_offset + 5]); \
|
||||
acc.s6 = dot(a_reg.s0123, b_lm[lm_offset + 6]); \
|
||||
acc.s7 = dot(a_reg.s0123, b_lm[lm_offset + 7]); \
|
||||
acc.s8 = dot(a_reg.s0123, b_lm[lm_offset + 8]); \
|
||||
acc.s9 = dot(a_reg.s0123, b_lm[lm_offset + 9]); \
|
||||
acc.sa = dot(a_reg.s0123, b_lm[lm_offset + 10]); \
|
||||
acc.sb = dot(a_reg.s0123, b_lm[lm_offset + 11]); \
|
||||
acc.sc = dot(a_reg.s0123, b_lm[lm_offset + 12]); \
|
||||
acc.sd = dot(a_reg.s0123, b_lm[lm_offset + 13]); \
|
||||
acc.se = dot(a_reg.s0123, b_lm[lm_offset + 14]); \
|
||||
acc.sf = dot(a_reg.s0123, b_lm[lm_offset + 15]); \
|
||||
acc.s0 += dot(a_reg.s4567, b_lm[lm_offset + 32]); \
|
||||
acc.s1 += dot(a_reg.s4567, b_lm[lm_offset + 33]); \
|
||||
acc.s2 += dot(a_reg.s4567, b_lm[lm_offset + 34]); \
|
||||
acc.s3 += dot(a_reg.s4567, b_lm[lm_offset + 35]); \
|
||||
acc.s4 += dot(a_reg.s4567, b_lm[lm_offset + 36]); \
|
||||
acc.s5 += dot(a_reg.s4567, b_lm[lm_offset + 37]); \
|
||||
acc.s6 += dot(a_reg.s4567, b_lm[lm_offset + 38]); \
|
||||
acc.s7 += dot(a_reg.s4567, b_lm[lm_offset + 39]); \
|
||||
acc.s8 += dot(a_reg.s4567, b_lm[lm_offset + 40]); \
|
||||
acc.s9 += dot(a_reg.s4567, b_lm[lm_offset + 41]); \
|
||||
acc.sa += dot(a_reg.s4567, b_lm[lm_offset + 42]); \
|
||||
acc.sb += dot(a_reg.s4567, b_lm[lm_offset + 43]); \
|
||||
acc.sc += dot(a_reg.s4567, b_lm[lm_offset + 44]); \
|
||||
acc.sd += dot(a_reg.s4567, b_lm[lm_offset + 45]); \
|
||||
acc.se += dot(a_reg.s4567, b_lm[lm_offset + 46]); \
|
||||
acc.sf += dot(a_reg.s4567, b_lm[lm_offset + 47]); \
|
||||
c_reg.lo += convert_float8(acc.lo); \
|
||||
c_reg.hi += convert_float8(acc.hi); \
|
||||
acc.s0 = dot(a_reg.s89ab, b_lm[lm_offset + 64]); \
|
||||
acc.s1 = dot(a_reg.s89ab, b_lm[lm_offset + 65]); \
|
||||
acc.s2 = dot(a_reg.s89ab, b_lm[lm_offset + 66]); \
|
||||
acc.s3 = dot(a_reg.s89ab, b_lm[lm_offset + 67]); \
|
||||
acc.s4 = dot(a_reg.s89ab, b_lm[lm_offset + 68]); \
|
||||
acc.s5 = dot(a_reg.s89ab, b_lm[lm_offset + 69]); \
|
||||
acc.s6 = dot(a_reg.s89ab, b_lm[lm_offset + 70]); \
|
||||
acc.s7 = dot(a_reg.s89ab, b_lm[lm_offset + 71]); \
|
||||
acc.s8 = dot(a_reg.s89ab, b_lm[lm_offset + 72]); \
|
||||
acc.s9 = dot(a_reg.s89ab, b_lm[lm_offset + 73]); \
|
||||
acc.sa = dot(a_reg.s89ab, b_lm[lm_offset + 74]); \
|
||||
acc.sb = dot(a_reg.s89ab, b_lm[lm_offset + 75]); \
|
||||
acc.sc = dot(a_reg.s89ab, b_lm[lm_offset + 76]); \
|
||||
acc.sd = dot(a_reg.s89ab, b_lm[lm_offset + 77]); \
|
||||
acc.se = dot(a_reg.s89ab, b_lm[lm_offset + 78]); \
|
||||
acc.sf = dot(a_reg.s89ab, b_lm[lm_offset + 79]); \
|
||||
acc.s0 += dot(a_reg.scdef, b_lm[lm_offset + 96]); \
|
||||
acc.s1 += dot(a_reg.scdef, b_lm[lm_offset + 97]); \
|
||||
acc.s2 += dot(a_reg.scdef, b_lm[lm_offset + 98]); \
|
||||
acc.s3 += dot(a_reg.scdef, b_lm[lm_offset + 99]); \
|
||||
acc.s4 += dot(a_reg.scdef, b_lm[lm_offset + 100]); \
|
||||
acc.s5 += dot(a_reg.scdef, b_lm[lm_offset + 101]); \
|
||||
acc.s6 += dot(a_reg.scdef, b_lm[lm_offset + 102]); \
|
||||
acc.s7 += dot(a_reg.scdef, b_lm[lm_offset + 103]); \
|
||||
acc.s8 += dot(a_reg.scdef, b_lm[lm_offset + 104]); \
|
||||
acc.s9 += dot(a_reg.scdef, b_lm[lm_offset + 105]); \
|
||||
acc.sa += dot(a_reg.scdef, b_lm[lm_offset + 106]); \
|
||||
acc.sb += dot(a_reg.scdef, b_lm[lm_offset + 107]); \
|
||||
acc.sc += dot(a_reg.scdef, b_lm[lm_offset + 108]); \
|
||||
acc.sd += dot(a_reg.scdef, b_lm[lm_offset + 109]); \
|
||||
acc.se += dot(a_reg.scdef, b_lm[lm_offset + 110]); \
|
||||
acc.sf += dot(a_reg.scdef, b_lm[lm_offset + 111]); \
|
||||
c_reg.lo += convert_float8(acc.lo); \
|
||||
c_reg.hi += convert_float8(acc.hi); \
|
||||
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_moe_q8_0_f32_ns(
|
||||
__global char * src0_q, // flat q8_0 quants [n_expert*ne01*ne00]
|
||||
__global half * src0_d, // flat q8_0 scales [n_expert*ne01*nb]
|
||||
__read_only image1d_buffer_t src1, // reordered activations (f32)
|
||||
__global uint * src2, // post-router out indices
|
||||
__global ushort * src2_emap,// expert per tile
|
||||
__write_only image1d_buffer_t dst,
|
||||
__global int * total_tiles,
|
||||
uint ne00,
|
||||
uint ne01
|
||||
) {
|
||||
uint block_id_m = get_global_id(1); // m_tile
|
||||
uint block_id_n = get_global_id(2); // n_tile
|
||||
|
||||
if (block_id_n >= total_tiles[0]) {
|
||||
return;
|
||||
}
|
||||
|
||||
__private half16 reg_a;
|
||||
__private float32 reg_c = (float32)(0);
|
||||
__local half4 shared_b[128];
|
||||
|
||||
const ushort expert_id = src2_emap[block_id_n];
|
||||
|
||||
const uint row = block_id_m * TILESIZE_M;
|
||||
const uint col = block_id_n * TILESIZE_N;
|
||||
|
||||
const uint nb = ne00 >> 5; // blocks per row (ne00/32)
|
||||
const uint w_row = expert_id * ne01 + row + get_local_id(0); // this lane's output row
|
||||
__global char * w_q = src0_q + (ulong)w_row * ne00; // char base for the row
|
||||
__global half * w_d = src0_d + (ulong)w_row * nb; // scale base for the row
|
||||
|
||||
uint sub_block_id_m = get_local_id(0);
|
||||
uint2 b_global_offset;
|
||||
b_global_offset.x = ((sub_block_id_m & 3) << 2) + (sub_block_id_m >> 2) * ne00;
|
||||
b_global_offset.y = b_global_offset.x + (16 * ne00);
|
||||
uint2 b_local_offset;
|
||||
b_local_offset.x = (sub_block_id_m & 3) * 32 + (sub_block_id_m >> 2);
|
||||
b_local_offset.y = b_local_offset.x + 16;
|
||||
|
||||
// Loop along K axis, 32 elements per iteration, split into 2 sub-blocks.
|
||||
for (uint step = 0; step < ne00; step += TILESIZE_K * 2) {
|
||||
half s = w_d[step >> 5]; // one q8_0 scale per 32-element block
|
||||
|
||||
// First sub-block: 16 weights (16 chars = one uint4) at K=step
|
||||
uint4 q8x16 = *((__global uint4 *)(w_q + step));
|
||||
|
||||
uint b_sub_offset = col * ne00 + step;
|
||||
float8 bx8_f32;
|
||||
bx8_f32.lo = read_imagef(src1, (b_sub_offset + b_global_offset.x) / 4);
|
||||
bx8_f32.hi = read_imagef(src1, (b_sub_offset + b_global_offset.y) / 4);
|
||||
half8 bx8_f16 = convert_half8(bx8_f32);
|
||||
shared_b[b_local_offset.x] = bx8_f16.lo;
|
||||
shared_b[b_local_offset.y] = bx8_f16.hi;
|
||||
|
||||
dequantize_q8_0(q8x16, reg_a, s);
|
||||
|
||||
sub_group_barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
half16 acc;
|
||||
dotx16_reduce8(reg_a, shared_b, reg_c.lo, 0);
|
||||
dotx16_reduce8(reg_a, shared_b, reg_c.hi, 16);
|
||||
|
||||
// Second sub-block: next 16 weights at K=step+16
|
||||
uint half_step = step + TILESIZE_K;
|
||||
q8x16 = *((__global uint4 *)(w_q + half_step));
|
||||
b_sub_offset = col * ne00 + half_step;
|
||||
|
||||
bx8_f32.lo = read_imagef(src1, (b_sub_offset + b_global_offset.x) / 4);
|
||||
bx8_f32.hi = read_imagef(src1, (b_sub_offset + b_global_offset.y) / 4);
|
||||
bx8_f16 = convert_half8(bx8_f32);
|
||||
shared_b[b_local_offset.x] = bx8_f16.lo;
|
||||
shared_b[b_local_offset.y] = bx8_f16.hi;
|
||||
|
||||
dequantize_q8_0(q8x16, reg_a, s);
|
||||
|
||||
sub_group_barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
dotx16_reduce8(reg_a, shared_b, reg_c.lo, 0);
|
||||
dotx16_reduce8(reg_a, shared_b, reg_c.hi, 16);
|
||||
}
|
||||
|
||||
if ((get_global_id(0) + block_id_m * TILESIZE_M) >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
__local uint out_idx[TILESIZE_N];
|
||||
|
||||
if (get_local_id(0) < TILESIZE_N) {
|
||||
uint idx = src2[block_id_n * TILESIZE_N + get_local_id(0)];
|
||||
if (idx == 0xFFFFFFFF) {
|
||||
idx = src2[block_id_n * TILESIZE_N + 0];
|
||||
}
|
||||
out_idx[get_local_id(0)] = idx * ne01;
|
||||
}
|
||||
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
uint m_offset = row + get_local_id(0);
|
||||
|
||||
write_imagef(dst, out_idx[1] + m_offset, (reg_c.s1));
|
||||
write_imagef(dst, out_idx[2] + m_offset, (reg_c.s2));
|
||||
write_imagef(dst, out_idx[3] + m_offset, (reg_c.s3));
|
||||
write_imagef(dst, out_idx[4] + m_offset, (reg_c.s4));
|
||||
write_imagef(dst, out_idx[5] + m_offset, (reg_c.s5));
|
||||
write_imagef(dst, out_idx[6] + m_offset, (reg_c.s6));
|
||||
write_imagef(dst, out_idx[7] + m_offset, (reg_c.s7));
|
||||
write_imagef(dst, out_idx[8] + m_offset, (reg_c.s8));
|
||||
write_imagef(dst, out_idx[9] + m_offset, (reg_c.s9));
|
||||
write_imagef(dst, out_idx[10] + m_offset, (reg_c.sa));
|
||||
write_imagef(dst, out_idx[11] + m_offset, (reg_c.sb));
|
||||
write_imagef(dst, out_idx[12] + m_offset, (reg_c.sc));
|
||||
write_imagef(dst, out_idx[13] + m_offset, (reg_c.sd));
|
||||
write_imagef(dst, out_idx[14] + m_offset, (reg_c.se));
|
||||
write_imagef(dst, out_idx[15] + m_offset, (reg_c.sf));
|
||||
write_imagef(dst, out_idx[16] + m_offset, (reg_c.sg));
|
||||
write_imagef(dst, out_idx[17] + m_offset, (reg_c.sh));
|
||||
write_imagef(dst, out_idx[18] + m_offset, (reg_c.si));
|
||||
write_imagef(dst, out_idx[19] + m_offset, (reg_c.sj));
|
||||
write_imagef(dst, out_idx[20] + m_offset, (reg_c.sk));
|
||||
write_imagef(dst, out_idx[21] + m_offset, (reg_c.sl));
|
||||
write_imagef(dst, out_idx[22] + m_offset, (reg_c.sm));
|
||||
write_imagef(dst, out_idx[23] + m_offset, (reg_c.sn));
|
||||
write_imagef(dst, out_idx[24] + m_offset, (reg_c.so));
|
||||
write_imagef(dst, out_idx[25] + m_offset, (reg_c.sp));
|
||||
write_imagef(dst, out_idx[26] + m_offset, (reg_c.sq));
|
||||
write_imagef(dst, out_idx[27] + m_offset, (reg_c.sr));
|
||||
write_imagef(dst, out_idx[28] + m_offset, (reg_c.ss));
|
||||
write_imagef(dst, out_idx[29] + m_offset, (reg_c.st));
|
||||
write_imagef(dst, out_idx[30] + m_offset, (reg_c.su));
|
||||
write_imagef(dst, out_idx[31] + m_offset, (reg_c.sv));
|
||||
|
||||
barrier(CLK_GLOBAL_MEM_FENCE);
|
||||
write_imagef(dst, out_idx[0] + m_offset, (reg_c.s0));
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
// Generic int8 dp4a MoE GEMM, specialized versions also exist
|
||||
// MOE_QT:
|
||||
// 4 (q4_K)/41(q4_1)/40(q4_0) NIBBLE image low nibbles -> EXP4
|
||||
// 5 (q5_K)/51(q5_1)/50(q5_0) NIBBLE+HI image nibbles + qh high-bit plane
|
||||
// 6 (q6_K) Q6 image nibbles + qh 2-bit -> SIGN6((nibble|hi2))
|
||||
// 80(q8_0)/82(mxfp4) INT8 global int8 codes (mxfp4: convert applies kvalues LUT)
|
||||
|
||||
#define TILESIZE_M 64
|
||||
#define TILESIZE_N 32
|
||||
#define QK_K 256
|
||||
|
||||
#ifndef MOE_QT
|
||||
#define MOE_QT 4
|
||||
#endif
|
||||
|
||||
// 4 nibbles in low 16 bits of u -> 4 bytes (value 0..15)
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
// 4 2-bit highs in byte b -> 4 bytes, bits 4-5 (q6_K)
|
||||
#define EXP2(b) ( (((uint)((b) & 0x03u)) << 4) | \
|
||||
(((uint)((b) & 0x0Cu)) << 10) | \
|
||||
(((uint)((b) & 0x30u)) << 16) | \
|
||||
(((uint)((b) & 0xC0u)) << 22) )
|
||||
|
||||
// q6 (0..63) -> (q6-32) signed int8/byte (no inter-byte carry)
|
||||
inline uint SIGN6(uint q6p){ uint x=q6p^0x20202020u; uint s=x&0x20202020u; return x|(s<<1)|(s<<2); }
|
||||
|
||||
// 4 high bits (one per element, in bits 0..3 of h) -> bit4 of each of 4 bytes (5-bit hi)
|
||||
#define EXP1(h) ( (((uint)((h) & 0x1u)) << 4) | \
|
||||
(((uint)((h) & 0x2u)) << 11) | \
|
||||
(((uint)((h) & 0x4u)) << 18) | \
|
||||
(((uint)((h) & 0x8u)) << 25) )
|
||||
|
||||
// per-type weight params + per-32-step unpack into qw[8] (8 int8 uints)
|
||||
#if MOE_QT == 4 || MOE_QT == 41 || MOE_QT == 40
|
||||
#define WEIGHT_PARAMS __read_only image1d_buffer_t src0_q,
|
||||
#define LOAD_QW(step, sub) \
|
||||
uint qw[8]; { \
|
||||
const uint qoff0 = row + ((ne01*(step))>>3) + ((expert_id*ne00*ne01)>>3); \
|
||||
const uint qoff1 = row + ((ne01*((step)+16))>>3) + ((expert_id*ne00*ne01)>>3); \
|
||||
const uint r0=read_imageui(src0_q,qoff0+lid).x, r1=read_imageui(src0_q,qoff0+lid+ne01).x; \
|
||||
const uint r2=read_imageui(src0_q,qoff1+lid).x, r3=read_imageui(src0_q,qoff1+lid+ne01).x; \
|
||||
qw[0]=EXP4(r0); qw[1]=EXP4(r0>>16); qw[2]=EXP4(r1); qw[3]=EXP4(r1>>16); \
|
||||
qw[4]=EXP4(r2); qw[5]=EXP4(r2>>16); qw[6]=EXP4(r3); qw[7]=EXP4(r3>>16); }
|
||||
|
||||
#elif MOE_QT == 5 || MOE_QT == 51 || MOE_QT == 50
|
||||
// low nibbles via image (q4_K layout) + high-bit plane src0_qh: 1 uint per 32-block
|
||||
// (bit i = high bit of element i). qh laid out [expert][block][row] to match the
|
||||
// existing q5_0 trans4 convert
|
||||
#define WEIGHT_PARAMS __read_only image1d_buffer_t src0_q, __global uint * src0_qh,
|
||||
#define LOAD_QW(step, sub) \
|
||||
uint qw[8]; { \
|
||||
const uint qoff0 = row + ((ne01*(step))>>3) + ((expert_id*ne00*ne01)>>3); \
|
||||
const uint qoff1 = row + ((ne01*((step)+16))>>3) + ((expert_id*ne00*ne01)>>3); \
|
||||
const uint r0=read_imageui(src0_q,qoff0+lid).x, r1=read_imageui(src0_q,qoff0+lid+ne01).x; \
|
||||
const uint r2=read_imageui(src0_q,qoff1+lid).x, r3=read_imageui(src0_q,qoff1+lid+ne01).x; \
|
||||
const uint h = src0_qh[row_idx + (sub)*ne01 + expert_id*(ne00>>5)*ne01]; \
|
||||
qw[0]=EXP4(r0)|EXP1(h); qw[1]=EXP4(r0>>16)|EXP1(h>>4); \
|
||||
qw[2]=EXP4(r1)|EXP1(h>>8); qw[3]=EXP4(r1>>16)|EXP1(h>>12); \
|
||||
qw[4]=EXP4(r2)|EXP1(h>>16); qw[5]=EXP4(r2>>16)|EXP1(h>>20); \
|
||||
qw[6]=EXP4(r3)|EXP1(h>>24); qw[7]=EXP4(r3>>16)|EXP1(h>>28); }
|
||||
|
||||
#elif MOE_QT == 6
|
||||
#define WEIGHT_PARAMS __read_only image1d_buffer_t src0_ql, __global uint * src0_qh,
|
||||
#define LOAD_QW(step, sub) \
|
||||
uint qw[8]; { \
|
||||
const uint qoff0 = row + ((ne01*(step))>>3) + ((expert_id*ne00*ne01)>>3); \
|
||||
const uint qoff1 = row + ((ne01*((step)+16))>>3) + ((expert_id*ne00*ne01)>>3); \
|
||||
const uint r0=read_imageui(src0_ql,qoff0+lid).x, r1=read_imageui(src0_ql,qoff0+lid+ne01).x; \
|
||||
const uint r2=read_imageui(src0_ql,qoff1+lid).x, r3=read_imageui(src0_ql,qoff1+lid+ne01).x; \
|
||||
const uint qhb = row + ((sub)*2)*ne01 + expert_id*((ne00>>5)*2)*ne01 + lid; \
|
||||
const uint qh1=src0_qh[qhb], qh2=src0_qh[qhb+ne01]; \
|
||||
qw[0]=SIGN6(EXP4(r0)|EXP2(qh1&0xFFu)); qw[1]=SIGN6(EXP4(r0>>16)|EXP2((qh1>>8)&0xFFu)); \
|
||||
qw[2]=SIGN6(EXP4(r1)|EXP2((qh1>>16)&0xFFu)); qw[3]=SIGN6(EXP4(r1>>16)|EXP2((qh1>>24)&0xFFu)); \
|
||||
qw[4]=SIGN6(EXP4(r2)|EXP2(qh2&0xFFu)); qw[5]=SIGN6(EXP4(r2>>16)|EXP2((qh2>>8)&0xFFu)); \
|
||||
qw[6]=SIGN6(EXP4(r3)|EXP2((qh2>>16)&0xFFu)); qw[7]=SIGN6(EXP4(r3>>16)|EXP2((qh2>>24)&0xFFu)); }
|
||||
|
||||
#elif MOE_QT == 80 || MOE_QT == 82
|
||||
// 8-bit direct: int8 codes 8 uints / 32-block, [expert][row][8*sub]. mxfp4: the
|
||||
// convert resolves kvalues_mxfp4[nibble] -> int8 and stores the e8m0_half scale.
|
||||
#define WEIGHT_PARAMS __global uint * src0_q8,
|
||||
#define LOAD_QW(step, sub) \
|
||||
uint qw[8]; { \
|
||||
const uint qb = (expert_id*ne01 + row_idx)*(ne00>>2) + (sub)*8; \
|
||||
qw[0]=src0_q8[qb+0]; qw[1]=src0_q8[qb+1]; qw[2]=src0_q8[qb+2]; qw[3]=src0_q8[qb+3]; \
|
||||
qw[4]=src0_q8[qb+4]; qw[5]=src0_q8[qb+5]; qw[6]=src0_q8[qb+6]; qw[7]=src0_q8[qb+7]; }
|
||||
#else
|
||||
#error "unknown MOE_QT"
|
||||
#endif
|
||||
|
||||
inline int dp4a4(uint w0,uint w1,uint w2,uint w3,uint a0,uint a1,uint a2,uint a3){
|
||||
int r=0; r=dot_acc_sat_4x8packed_ss_int(w0,a0,r); r=dot_acc_sat_4x8packed_ss_int(w1,a1,r);
|
||||
r=dot_acc_sat_4x8packed_ss_int(w2,a2,r); r=dot_acc_sat_4x8packed_ss_int(w3,a3,r); return r; }
|
||||
|
||||
// One token's two-half dp4a + uniform scale/min epilogue into acc[t].
|
||||
#define MOE_DP4A_T(t) do { \
|
||||
const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], sh_qa[t][0],sh_qa[t][1],sh_qa[t][2],sh_qa[t][3]); \
|
||||
const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], sh_qa[t][4],sh_qa[t][5],sh_qa[t][6],sh_qa[t][7]); \
|
||||
const float a_d = (float)sh_d[t]; \
|
||||
acc[t] += sc0*a_d*(float)raw1 + sc1*a_d*(float)raw2 - mn*(float)sh_s[t]; \
|
||||
} while (0)
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_moe_q8_1_dp4a(
|
||||
WEIGHT_PARAMS // per-type native weight buffer(s)
|
||||
__global half * src0_scale,// uniform f16 16/superblock (per-16), [expert,row]
|
||||
__global half * src0_min, // uniform f16 8/superblock (per-32), [expert,row]
|
||||
__global uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem)
|
||||
__global half * src1_da, // q8_1 per-block scale [tok_slot * ne00/32]
|
||||
__global half * src1_sa, // q8_1 per-block sum*d [tok_slot * ne00/32]
|
||||
__global uint * src2, // post-router (orig out positions)
|
||||
__global ushort * src2_emap, // tile -> expert id
|
||||
__write_only image1d_buffer_t dst,
|
||||
__global int * total_tiles,
|
||||
uint ne00,
|
||||
uint ne01,
|
||||
int is_ragged,
|
||||
int has_min // 0 for symmetric types (q8_0/q6_K/q4_0/...): skip min read
|
||||
) {
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
if (block_id_n >= total_tiles[0]) return;
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> output row within M-tile
|
||||
const ushort expert_id = src2_emap[block_id_n];
|
||||
const uint row = block_id_m * TILESIZE_M;
|
||||
const uint col = block_id_n * TILESIZE_N;
|
||||
const uint row_idx = row + lid;
|
||||
|
||||
// Scale/min are laid out FLAT per-32-block (2 per-16-segment scales + 1 min per
|
||||
// 32-block), so K only needs to be a multiple of 32 — works for the 32-block
|
||||
// types (q8_0/q5_0/q4_0/...) as well as the K-quants (K%256==0, same bytes).
|
||||
const uint nblk32 = ne00 / 32;
|
||||
const uint sc_per_row = nblk32 * 2;
|
||||
const uint mn_per_row = nblk32;
|
||||
const uint ne00_u = ne00 >> 2;
|
||||
const uint ne00_b = ne00 >> 5;
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
__local uint sh_src2[TILESIZE_N];
|
||||
__local int sh_nreal;
|
||||
if (lid < TILESIZE_N) sh_src2[lid] = src2[col + lid];
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (lid == 0) {
|
||||
int nr = TILESIZE_N;
|
||||
if (is_ragged) { nr = 0;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) if (sh_src2[t] != 0xFFFFFFFFu) ++nr; }
|
||||
sh_nreal = nr;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
const int n_real = sh_nreal;
|
||||
|
||||
float acc[TILESIZE_N];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
|
||||
|
||||
for (uint step = 0; step < ne00; step += 32) {
|
||||
const uint sub = step >> 5; // 32-block index along K
|
||||
|
||||
// uniform pre-decoded scale (2 per-16-seg) + min (1) for this row, this 32-block
|
||||
__global half * scl = src0_scale + (expert_id*ne01 + row_idx)*sc_per_row + sub*2;
|
||||
const float sc0 = (float)scl[0];
|
||||
const float sc1 = (float)scl[1];
|
||||
float mn = 0.0f;
|
||||
if (has_min) mn = (float)src0_min[(expert_id*ne01 + row_idx)*mn_per_row + sub];
|
||||
|
||||
LOAD_QW(step, sub)
|
||||
|
||||
const uint stage_lim = (uint)n_real * 8;
|
||||
for (uint idx = lid; idx < stage_lim; idx += 64) {
|
||||
const uint t = idx >> 3, u = idx & 7;
|
||||
sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
|
||||
}
|
||||
if (lid < (uint)n_real) {
|
||||
sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];
|
||||
sh_s[lid] = src1_sa[(col + lid) * ne00_b + sub];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 0; t < TILESIZE_N; ++t) { MOE_DP4A_T(t); }
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int t = 0; t < n_real; ++t) { MOE_DP4A_T(t); }
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (row_idx >= ne01) return;
|
||||
|
||||
__local uint out_idx[TILESIZE_N];
|
||||
if (lid < TILESIZE_N) {
|
||||
uint idx = sh_src2[lid];
|
||||
if (idx == 0xFFFFFFFF) idx = sh_src2[0];
|
||||
out_idx[lid] = idx * ne01;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const uint m_offset = row + lid;
|
||||
if (n_real == TILESIZE_N) {
|
||||
#pragma unroll
|
||||
for (int t = 1; t < TILESIZE_N; ++t) write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
barrier(CLK_GLOBAL_MEM_FENCE);
|
||||
write_imagef(dst, out_idx[0] + m_offset, acc[0]);
|
||||
} else {
|
||||
for (int t = 0; t < n_real; ++t) write_imagef(dst, out_idx[t] + m_offset, acc[t]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
// Weight layout, feature-major:
|
||||
// src0_q[row + (k/4)*m] ushort = 4 nibbles (K = 4*grp .. +3)
|
||||
// src0_d[row + (k/32)*m] half = per-32-block scale
|
||||
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// IQ4_NL non-linear codebook as signed int8, packed 4 codes per uint.
|
||||
// divergent nibble lookups read a small __constant uint array + shift,
|
||||
// never a byte array because byte-indexed __constant loads serialize on Adreno and tank perf
|
||||
// idx 0-3: -127,-104,-83,-65 = 0x81,0x98,0xAD,0xBF
|
||||
// idx 4-7: -49,-35,-22,-10 = 0xCF,0xDD,0xEA,0xF6
|
||||
// idx 8-11: 1, 13, 25, 38 = 0x01,0x0D,0x19,0x26
|
||||
// idx 12-15: 53, 69, 89,113 = 0x35,0x45,0x59,0x71
|
||||
__constant uint kvalues_iq4nl_i8x4[4] = {
|
||||
0xBFAD9881u, 0xF6EADDCFu, 0x26190D01u, 0x71594535u
|
||||
};
|
||||
|
||||
// nibble (0..15) -> its codebook byte in the low 8 bits.
|
||||
inline uint iq4nl_code(uint n) {
|
||||
return (kvalues_iq4nl_i8x4[n >> 2] >> ((n & 3u) * 8u)) & 0xFFu;
|
||||
}
|
||||
|
||||
// 4 nibbles in low 16 bits of u -> 4 codebook int8, packed for dp4a.
|
||||
inline uint iq4nl_pack(ushort u) {
|
||||
return iq4nl_code((uint)( u & 0xF))
|
||||
| (iq4nl_code((uint)((u >> 4) & 0xF)) << 8)
|
||||
| (iq4nl_code((uint)((u >> 8) & 0xF)) << 16)
|
||||
| (iq4nl_code((uint)((u >> 12) & 0xF)) << 24);
|
||||
}
|
||||
|
||||
inline int dot8_q8a(uint8 qw, __local const uint * a) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a(
|
||||
__global const ushort * src0_q, // IQ4_NL nibbles (4/ushort, feature-major)
|
||||
__global const half * src0_d, // per-32-block scale, feature-major
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k // K (== ne00)
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
|
||||
const float d_w = (float)src0_d[rrow + sub * (uint)m];
|
||||
|
||||
// 8 weight uints (32 codebook int8) for this row, this 32-block.
|
||||
const uint qsbase = rrow + (step >> 2) * (uint)m;
|
||||
uint8 qw;
|
||||
qw.s0 = iq4nl_pack(src0_q[qsbase + 0 * m]);
|
||||
qw.s1 = iq4nl_pack(src0_q[qsbase + 1 * m]);
|
||||
qw.s2 = iq4nl_pack(src0_q[qsbase + 2 * m]);
|
||||
qw.s3 = iq4nl_pack(src0_q[qsbase + 3 * m]);
|
||||
qw.s4 = iq4nl_pack(src0_q[qsbase + 4 * m]);
|
||||
qw.s5 = iq4nl_pack(src0_q[qsbase + 5 * m]);
|
||||
qw.s6 = iq4nl_pack(src0_q[qsbase + 6 * m]);
|
||||
qw.s7 = iq4nl_pack(src0_q[qsbase + 7 * m]);
|
||||
|
||||
// cooperatively stage the 32-token x 32-K int8 activations to lm
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += d_w * LD4(sh_d, b) * rf;
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
// dst is [token, feature] row-major (stride m): dst[col*m + row].
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// Expand the 4 nibbles in the low 16 bits of u into 4 bytes (value 0..15),
|
||||
// packed for the int8 dp4a. The -8 zero-point is applied via the sum term.
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
inline int dot8_q8a(uint8 qw, __local const uint * a) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q4_0_q8_1_dp4a(
|
||||
__global const ushort * src0_q, // q4_0 nibbles (4/ushort, feature-major)
|
||||
__global const half * src0_d, // per-32-block scale, feature-major
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global const half * src1_sa, // q8_1 per-block sum*d [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k // K (== ne00)
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
|
||||
const float d_w = (float)src0_d[rrow + sub * (uint)m];
|
||||
|
||||
// 8 weight uints (32 nibbles) for this row, this 32-block. Feature-major:
|
||||
// src0_q[row + (k/4 + u)*m], k/4 = step/4 (= step>>2). EXP4 -> dp4a int8.
|
||||
const uint qsbase = rrow + (step >> 2) * (uint)m;
|
||||
uint8 qw;
|
||||
qw.s0 = EXP4(src0_q[qsbase + 0 * m]);
|
||||
qw.s1 = EXP4(src0_q[qsbase + 1 * m]);
|
||||
qw.s2 = EXP4(src0_q[qsbase + 2 * m]);
|
||||
qw.s3 = EXP4(src0_q[qsbase + 3 * m]);
|
||||
qw.s4 = EXP4(src0_q[qsbase + 4 * m]);
|
||||
qw.s5 = EXP4(src0_q[qsbase + 5 * m]);
|
||||
qw.s6 = EXP4(src0_q[qsbase + 6 * m]);
|
||||
qw.s7 = EXP4(src0_q[qsbase + 7 * m]);
|
||||
|
||||
// cooperatively stage the 32-token x 32-K int8 activations to LDS
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
// q4_0: w = d*(q-8) -> d_w * (a_d * dp4a(q,qa) - 8 * a_s)
|
||||
acc[g] += d_w * (LD4(sh_d, b) * rf - 8.0f * LD4(sh_s, b));
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
// dst is [token, feature] row-major (stride m): dst[col*m + row].
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#ifndef TILESIZE_N
|
||||
#define TILESIZE_N 32
|
||||
#endif
|
||||
#define QK_K 256
|
||||
#define K_SCALE_SIZE 12
|
||||
|
||||
inline void get_scale_min_k4(
|
||||
int j,
|
||||
global const uchar * q,
|
||||
uchar * d,
|
||||
uchar * m,
|
||||
uchar mask_d6,
|
||||
uchar mask_d4,
|
||||
uchar mask_hi2
|
||||
) {
|
||||
if (j < 4) {
|
||||
*d = q[j] & mask_d6;
|
||||
*m = q[j+4] & mask_d6;
|
||||
} else {
|
||||
*d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
|
||||
*m = ((q[j+4] >> 4) & mask_d4) | ((q[j] & mask_hi2) >> 2);
|
||||
}
|
||||
}
|
||||
|
||||
// Expand the 4 nibbles in the low 16 bits of `u` into 4 bytes (one nibble per
|
||||
// byte, value 0..15), packed for the int8 dp4a.
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
// 32-K dp4a dot of one token's int8 activations (8 packed uints in lm) against the
|
||||
// row's 8 packed weight uints. qw passed by value as a uint8 (register), not an array.
|
||||
inline int dot8_q8a(uint8 qw, __local const uint * a) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a(
|
||||
__global const ushort * src0_q, // q4_K weights (noshuffle, packed nibbles)
|
||||
__global const uchar * src0_s, // 6-bit scale/min codes
|
||||
__global const half * src0_d, // per-superblock scale
|
||||
__global const half * src0_dm, // per-superblock min
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global const half * src1_sa, // q8_1 per-block sum*d [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k, // K (== ne00)
|
||||
uchar mask_d6,
|
||||
uchar mask_d4,
|
||||
uchar mask_hi2
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
const uint num_superblocks = (uint)k / QK_K;
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
// One float4 vector-register accumulator per group of 4 tokens (NGROUPS = TILESIZE_N/4).
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) { acc[g] = (float4)(0.0f); }
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
const uint sb_idx = step / QK_K;
|
||||
const uint sub_idx = sub & 7;
|
||||
|
||||
// weight scale/min for this WI's row, this subblock
|
||||
const float dd = (float)src0_d [rrow + sb_idx * m];
|
||||
const float dmm = (float)src0_dm[rrow + sb_idx * m];
|
||||
global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
|
||||
uchar sv, mn;
|
||||
get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
|
||||
const float scale = dd * (float)sv;
|
||||
const float minv = dmm * (float)mn;
|
||||
|
||||
// repack this row's 32 weight nibbles into 8 dp4a uints. The packed q4_K
|
||||
// layout stores one ushort = 4 consecutive-K nibbles for a row at
|
||||
// src0_q[row + (K_group)*m], K_group = step/4 + u.
|
||||
const uint wbase = rrow + (step >> 2) * (uint)m;
|
||||
uint8 qw;
|
||||
qw.s0 = EXP4(src0_q[wbase + 0 * m]);
|
||||
qw.s1 = EXP4(src0_q[wbase + 1 * m]);
|
||||
qw.s2 = EXP4(src0_q[wbase + 2 * m]);
|
||||
qw.s3 = EXP4(src0_q[wbase + 3 * m]);
|
||||
qw.s4 = EXP4(src0_q[wbase + 4 * m]);
|
||||
qw.s5 = EXP4(src0_q[wbase + 5 * m]);
|
||||
qw.s6 = EXP4(src0_q[wbase + 6 * m]);
|
||||
qw.s7 = EXP4(src0_q[wbase + 7 * m]);
|
||||
|
||||
// cooperatively stage the 32-token x 32-K int8 activations to lm
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
// dst is [token, feature] row-major (stride m): dst[col*m + row]. Scatter each
|
||||
// lane with a per-token padding guard (dst is non-contiguous in token).
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg(
|
||||
__read_only image1d_buffer_t src0_q_img, // q4_K weights as uint32 texels (2 ushorts/texel)
|
||||
__global const uchar * src0_s, // 6-bit scale/min codes
|
||||
__global const half * src0_d, // per-superblock scale
|
||||
__global const half * src0_dm, // per-superblock min
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global const half * src1_sa, // q8_1 per-block sum*d [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k, // K (== ne00)
|
||||
uchar mask_d6,
|
||||
uchar mask_d4,
|
||||
uchar mask_hi2
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
// Constant per WI: the ushort the row needs always sits in the same half of
|
||||
// its uint32 texel (m even => index parity == rrow parity). Hoist the shift.
|
||||
const uint sel = (rrow & 1u) * 16u;
|
||||
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
const uint num_superblocks = (uint)k / QK_K;
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
const uint sb_idx = step / QK_K;
|
||||
const uint sub_idx = sub & 7;
|
||||
|
||||
const float dd = (float)src0_d [rrow + sb_idx * m];
|
||||
const float dmm = (float)src0_dm[rrow + sb_idx * m];
|
||||
global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
|
||||
uchar sv, mn;
|
||||
get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
|
||||
const float scale = dd * (float)sv;
|
||||
const float minv = dmm * (float)mn;
|
||||
|
||||
const uint wbase = rrow + (step >> 2) * (uint)m;
|
||||
uint8 qw;
|
||||
qw.s0 = EXP4(read_imageui(src0_q_img, (int)((wbase + 0 * m) >> 1)).x >> sel);
|
||||
qw.s1 = EXP4(read_imageui(src0_q_img, (int)((wbase + 1 * m) >> 1)).x >> sel);
|
||||
qw.s2 = EXP4(read_imageui(src0_q_img, (int)((wbase + 2 * m) >> 1)).x >> sel);
|
||||
qw.s3 = EXP4(read_imageui(src0_q_img, (int)((wbase + 3 * m) >> 1)).x >> sel);
|
||||
qw.s4 = EXP4(read_imageui(src0_q_img, (int)((wbase + 4 * m) >> 1)).x >> sel);
|
||||
qw.s5 = EXP4(read_imageui(src0_q_img, (int)((wbase + 5 * m) >> 1)).x >> sel);
|
||||
qw.s6 = EXP4(read_imageui(src0_q_img, (int)((wbase + 6 * m) >> 1)).x >> sel);
|
||||
qw.s7 = EXP4(read_imageui(src0_q_img, (int)((wbase + 7 * m) >> 1)).x >> sel);
|
||||
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -0,0 +1,235 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
// Weight layout
|
||||
// src0_qs[row + (k/4)*m] ushort = 4 low nibbles (K = 4*grp .. +3)
|
||||
// src0_qh[row + (k/8)*m] uchar = 8 high bits (one per element)
|
||||
// src0_d [row + (k/32)*m] half = per-32-block scale
|
||||
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// 4 nibbles in low 16 bits of u -> 4 bytes (value 0..15)
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
// 4 high bits (one per element, in bits 0..3 of h) -> bit4 of each of 4 bytes
|
||||
#define EXP1(h) ( (((uint)((h) & 0x1u)) << 4) | \
|
||||
(((uint)((h) & 0x2u)) << 11) | \
|
||||
(((uint)((h) & 0x4u)) << 18) | \
|
||||
(((uint)((h) & 0x8u)) << 25) )
|
||||
|
||||
inline int dot8_q8a(uint8 qw, __local const uint * a) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q5_0_q8_1_dp4a(
|
||||
__global const ushort * src0_qs, // q5_0 low nibbles (4/ushort, feature-major)
|
||||
__global const uchar * src0_qh, // q5_0 high-bit plane (8/uchar, feature-major)
|
||||
__global const half * src0_d, // per-32-block scale, feature-major
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global const half * src1_sa, // q8_1 per-block sum*d [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k // K (== ne00)
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
|
||||
const float d_w = (float)src0_d[rrow + sub * (uint)m];
|
||||
const float minv = d_w * 16.0f; // -16 centering -> subtract via q8_1 sum
|
||||
|
||||
// 8 weight uints (32 elements) for this row, this 32-block.
|
||||
// nibbles: src0_qs[row + (step/4 + u)*m]; high bits: src0_qh[row + (step/8 + u/2)*m],
|
||||
// 4-bit group selected by (u&1)*4.
|
||||
const uint qsbase = rrow + (step >> 2) * (uint)m;
|
||||
const uint qhbase = rrow + (step >> 3) * (uint)m;
|
||||
uint8 qw;
|
||||
#define QW(u) (EXP4(src0_qs[qsbase + (u) * m]) | \
|
||||
EXP1((uint)(src0_qh[qhbase + ((u) >> 1) * m] >> (((u) & 1u) * 4u)) & 0xFu))
|
||||
qw.s0 = QW(0); qw.s1 = QW(1); qw.s2 = QW(2); qw.s3 = QW(3);
|
||||
qw.s4 = QW(4); qw.s5 = QW(5); qw.s6 = QW(6); qw.s7 = QW(7);
|
||||
#undef QW
|
||||
|
||||
// cooperatively stage the 32-token x 32-K int8 activations to lm
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += d_w * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q5_0_q8_1_dp4a_wimg(
|
||||
__read_only image1d_buffer_t src0_qs_img, // q5_0 low nibbles as uint32 texels (2 ushorts/texel)
|
||||
__global const uchar * src0_qh,
|
||||
__global const half * src0_d,
|
||||
__global const uint * src1_qa,
|
||||
__global const half * src1_da,
|
||||
__global const half * src1_sa,
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m,
|
||||
int n_no_padding,
|
||||
int k
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0);
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0;
|
||||
|
||||
const uint sel = (rrow & 1u) * 16u; // constant per WI: qs ushort half in its uint32 texel
|
||||
|
||||
const uint k_u = (uint)k >> 2;
|
||||
const uint k_b = (uint)k >> 5;
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
|
||||
const float d_w = (float)src0_d[rrow + sub * (uint)m];
|
||||
const float minv = d_w * 16.0f;
|
||||
|
||||
const uint qsbase = rrow + (step >> 2) * (uint)m; // ushort index
|
||||
const uint qhbase = rrow + (step >> 3) * (uint)m;
|
||||
uint8 qw;
|
||||
// qs ushort via texture: uint32 texel = ushort_index>>1, half = sel.
|
||||
#define QSU(u) ((read_imageui(src0_qs_img, (int)((qsbase + (u) * m) >> 1)).x >> sel) & 0xFFFFu)
|
||||
#define QW(u) (EXP4(QSU(u)) | \
|
||||
EXP1((uint)(src0_qh[qhbase + ((u) >> 1) * m] >> (((u) & 1u) * 4u)) & 0xFu))
|
||||
qw.s0 = QW(0); qw.s1 = QW(1); qw.s2 = QW(2); qw.s3 = QW(3);
|
||||
qw.s4 = QW(4); qw.s5 = QW(5); qw.s6 = QW(6); qw.s7 = QW(7);
|
||||
#undef QW
|
||||
#undef QSU
|
||||
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += d_w * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#define TILESIZE_N 32
|
||||
#define QK_K 256
|
||||
#define K_SCALE_SIZE 12
|
||||
|
||||
inline void get_scale_min_k4(
|
||||
int j,
|
||||
global const uchar * q,
|
||||
uchar * d,
|
||||
uchar * m,
|
||||
uchar mask_d6,
|
||||
uchar mask_d4,
|
||||
uchar mask_hi2
|
||||
) {
|
||||
if (j < 4) {
|
||||
*d = q[j] & mask_d6;
|
||||
*m = q[j+4] & mask_d6;
|
||||
} else {
|
||||
*d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
|
||||
*m = ((q[j+4] >> 4) & mask_d4) | ((q[j] & mask_hi2) >> 2);
|
||||
}
|
||||
}
|
||||
|
||||
// 4 nibbles in the low 16 bits of `u` -> 4 bytes (value 0..15, bits 0-3).
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
// 4 high bits (one per element, in bits 0-3 of h) -> bit 4 of each of 4 bytes,
|
||||
// so OR with EXP4 forms the 5-bit q5_K code 0..31.
|
||||
#define EXP1(h) ( (((uint)((h) & 0x1u)) << 4) | \
|
||||
(((uint)((h) & 0x2u)) << 11) | \
|
||||
(((uint)((h) & 0x4u)) << 18) | \
|
||||
(((uint)((h) & 0x8u)) << 25) )
|
||||
|
||||
inline int dot8_q8a(uint8 qw, __local const uint * a) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q5_k_q8_1_dp4a(
|
||||
__global const ushort * src0_q, // q5_K low nibbles (transposed, ushort = 4 nibbles)
|
||||
__global const uchar * src0_qh, // q5_K high bits (transposed, uchar = 8 elems/byte)
|
||||
__global const uchar * src0_s, // 6-bit scale/min codes [row][superblock][12]
|
||||
__global const half * src0_d, // per-superblock scale (transposed)
|
||||
__global const half * src0_dm, // per-superblock min (transposed)
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global const half * src1_sa, // q8_1 per-block sum*d [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k, // K (== ne00)
|
||||
uchar mask_d6,
|
||||
uchar mask_d4,
|
||||
uchar mask_hi2
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0;
|
||||
|
||||
const uint num_superblocks = (uint)k / QK_K;
|
||||
const uint k_u = (uint)k >> 2;
|
||||
const uint k_b = (uint)k >> 5;
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
__local half sh_s[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
const uint sb_idx = step / QK_K;
|
||||
const uint sub_idx = sub & 7;
|
||||
|
||||
const float dd = (float)src0_d [rrow + sb_idx * m];
|
||||
const float dmm = (float)src0_dm[rrow + sb_idx * m];
|
||||
global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
|
||||
uchar sv, mn;
|
||||
get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
|
||||
const float scale = dd * (float)sv;
|
||||
const float minv = dmm * (float)mn;
|
||||
|
||||
// repack this row's 32 weights (nibble | high-bit) into 8 dp4a uints.
|
||||
// ushort u -> 4 elements at K = step + u*4; its 4 high bits are nibble
|
||||
// (u&1) of qh byte (step/8 + u/2).
|
||||
const uint wbase = rrow + (step >> 2) * (uint)m;
|
||||
const uint qhbase = rrow + (step >> 3) * (uint)m;
|
||||
uint8 qw;
|
||||
#define QWU(u) ( EXP4((uint)src0_q[wbase + (uint)(u) * m]) \
|
||||
| EXP1( (uint)((src0_qh[qhbase + (uint)((u) >> 1) * m] >> (((u) & 1) * 4)) & 0x0Fu) ) )
|
||||
qw.s0 = QWU(0); qw.s1 = QWU(1); qw.s2 = QWU(2); qw.s3 = QWU(3);
|
||||
qw.s4 = QWU(4); qw.s5 = QWU(5); qw.s6 = QWU(6); qw.s7 = QWU(7);
|
||||
#undef QWU
|
||||
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
#define TILESIZE_N 32
|
||||
#define QK_K 256
|
||||
|
||||
// 4 nibbles in the low 16 bits of `u` -> 4 bytes (value 0..15, in bits 0-3).
|
||||
#define EXP4(u) ( ((uint)((u) & 0x000Fu)) | \
|
||||
(((uint)((u) & 0x00F0u)) << 4) | \
|
||||
(((uint)((u) & 0x0F00u)) << 8) | \
|
||||
(((uint)((u) & 0xF000u)) << 12) )
|
||||
|
||||
// 4 2-bit highs in byte `b` -> 4 bytes, value 0..3 in bits 4-5 (pre-multiplied
|
||||
// by 16 so it ORs with the EXP4 nibble to form q6 in 0..63).
|
||||
#define EXP2(b) ( (((uint)((b) & 0x03u)) << 4) | \
|
||||
(((uint)((b) & 0x0Cu)) << 10) | \
|
||||
(((uint)((b) & 0x30u)) << 16) | \
|
||||
(((uint)((b) & 0xC0u)) << 22) )
|
||||
|
||||
// q6 (0..63, bits 0-5 of each byte) -> (q6-32) as a signed int8 per byte.
|
||||
inline uint SIGN6(uint q6p) {
|
||||
uint x = q6p ^ 0x20202020u;
|
||||
uint s = x & 0x20202020u;
|
||||
return x | (s << 1) | (s << 2);
|
||||
}
|
||||
|
||||
// 16-K dp4a dot: 4 packed weight uints against 4 packed int8 activation uints.
|
||||
inline int dot4_q8a(uint w0, uint w1, uint w2, uint w3,
|
||||
uint a0, uint a1, uint a2, uint a3) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(w0, a0, r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(w1, a1, r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(w2, a2, r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(w3, a3, r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q6_k_q8_1_dp4a(
|
||||
__global const ushort * src0_ql, // q6_K low nibbles (noshuffle)
|
||||
__global const uchar * src0_qh, // q6_K high 2-bit (uchar, 4 highs/elem)
|
||||
__global const ushort * src0_s, // int8 scale codes (2 chars/ushort, per 16)
|
||||
__global const half * src0_d, // per-superblock scale
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k // K (== ne00)
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5; // 32-block index along K
|
||||
const uint sb_idx = step / QK_K; // superblock index
|
||||
|
||||
// q6_K superblock scale + the two int8 sub-scales spanning this 32-block
|
||||
const float dd = (float)src0_d[rrow + sb_idx * m];
|
||||
const char2 sc = as_char2(src0_s[rrow + sub * m]);
|
||||
const float scale0 = dd * (float)sc.s0; // K step..step+15
|
||||
const float scale1 = dd * (float)sc.s1; // K step+16..step+31
|
||||
|
||||
// repack this row's 32 weights into 8 dp4a uints (4 K each). ql ushort +
|
||||
// qh uchar are co-located at src0_*[row + (step/4 + u)*m].
|
||||
const uint wbase = rrow + (step >> 2) * (uint)m;
|
||||
uint qw[8];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 8; ++u) {
|
||||
const uint o = wbase + (uint)u * (uint)m;
|
||||
qw[u] = SIGN6(EXP4((uint)src0_ql[o]) | EXP2((uint)src0_qh[o]));
|
||||
}
|
||||
|
||||
// cooperatively stage the 32-token x 32-K int8 activations + scale
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
#define DOT_TOK(j) { \
|
||||
__local const uint * a = sh_qa[b + (j)]; \
|
||||
const int raw1 = dot4_q8a(qw[0], qw[1], qw[2], qw[3], a[0], a[1], a[2], a[3]); \
|
||||
const int raw2 = dot4_q8a(qw[4], qw[5], qw[6], qw[7], a[4], a[5], a[6], a[7]); \
|
||||
rf.s##j = scale0 * (float)raw1 + scale1 * (float)raw2; \
|
||||
}
|
||||
DOT_TOK(0); DOT_TOK(1); DOT_TOK(2); DOT_TOK(3);
|
||||
#undef DOT_TOK
|
||||
const float4 ad = (float4)((float)sh_d[b+0], (float)sh_d[b+1], (float)sh_d[b+2], (float)sh_d[b+3]);
|
||||
acc[g] += ad * rf;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
// dst is [token, feature] row-major (stride m): dst[col*m + row].
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||
#ifdef cl_khr_integer_dot_product
|
||||
#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
|
||||
#endif
|
||||
|
||||
// ne1<=8 keeps the f16 / bin small-batch path.
|
||||
|
||||
#define TILESIZE_N 32
|
||||
|
||||
// 32-K dp4a dot of one token's int8 activations (8 packed uints in lm) against
|
||||
// 8 packed weight uints. q8_0 weights are already dp4a-format signed int8.
|
||||
inline int dot8_q8a(uint8 qw, __local const uint * a) {
|
||||
int r = 0;
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
|
||||
r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
|
||||
return r;
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a(
|
||||
__global const uint * src0_q, // q8_0 weights: signed int8, 4/uint, feature-major
|
||||
__global const half * src0_d, // per-32-block scale, feature-major [row + (k/32)*m]
|
||||
__global const uint * src1_qa, // q8_1 activations int8 (as uint, 4/elem) [N, K]
|
||||
__global const half * src1_da, // q8_1 per-block scale [N, K/32]
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m, // output features (rows)
|
||||
int n_no_padding, // tokens (cols)
|
||||
int k // K (== ne00)
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0); // 0..63 -> row within the M-tile
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0; // clamp OOB rows; their writes are masked
|
||||
|
||||
const uint k_u = (uint)k >> 2; // K in uint (int8x4) units
|
||||
const uint k_b = (uint)k >> 5; // blocks-of-32 along K
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
|
||||
const float d_w = (float)src0_d[rrow + sub * (uint)m];
|
||||
|
||||
// 8 weight uints (32 int8) for this row, this 32-block. Feature-major:
|
||||
// src0_q[row + (k/4 + u)*m], k/4 = step/4 (= step>>2).
|
||||
const uint wbase = rrow + (step >> 2) * (uint)m;
|
||||
uint8 qw;
|
||||
qw.s0 = src0_q[wbase + 0 * m];
|
||||
qw.s1 = src0_q[wbase + 1 * m];
|
||||
qw.s2 = src0_q[wbase + 2 * m];
|
||||
qw.s3 = src0_q[wbase + 3 * m];
|
||||
qw.s4 = src0_q[wbase + 4 * m];
|
||||
qw.s5 = src0_q[wbase + 5 * m];
|
||||
qw.s6 = src0_q[wbase + 6 * m];
|
||||
qw.s7 = src0_q[wbase + 7 * m];
|
||||
|
||||
// cooperatively stage the 32-token x 32-K int8 activations to LDS
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += d_w * LD4(sh_d, b) * rf;
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
// dst is [token, feature] row-major (stride m): dst[col*m + row].
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
|
||||
__attribute__((qcom_wave_pair_mode(1)))
|
||||
kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg(
|
||||
__read_only image1d_buffer_t src0_q_img, // q8_0 weights as uint32 texels (4 int8/texel)
|
||||
__global const half * src0_d,
|
||||
__global const uint * src1_qa,
|
||||
__global const half * src1_da,
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int m,
|
||||
int n_no_padding,
|
||||
int k
|
||||
) {
|
||||
dst = (global float *)((global char *)dst + offsetd);
|
||||
|
||||
const uint lid = get_local_id(0);
|
||||
const uint block_id_m = get_global_id(1);
|
||||
const uint block_id_n = get_global_id(2);
|
||||
|
||||
const uint row = block_id_m * 64 + lid;
|
||||
const uint col_base = block_id_n * TILESIZE_N;
|
||||
const bool row_valid = row < (uint)m;
|
||||
const uint rrow = row_valid ? row : 0;
|
||||
|
||||
const uint k_u = (uint)k >> 2;
|
||||
const uint k_b = (uint)k >> 5;
|
||||
|
||||
__local uint sh_qa[TILESIZE_N][8];
|
||||
__local half sh_d[TILESIZE_N];
|
||||
|
||||
#define NGROUPS (TILESIZE_N / 4)
|
||||
float4 acc[NGROUPS];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
|
||||
|
||||
for (uint step = 0; step < (uint)k; step += 32) {
|
||||
const uint sub = step >> 5;
|
||||
|
||||
const float d_w = (float)src0_d[rrow + sub * (uint)m];
|
||||
|
||||
const uint wbase = rrow + (step >> 2) * (uint)m;
|
||||
uint8 qw;
|
||||
qw.s0 = read_imageui(src0_q_img, (int)(wbase + 0 * m)).x;
|
||||
qw.s1 = read_imageui(src0_q_img, (int)(wbase + 1 * m)).x;
|
||||
qw.s2 = read_imageui(src0_q_img, (int)(wbase + 2 * m)).x;
|
||||
qw.s3 = read_imageui(src0_q_img, (int)(wbase + 3 * m)).x;
|
||||
qw.s4 = read_imageui(src0_q_img, (int)(wbase + 4 * m)).x;
|
||||
qw.s5 = read_imageui(src0_q_img, (int)(wbase + 5 * m)).x;
|
||||
qw.s6 = read_imageui(src0_q_img, (int)(wbase + 6 * m)).x;
|
||||
qw.s7 = read_imageui(src0_q_img, (int)(wbase + 7 * m)).x;
|
||||
|
||||
for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
|
||||
const uint t = idx >> 3;
|
||||
const uint u = idx & 7;
|
||||
const uint c = col_base + t;
|
||||
sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
|
||||
}
|
||||
if (lid < TILESIZE_N) {
|
||||
const uint c = col_base + lid;
|
||||
sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const int b = g * 4;
|
||||
float4 rf;
|
||||
rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]); rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
|
||||
rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]); rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
|
||||
acc[g] += d_w * LD4(sh_d, b) * rf;
|
||||
}
|
||||
#undef LD4
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
}
|
||||
|
||||
if (!row_valid) {
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int g = 0; g < NGROUPS; ++g) {
|
||||
const uint b = (uint)(g * 4);
|
||||
const float4 a = acc[g];
|
||||
const uint c0 = col_base + b;
|
||||
if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
|
||||
if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
|
||||
if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
|
||||
if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
|
||||
}
|
||||
#undef NGROUPS
|
||||
}
|
||||
@@ -163,3 +163,95 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns(
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
__attribute__((qcom_reqd_sub_group_size("half")))
|
||||
__kernel void kernel_gemv_moe_mxfp4_f32_ns_wimg(
|
||||
__read_only image1d_buffer_t src0_q,
|
||||
__global uchar * src0_e,
|
||||
__read_only image1d_buffer_t src1,
|
||||
__global uint * src2,
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne11
|
||||
) {
|
||||
uint i01 = get_global_id(0);
|
||||
uint i20 = get_global_id(2);
|
||||
uint sgid = get_local_id(1);
|
||||
uint slid = get_sub_group_local_id();
|
||||
|
||||
if (i01 >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
uint expert_offset = expert_id * ne00 * ne01 / 32;
|
||||
|
||||
__private float sum = 0.0f;
|
||||
|
||||
for (uint ib00 = sgid; ib00 < (ne00 / QK_MXFP4); ib00 += N_SIMDGROUP) {
|
||||
|
||||
uint4 regQ;
|
||||
uint block_offset = expert_offset * 4 + ib00 * ne01 * 4 + i01;
|
||||
|
||||
regQ.s0 = read_imageui(src0_q, (int)(block_offset)).x;
|
||||
regQ.s1 = read_imageui(src0_q, (int)(block_offset + ne01)).x;
|
||||
regQ.s2 = read_imageui(src0_q, (int)(block_offset + ne01 * 2)).x;
|
||||
regQ.s3 = read_imageui(src0_q, (int)(block_offset + ne01 * 3)).x;
|
||||
|
||||
uint offset = i11 * ne00 / 4 + ib00 * 8;
|
||||
|
||||
half8 fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s0));
|
||||
|
||||
float4 shared_y4;
|
||||
shared_y4 = read_imagef(src1, (offset + 0));
|
||||
float4 acc = shared_y4 * convert_float4(fp16x8.lo);
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 1));
|
||||
acc += shared_y4 * convert_float4(fp16x8.hi);
|
||||
|
||||
fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s1));
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 2));
|
||||
acc += shared_y4 * convert_float4(fp16x8.lo);
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 3));
|
||||
acc += shared_y4 * convert_float4(fp16x8.hi);
|
||||
|
||||
fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s2));
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 4));
|
||||
acc += shared_y4 * convert_float4(fp16x8.lo);
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 5));
|
||||
acc += shared_y4 * convert_float4(fp16x8.hi);
|
||||
|
||||
fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s3));
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 6));
|
||||
acc += shared_y4 * convert_float4(fp16x8.lo);
|
||||
|
||||
shared_y4 = read_imagef(src1, (offset + 7));
|
||||
acc += shared_y4 * convert_float4(fp16x8.hi);
|
||||
|
||||
uchar regE = src0_e[ib00 * ne01 + i01 + expert_offset];
|
||||
sum += e8m0_to_fp32(regE) * ((acc.s0 + acc.s1) + (acc.s2 + acc.s3));
|
||||
}
|
||||
|
||||
__local float reduceLM[SIMDGROUP_WIDTH * (N_SIMDGROUP - 1)];
|
||||
if (sgid == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = sum;
|
||||
if (sgid == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = sum;
|
||||
if (sgid == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = sum;
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
|
||||
if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
|
||||
if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
|
||||
|
||||
if (sgid == 0) {
|
||||
dst = dst + (offsetd >> 2);
|
||||
dst[i01 + i20 * ne01] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,3 +153,114 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns(
|
||||
dst[i01 + i20 * ne01] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
__attribute__((qcom_reqd_sub_group_size("half")))
|
||||
__kernel void kernel_gemv_moe_q4_k_f32_ns_wimg(
|
||||
__read_only image1d_buffer_t src0_q,
|
||||
__global half * src0_d,
|
||||
__global half * src0_dm,
|
||||
__global uchar * src0_s,
|
||||
__read_only image1d_buffer_t src1,
|
||||
__global uint * src2,
|
||||
__global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne11
|
||||
) {
|
||||
uint i01 = get_global_id(0);
|
||||
uint i20 = get_global_id(2);
|
||||
uint sgid = get_local_id(1);
|
||||
uint slid = get_sub_group_local_id();
|
||||
|
||||
if (i01 >= ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
uint i11 = i20 % ne11;
|
||||
|
||||
uint expert_id = src2[i20];
|
||||
|
||||
int num_superblocks = ne00 / QK_K;
|
||||
int num_subblocks = ne00 / 32;
|
||||
int scales_per_row = num_superblocks * K_SCALE_SIZE;
|
||||
|
||||
uint expert_q_offset = expert_id * (ne00 / 8) * ne01;
|
||||
uint expert_d_offset = expert_id * num_superblocks * ne01;
|
||||
|
||||
__private float sum = 0.0f;
|
||||
|
||||
for (uint ib = sgid; ib < num_subblocks; ib += N_SIMDGROUP) {
|
||||
uint sb = ib / 8;
|
||||
uint j = ib % 8;
|
||||
|
||||
half d_val = src0_d[expert_d_offset + sb * ne01 + i01];
|
||||
half dm_val = src0_dm[expert_d_offset + sb * ne01 + i01];
|
||||
|
||||
global const uchar * sc = src0_s + (expert_id * ne01 + i01) * scales_per_row + sb * K_SCALE_SIZE;
|
||||
uchar sv, mn;
|
||||
get_scale_min_k4(j, sc, &sv, &mn);
|
||||
|
||||
float scale = (float)d_val * (float)sv;
|
||||
float minv = (float)dm_val * (float)mn;
|
||||
|
||||
uint q_base = expert_q_offset + ib * ne01 * 4 + i01;
|
||||
|
||||
uint4 regQ;
|
||||
regQ.s0 = read_imageui(src0_q, (int)(q_base)).x;
|
||||
regQ.s1 = read_imageui(src0_q, (int)(q_base + ne01)).x;
|
||||
regQ.s2 = read_imageui(src0_q, (int)(q_base + ne01 * 2)).x;
|
||||
regQ.s3 = read_imageui(src0_q, (int)(q_base + ne01 * 3)).x;
|
||||
|
||||
uint y_offset = i11 * ne00 / 4 + ib * 8;
|
||||
|
||||
float8 fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s0), scale, minv);
|
||||
|
||||
float4 shared_y4;
|
||||
shared_y4 = read_imagef(src1, (y_offset + 0));
|
||||
float4 acc = shared_y4 * fp32x8.lo;
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 1));
|
||||
acc += shared_y4 * fp32x8.hi;
|
||||
|
||||
fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s1), scale, minv);
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 2));
|
||||
acc += shared_y4 * fp32x8.lo;
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 3));
|
||||
acc += shared_y4 * fp32x8.hi;
|
||||
|
||||
fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s2), scale, minv);
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 4));
|
||||
acc += shared_y4 * fp32x8.lo;
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 5));
|
||||
acc += shared_y4 * fp32x8.hi;
|
||||
|
||||
fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s3), scale, minv);
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 6));
|
||||
acc += shared_y4 * fp32x8.lo;
|
||||
|
||||
shared_y4 = read_imagef(src1, (y_offset + 7));
|
||||
acc += shared_y4 * fp32x8.hi;
|
||||
|
||||
sum += ((acc.s0 + acc.s1) + (acc.s2 + acc.s3));
|
||||
}
|
||||
|
||||
__local float reduceLM[SIMDGROUP_WIDTH * (N_SIMDGROUP - 1)];
|
||||
if (sgid == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = sum;
|
||||
if (sgid == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = sum;
|
||||
if (sgid == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = sum;
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
|
||||
if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
|
||||
if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
|
||||
|
||||
if (sgid == 0) {
|
||||
dst = dst + (offsetd >> 2);
|
||||
dst[i01 + i20 * ne01] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
// Fused MoE combine epilogue: replaces the router-weight MUL + the (n_expert_used-1)
|
||||
// cross-expert ADD chain with ONE weighted-sum-across-experts pass.
|
||||
// dst[row, tok] = sum_e experts[row, e, tok] * weights[0, e, tok]
|
||||
// experts: [n_embd, n_expert_used, n_tokens] f32 (contiguous after down-proj GEMM)
|
||||
// weights: [1, n_expert_used, n_tokens] f32
|
||||
// dst: [n_embd, n_tokens] f32
|
||||
// One read of experts + one write of dst (eliminates the intermediate weighted
|
||||
// buffer and the k-1 elementwise add round-trips). Vectorized float4 over rows.
|
||||
// strides e1/e2/w1/w2/d1 are in ELEMENTS (floats).
|
||||
|
||||
__kernel void kernel_moe_combine_f32(
|
||||
__global const char * e_buf, ulong off_e,
|
||||
__global const char * w_buf, ulong off_w,
|
||||
__global char * d_buf, ulong off_d,
|
||||
int n_embd4, // n_embd / 4
|
||||
int k, // n_expert_used
|
||||
int n_tokens,
|
||||
uint e1, uint e2, // experts strides (elements): per-expert, per-token
|
||||
uint w1, uint w2, // weights strides (elements)
|
||||
uint d1) // dst per-token stride (elements)
|
||||
{
|
||||
const uint r4 = get_global_id(0);
|
||||
const uint tok = get_global_id(1);
|
||||
if (r4 >= (uint)n_embd4 || tok >= (uint)n_tokens) return;
|
||||
|
||||
__global const float * E = (__global const float *)(e_buf + off_e) + tok*e2 + r4*4u;
|
||||
__global const float * W = (__global const float *)(w_buf + off_w) + tok*w2;
|
||||
|
||||
float4 acc = (float4)(0.0f);
|
||||
for (int e = 0; e < k; ++e) {
|
||||
acc = mad(vload4(0, E + (uint)e*e1), (float4)(W[(uint)e*w1]), acc);
|
||||
}
|
||||
|
||||
__global float * D = (__global float *)(d_buf + off_d) + tok*d1 + r4*4u;
|
||||
vstore4(acc, 0, D);
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
|
||||
// Fused MoE activation reorder + q8_1 quantization for the dp4a prefill GEMM.
|
||||
// Combines kernel_moe_reorder_b (gather src1 rows per the post-router map) with
|
||||
// the q8_1 quant pre-pass, so the f32 reordered-activation tile buffer is never
|
||||
// materialised (saves a full write + read of [tok_slots * ne00] floats).
|
||||
//
|
||||
// One work-item per (token_slot, 32-block). Padding lanes (router 0xFFFFFFFF)
|
||||
// emit d=0,s=0,qs=0 so they contribute nothing to the GEMM, exactly as the
|
||||
// reorder zero-fill did. Output layout matches kernel_moe_quant_a_q8_1:
|
||||
// qa[token_slot*K + blk*32 + i], da/sa[token_slot*(K/32) + blk].
|
||||
__kernel void kernel_moe_reorder_quant_a_q8_1(
|
||||
__global const float * src, // original activations (offset applied)
|
||||
__global const uint * router, // post-router indices [tok_slots]
|
||||
__global char * qa,
|
||||
__global half * da,
|
||||
__global half * sa,
|
||||
__global const int * total_tiles,
|
||||
uint K,
|
||||
ushort map_ratio,
|
||||
uint tile_size,
|
||||
uint n_kblocks // K / 32
|
||||
) {
|
||||
const uint blk = get_global_id(0); // 32-block along K
|
||||
const uint tok = get_global_id(1); // token slot (post_router_idx)
|
||||
|
||||
if (blk >= n_kblocks || tok >= (uint)total_tiles[0] * tile_size) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint out_base = tok * K + blk * 32;
|
||||
const uint bidx = tok * n_kblocks + blk;
|
||||
|
||||
const uint router_idx = router[tok];
|
||||
|
||||
float v[32];
|
||||
float amax = 0.0f;
|
||||
if (router_idx == 0xFFFFFFFF) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) v[i] = 0.0f;
|
||||
} else {
|
||||
const uint act_idx = router_idx / map_ratio;
|
||||
const uint in_base = act_idx * K + blk * 32;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
v[i] = src[in_base + i];
|
||||
amax = fmax(amax, fabs(v[i]));
|
||||
}
|
||||
}
|
||||
|
||||
const float d = amax / 127.0f;
|
||||
const float id = (amax > 0.0f) ? (127.0f / amax) : 0.0f;
|
||||
|
||||
int sum = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
const int q = (int)rint(v[i] * id);
|
||||
qa[out_base + i] = (char)q;
|
||||
sum += q;
|
||||
}
|
||||
|
||||
da[bidx] = (half)d;
|
||||
sa[bidx] = (half)(d * (float)sum);
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||||
|
||||
// Quantize a contiguous [N, K] f32 activation buffer (token-major, K contiguous
|
||||
// per token) into q8_1 blocks of 32: int8 quants + per-block scale d + per-block
|
||||
// sum s (= d * Sum(qs)). Consumed by kernel_gemm_noshuffle_q4_k_q8_1_dp4a for the
|
||||
// dp4a (int8) dense q4_K prefill GEMM. One work-item per 32-element block.
|
||||
__kernel void kernel_quant_a_q8_1(
|
||||
__global const float * src, // [N * K]
|
||||
__global char * qa, // [N * K]
|
||||
__global half * da, // [N * (K/32)]
|
||||
__global half * sa, // [N * (K/32)]
|
||||
int total_blocks // N * (K/32)
|
||||
) {
|
||||
const int blk = get_global_id(0);
|
||||
if (blk >= total_blocks) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int base = blk * 32;
|
||||
|
||||
float v[32];
|
||||
float amax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
v[i] = src[base + i];
|
||||
amax = fmax(amax, fabs(v[i]));
|
||||
}
|
||||
|
||||
const float d = amax / 127.0f;
|
||||
const float id = (amax > 0.0f) ? (127.0f / amax) : 0.0f;
|
||||
|
||||
int sum = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
const int q = (int)rint(v[i] * id);
|
||||
qa[base + i] = (char)q;
|
||||
sum += q;
|
||||
}
|
||||
|
||||
da[blk] = (half)d;
|
||||
sa[blk] = (half)(d * (float)sum);
|
||||
}
|
||||
+120
-1
@@ -377,6 +377,104 @@ static void dequantize_mul_mat_vec_q2_k(const void *__restrict__ vx,
|
||||
}
|
||||
}
|
||||
|
||||
static void dequantize_mul_mat_vec_q2_k_reorder(const void *__restrict__ vx,
|
||||
const float *__restrict__ yy,
|
||||
float *__restrict__ dst,
|
||||
const int ncols, int nrows,
|
||||
const sycl::nd_item<3> &item_ct1) {
|
||||
|
||||
static_assert(16%K_QUANTS_PER_ITERATION == 0, "16 must be divisible by K_QUANTS_PER_ITERATION");
|
||||
|
||||
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
|
||||
item_ct1.get_local_id(1);
|
||||
if (row > nrows) return;
|
||||
|
||||
const int num_blocks_per_row = ncols / QK_K;
|
||||
const int ib0 = row*num_blocks_per_row;
|
||||
|
||||
// SOA base pointers for the reordered layout:
|
||||
// [qs: nb * (QK_K/4)] [scales: nb * (QK_K/16)] [dm: nb * sizeof(half2)]
|
||||
const int nb = nrows * num_blocks_per_row;
|
||||
const uint8_t * qs_base = (const uint8_t *)vx;
|
||||
const uint8_t * scales_base = qs_base + (size_t)nb * (QK_K / 4);
|
||||
const sycl::half2 * dm_base = (const sycl::half2 *)(scales_base + (size_t)nb * (QK_K / 16));
|
||||
|
||||
float tmp = 0; // partial sum for thread in warp
|
||||
|
||||
#if QK_K == 256
|
||||
const int tid =
|
||||
item_ct1.get_local_id(2) / K_QUANTS_PER_ITERATION; // 0...7 or 0...15
|
||||
const int ix =
|
||||
item_ct1.get_local_id(2) % K_QUANTS_PER_ITERATION; // 0 or 0,1
|
||||
|
||||
const int step = 16/K_QUANTS_PER_ITERATION;
|
||||
|
||||
const int in = tid % step; // 0...15 or 0...7
|
||||
|
||||
const int l0 = K_QUANTS_PER_ITERATION*in; // 0...15 or 0...14 in steps of 2
|
||||
|
||||
uint32_t aux[4];
|
||||
const uint8_t * d = (const uint8_t *)aux;
|
||||
const uint8_t * m = (const uint8_t *)(aux + 2);
|
||||
|
||||
for (int i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) {
|
||||
const int bi = ib0 + i;
|
||||
|
||||
const sycl::half2 dm_val = dm_base[bi];
|
||||
const float dall = dm_val[0];
|
||||
const float dmin = dm_val[1];
|
||||
|
||||
for (int im = 0; im < 2; ++im) {
|
||||
const int q_offset = 32*im + l0;
|
||||
const int s_offset = 8*im;
|
||||
const int y_offset = 128*im + l0;
|
||||
|
||||
const float * y = yy + i * QK_K + y_offset;
|
||||
const uint8_t * q = qs_base + bi * (QK_K / 4) + q_offset;
|
||||
|
||||
const uint32_t * a = (const uint32_t *)(scales_base + bi * (QK_K / 16) + s_offset);
|
||||
aux[0] = a[0] & 0x0f0f0f0f;
|
||||
aux[1] = a[1] & 0x0f0f0f0f;
|
||||
aux[2] = (a[0] >> 4) & 0x0f0f0f0f;
|
||||
aux[3] = (a[1] >> 4) & 0x0f0f0f0f;
|
||||
|
||||
float sum1 = 0, sum2 = 0;
|
||||
for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) {
|
||||
sum1 += y[l+ 0] * d[0] * ((q[l+ 0] >> 0) & 3)
|
||||
+ y[l+32] * d[2] * ((q[l+ 0] >> 2) & 3)
|
||||
+ y[l+64] * d[4] * ((q[l+ 0] >> 4) & 3)
|
||||
+ y[l+96] * d[6] * ((q[l+ 0] >> 6) & 3)
|
||||
+ y[l+16] * d[1] * ((q[l+16] >> 0) & 3)
|
||||
+ y[l+48] * d[3] * ((q[l+16] >> 2) & 3)
|
||||
+ y[l+80] * d[5] * ((q[l+16] >> 4) & 3)
|
||||
+y[l+112] * d[7] * ((q[l+16] >> 6) & 3);
|
||||
sum2 += y[l+ 0] * m[0] + y[l+32] * m[2] + y[l+64] * m[4] + y[ l+96] * m[6]
|
||||
+ y[l+16] * m[1] + y[l+48] * m[3] + y[l+80] * m[5] + y[l+112] * m[7];
|
||||
|
||||
}
|
||||
tmp += dall * sum1 - dmin * sum2;
|
||||
}
|
||||
}
|
||||
#else
|
||||
GGML_UNUSED(vx);
|
||||
GGML_UNUSED(yy);
|
||||
GGML_UNUSED(ncols);
|
||||
GGML_UNUSED(item_ct1);
|
||||
GGML_ABORT("Q2_K reorder DMMV not supported for QK_K != 256");
|
||||
#endif
|
||||
|
||||
// sum up partial sums and write back result
|
||||
#pragma unroll
|
||||
for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
|
||||
tmp +=
|
||||
dpct::permute_sub_group_by_xor(item_ct1.get_sub_group(), tmp, mask);
|
||||
}
|
||||
|
||||
if (item_ct1.get_local_id(2) == 0) {
|
||||
dst[row] = tmp;
|
||||
}
|
||||
}
|
||||
|
||||
static void dequantize_mul_mat_vec_q3_k(const void *__restrict__ vx,
|
||||
const float *__restrict__ yy,
|
||||
float *__restrict__ dst,
|
||||
@@ -1664,6 +1762,22 @@ static void dequantize_mul_mat_vec_q2_K_sycl(const void *vx, const float *y,
|
||||
});
|
||||
}
|
||||
|
||||
static void dequantize_mul_mat_vec_q2_K_sycl_reorder(const void *vx, const float *y,
|
||||
float *dst, const int ncols,
|
||||
const int nrows,
|
||||
dpct::queue_ptr stream) {
|
||||
GGML_ASSERT(ncols % QK_K == 0);
|
||||
const int ny = 2 / K_QUANTS_PER_ITERATION;
|
||||
const int block_num_y = (nrows + ny - 1) / ny;
|
||||
const sycl::range<3> block_nums(1, 1, block_num_y);
|
||||
const sycl::range<3> block_dims(1, ny, WARP_SIZE);
|
||||
stream->parallel_for(
|
||||
sycl::nd_range<3>(block_nums * block_dims, block_dims),
|
||||
[=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
|
||||
dequantize_mul_mat_vec_q2_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
|
||||
});
|
||||
}
|
||||
|
||||
static void dequantize_mul_mat_vec_q3_K_sycl(const void *vx, const float *y,
|
||||
float *dst, const int ncols,
|
||||
const int nrows,
|
||||
@@ -1859,7 +1973,12 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
|
||||
}
|
||||
break;
|
||||
case GGML_TYPE_Q2_K:
|
||||
dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
|
||||
if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
|
||||
((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
|
||||
dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
|
||||
} else {
|
||||
dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
|
||||
}
|
||||
break;
|
||||
case GGML_TYPE_Q3_K:
|
||||
if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
|
||||
|
||||
@@ -547,6 +547,7 @@ ggml_backend_sycl_buffer_init_tensor(ggml_backend_buffer_t buffer,
|
||||
switch (tensor->type) {
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_Q2_K:
|
||||
case GGML_TYPE_Q3_K:
|
||||
case GGML_TYPE_Q4_K:
|
||||
case GGML_TYPE_Q5_K:
|
||||
@@ -3675,6 +3676,7 @@ inline bool ggml_sycl_supports_reorder_mul_mat_sycl(enum ggml_type type) {
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q8_0:
|
||||
return true;
|
||||
case GGML_TYPE_Q2_K:
|
||||
case GGML_TYPE_Q3_K:
|
||||
case GGML_TYPE_Q4_K:
|
||||
case GGML_TYPE_Q5_K:
|
||||
@@ -3690,6 +3692,7 @@ inline bool ggml_sycl_supports_reorder_dmmv(enum ggml_type type) {
|
||||
case GGML_TYPE_Q1_0:
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q8_0:
|
||||
case GGML_TYPE_Q2_K:
|
||||
case GGML_TYPE_Q3_K:
|
||||
case GGML_TYPE_Q4_K:
|
||||
case GGML_TYPE_Q5_K:
|
||||
@@ -4069,6 +4072,49 @@ static bool reorder_qw_q6_k_moe(uint8_t * data_device, size_t expert_bytes, int6
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool reorder_qw_q2_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
|
||||
GGML_ASSERT(size % sizeof(block_q2_K) == 0);
|
||||
GGML_ASSERT(offset % sizeof(block_q2_K) == 0);
|
||||
|
||||
const int nblocks = size / sizeof(block_q2_K);
|
||||
|
||||
sycl_reorder_temp_buffer tmp(stream, size);
|
||||
if (!tmp) {
|
||||
GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
|
||||
return false;
|
||||
}
|
||||
uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
|
||||
|
||||
sycl::event copy_event;
|
||||
SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
|
||||
if (!g_ggml_sycl_use_async_mem_op) {
|
||||
copy_event.wait();
|
||||
}
|
||||
|
||||
auto * qs_ptr = data_device;
|
||||
auto * scales_ptr = qs_ptr + (QK_K / 4) * nblocks;
|
||||
sycl::half2 * dm_ptr = (sycl::half2 *) (scales_ptr + (QK_K / 16) * nblocks);
|
||||
|
||||
auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
|
||||
const block_q2_K * x = (const block_q2_K *) tmp_buf;
|
||||
const int ib = i;
|
||||
|
||||
for (int j = 0; j < QK_K / 4; ++j) {
|
||||
qs_ptr[ib * (QK_K / 4) + j] = x[ib].qs[j];
|
||||
}
|
||||
|
||||
for (int j = 0; j < QK_K / 16; ++j) {
|
||||
scales_ptr[ib * (QK_K / 16) + j] = x[ib].scales[j];
|
||||
}
|
||||
|
||||
dm_ptr[ib] = x[ib].dm;
|
||||
});
|
||||
if (!g_ggml_sycl_use_async_mem_op) {
|
||||
reorder_event.wait_and_throw();
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool reorder_qw_q3_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
|
||||
GGML_ASSERT(size % sizeof(block_q3_K) == 0);
|
||||
GGML_ASSERT(offset % sizeof(block_q3_K) == 0);
|
||||
@@ -4245,6 +4291,8 @@ static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) {
|
||||
return reorder_qw_q4_0(data_device, ncols, nrows, size, 0, stream);
|
||||
case GGML_TYPE_Q8_0:
|
||||
return reorder_qw_q8_0(data_device, ncols, nrows, size, 0, stream);
|
||||
case GGML_TYPE_Q2_K:
|
||||
return reorder_qw_q2_k(data_device, size, 0, stream);
|
||||
case GGML_TYPE_Q3_K:
|
||||
return reorder_qw_q3_k(data_device, size, 0, stream);
|
||||
case GGML_TYPE_Q4_K:
|
||||
|
||||
@@ -6501,6 +6501,14 @@ static vk_device ggml_vk_get_device(size_t idx) {
|
||||
device->mul_mat_id_m[i] = true;
|
||||
device->mul_mat_id_s[i] = false;
|
||||
break;
|
||||
case VK_VENDOR_ID_QUALCOMM:
|
||||
device->mul_mat_l[i] = false;
|
||||
device->mul_mat_m[i] = true;
|
||||
device->mul_mat_s[i] = true;
|
||||
device->mul_mat_id_l[i] = false;
|
||||
device->mul_mat_id_m[i] = true;
|
||||
device->mul_mat_id_s[i] = true;
|
||||
break;
|
||||
#endif
|
||||
default:
|
||||
device->mul_mat_l[i] = true;
|
||||
|
||||
+40
-2
@@ -1079,6 +1079,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
|
||||
"RWKV_WKV7",
|
||||
"SOLVE_TRI",
|
||||
"GATED_DELTA_NET",
|
||||
"LIGHTNING_INDEXER",
|
||||
|
||||
"UNARY",
|
||||
|
||||
@@ -1096,7 +1097,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
|
||||
"GLU",
|
||||
};
|
||||
|
||||
static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
|
||||
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
|
||||
|
||||
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
|
||||
"none",
|
||||
@@ -1190,6 +1191,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
|
||||
"rwkv_wkv7(r, w, k, v, a, b, s)",
|
||||
"A X = B, A triangular, solve X",
|
||||
"gated_delta_net(q, k, v, g, beta, s)",
|
||||
"lightning_indexer(q, k, weights, mask)",
|
||||
|
||||
"unary(x)",
|
||||
|
||||
@@ -1207,7 +1209,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
|
||||
"glu(x)",
|
||||
};
|
||||
|
||||
static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
|
||||
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
|
||||
|
||||
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
|
||||
|
||||
@@ -6287,6 +6289,42 @@ struct ggml_tensor * ggml_gated_delta_net(
|
||||
return result;
|
||||
}
|
||||
|
||||
// ggml_lightning_indexer
|
||||
|
||||
struct ggml_tensor * ggml_lightning_indexer(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * q,
|
||||
struct ggml_tensor * k,
|
||||
struct ggml_tensor * weights,
|
||||
struct ggml_tensor * mask) {
|
||||
|
||||
GGML_ASSERT( q->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( weights->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( mask->type == GGML_TYPE_F16);
|
||||
GGML_ASSERT( q->ne[0] == k->ne[0]);
|
||||
GGML_ASSERT( mask->ne[0] == k->ne[2]);
|
||||
GGML_ASSERT( q->ne[1] == weights->ne[0]);
|
||||
GGML_ASSERT( k->ne[1] == 1);
|
||||
GGML_ASSERT( mask->ne[1] == q->ne[2]);
|
||||
GGML_ASSERT( q->ne[2] == weights->ne[1]);
|
||||
GGML_ASSERT(weights->ne[2] == 1);
|
||||
GGML_ASSERT( mask->ne[2] == 1);
|
||||
GGML_ASSERT( q->ne[3] == k->ne[3]);
|
||||
GGML_ASSERT( k->ne[3] == weights->ne[3]);
|
||||
GGML_ASSERT(weights->ne[3] % mask->ne[3] == 0);
|
||||
|
||||
int64_t ne[4] = { k->ne[2], q->ne[2], 1, q->ne[3] };
|
||||
struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
|
||||
|
||||
result->op = GGML_OP_LIGHTNING_INDEXER;
|
||||
result->src[0] = q;
|
||||
result->src[1] = k;
|
||||
result->src[2] = weights;
|
||||
result->src[3] = mask;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct ggml_hash_set ggml_hash_set_new(size_t size) {
|
||||
|
||||
@@ -557,6 +557,10 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr
|
||||
GGML_LOG_ERROR("%s: encountered bad_alloc error while reading key %" PRIi64 "\n", __func__, i);
|
||||
ok = false;
|
||||
}
|
||||
if (ok && key.empty()) {
|
||||
GGML_LOG_ERROR("%s: key %" PRIi64 " is empty\n", __func__, i);
|
||||
ok = false;
|
||||
}
|
||||
for (size_t j = 0; ok && j < ctx->kv.size(); ++j) {
|
||||
if (key == ctx->kv[j].key) {
|
||||
GGML_LOG_ERROR("%s: duplicate key '%s' for tensors %zu and %" PRIi64 " \n", __func__, key.c_str(), j, i);
|
||||
|
||||
@@ -5,7 +5,7 @@ import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
HTTPLIB_VERSION = "refs/tags/v0.49.0"
|
||||
HTTPLIB_VERSION = "refs/tags/v0.50.1"
|
||||
|
||||
vendor = {
|
||||
"https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp",
|
||||
|
||||
@@ -55,6 +55,12 @@ static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = {
|
||||
/*.n_tokens_per_seq =*/ 16,
|
||||
};
|
||||
|
||||
static const llm_fused_op_probe llm_fused_op_lid_probe = {
|
||||
/*.op =*/ LLM_FUSED_OP_LIGHTNING_INDEXER,
|
||||
/*.name =*/ "Lightning Indexer",
|
||||
/*.n_tokens_per_seq =*/ 1,
|
||||
};
|
||||
|
||||
llama_context::llama_context(
|
||||
const llama_model & model,
|
||||
llama_context_params params) :
|
||||
@@ -226,6 +232,9 @@ llama_context::llama_context(
|
||||
cparams.fused_gdn_ch = true;
|
||||
cparams.auto_fgdn = true;
|
||||
|
||||
cparams.fused_lid = true;
|
||||
cparams.auto_flid = true;
|
||||
|
||||
// with causal attention, the batch size is limited by the context size
|
||||
cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
|
||||
|
||||
@@ -522,6 +531,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
|
||||
resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);
|
||||
cparams.auto_fgdn = false;
|
||||
}
|
||||
|
||||
if (cparams.auto_flid) {
|
||||
LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func);
|
||||
resolve(llm_fused_op_lid_probe, cparams.fused_lid);
|
||||
cparams.auto_flid = false;
|
||||
}
|
||||
}
|
||||
|
||||
void llama_context::sched_reserve() {
|
||||
|
||||
@@ -41,6 +41,8 @@ struct llama_cparams {
|
||||
bool fused_gdn_ar; // use fused gated delta net (autoregressive)
|
||||
bool fused_gdn_ch; // use fused gated delta net (chunked)
|
||||
bool auto_fgdn;
|
||||
bool fused_lid; // use fused lightning indexer
|
||||
bool auto_flid;
|
||||
bool no_perf;
|
||||
bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
|
||||
bool op_offload;
|
||||
|
||||
+3
-3
@@ -842,7 +842,7 @@ static void dsv4_build_comp_inputs(
|
||||
GGML_ASSERT(n_stream > 0);
|
||||
GGML_ASSERT(n_tokens%n_stream == 0);
|
||||
|
||||
inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream);
|
||||
inp.kq_mask = ggml_new_tensor_4d(ctx, (strcmp(name, "lid") != 0 && cparams.flash_attn) || (strcmp(name, "lid") == 0 && cparams.fused_lid) ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream);
|
||||
ggml_set_input(inp.kq_mask);
|
||||
ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str());
|
||||
}
|
||||
@@ -3025,9 +3025,9 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
|
||||
{
|
||||
inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch);
|
||||
|
||||
// ensure F32 mask
|
||||
// ensure that mask type matches fused lightning indexer use (requires f16 mask)
|
||||
auto cparams_copy = cparams;
|
||||
cparams_copy.flash_attn = false;
|
||||
cparams_copy.flash_attn = cparams.fused_lid;
|
||||
|
||||
inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy);
|
||||
inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid;
|
||||
|
||||
@@ -42,6 +42,7 @@ enum llm_fused_op {
|
||||
LLM_FUSED_OP_FLASH_ATTN,
|
||||
LLM_FUSED_OP_GDN_AR,
|
||||
LLM_FUSED_OP_GDN_CH,
|
||||
LLM_FUSED_OP_LIGHTNING_INDEXER,
|
||||
};
|
||||
|
||||
enum llm_ffn_op_type : int {
|
||||
|
||||
+59
-15
@@ -29,6 +29,15 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
|
||||
return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);
|
||||
}
|
||||
|
||||
static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) {
|
||||
GGML_ASSERT(ggml_is_contiguous(tensor));
|
||||
GGML_ASSERT(tensor->ne[3] == 1);
|
||||
GGML_ASSERT(stream < (uint32_t) tensor->ne[2]);
|
||||
|
||||
const size_t stream_size = tensor->nb[2];
|
||||
ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size);
|
||||
}
|
||||
|
||||
static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {
|
||||
if (n_stream <= 1) {
|
||||
return 0;
|
||||
@@ -781,11 +790,20 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
|
||||
__func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
|
||||
}
|
||||
|
||||
void llama_dsv4_comp_state::clear(bool data) {
|
||||
void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
|
||||
if (!data) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (seq_id >= 0) {
|
||||
GGML_ASSERT((uint32_t) seq_id < n_stream);
|
||||
for (const auto & layer : layers) {
|
||||
dsv4_clear_tensor_stream(layer.kv, (uint32_t) seq_id);
|
||||
dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
for (auto & [_, buf] : ctxs_bufs) {
|
||||
ggml_backend_buffer_clear(buf.get(), 0);
|
||||
}
|
||||
@@ -1034,7 +1052,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
|
||||
// graph does not necessarily overwrite; uninitialized buffer contents would
|
||||
// otherwise leak in (instance-specific garbage) and corrupt recall. Zero all
|
||||
// compressed buffers up front so reads of un-written rows are deterministic.
|
||||
clear_compressed(true);
|
||||
clear_compressed(-1, true);
|
||||
}
|
||||
|
||||
llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
|
||||
@@ -1147,7 +1165,7 @@ bool llama_kv_cache_dsv4::get_can_shift() const {
|
||||
|
||||
void llama_kv_cache_dsv4::clear(bool data) {
|
||||
kv_raw->clear(data);
|
||||
clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows
|
||||
clear_compressed(-1, true); // DSV4 compressed buffers must never expose stale/uninit rows
|
||||
}
|
||||
|
||||
bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
@@ -1169,7 +1187,7 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
|
||||
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
|
||||
|
||||
if (res) {
|
||||
clear_compressed(true);
|
||||
clear_compressed(seq_id, true);
|
||||
}
|
||||
|
||||
return res;
|
||||
@@ -1177,22 +1195,29 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
|
||||
|
||||
void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
|
||||
kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1);
|
||||
clear_compressed(true);
|
||||
}
|
||||
|
||||
void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
|
||||
GGML_ASSERT(seq_id >= 0 && (uint32_t) seq_id < n_seq_max);
|
||||
|
||||
kv_raw->seq_keep(seq_id);
|
||||
clear_compressed(true);
|
||||
|
||||
for (llama_seq_id id = 0; id < (llama_seq_id) n_seq_max; ++id) {
|
||||
if (id == seq_id) {
|
||||
continue;
|
||||
}
|
||||
|
||||
kv_raw->seq_rm(id, -1, -1);
|
||||
clear_compressed(id, true);
|
||||
}
|
||||
}
|
||||
|
||||
void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
|
||||
kv_raw->seq_add(seq_id, p0, p1, shift);
|
||||
clear_compressed(true);
|
||||
}
|
||||
|
||||
void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
|
||||
kv_raw->seq_div(seq_id, p0, p1, d);
|
||||
clear_compressed(true);
|
||||
}
|
||||
|
||||
llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const {
|
||||
@@ -1328,13 +1353,32 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
|
||||
return lid_state.get();
|
||||
}
|
||||
|
||||
void llama_kv_cache_dsv4::clear_compressed(bool data) {
|
||||
kv_csa->clear(data);
|
||||
kv_hca->clear(data);
|
||||
kv_lid->clear(data);
|
||||
csa_state->clear(data);
|
||||
hca_state->clear(data);
|
||||
lid_state->clear(data);
|
||||
void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
|
||||
if (seq_id < 0) {
|
||||
kv_csa->clear(data);
|
||||
kv_hca->clear(data);
|
||||
kv_lid->clear(data);
|
||||
} else {
|
||||
GGML_ASSERT((uint32_t) seq_id < n_seq_max);
|
||||
|
||||
const auto clear_seq = [seq_id, data](llama_kv_cache * kv) {
|
||||
kv->seq_rm(seq_id, -1, -1);
|
||||
|
||||
if (data) {
|
||||
for (uint32_t il : kv->get_layer_ids()) {
|
||||
dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
clear_seq(kv_csa.get());
|
||||
clear_seq(kv_hca.get());
|
||||
clear_seq(kv_lid.get());
|
||||
}
|
||||
|
||||
csa_state->clear(seq_id, data);
|
||||
hca_state->clear(seq_id, data);
|
||||
lid_state->clear(seq_id, data);
|
||||
}
|
||||
|
||||
//
|
||||
|
||||
@@ -21,7 +21,7 @@ public:
|
||||
const char * name,
|
||||
const llama_memory_i::layer_filter_cb & filter);
|
||||
|
||||
void clear(bool data);
|
||||
void clear(llama_seq_id seq_id, bool data);
|
||||
|
||||
uint32_t get_ratio() const;
|
||||
uint32_t get_state_size() const;
|
||||
@@ -67,6 +67,8 @@ private:
|
||||
// DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches.
|
||||
// The compressed caches are storage only; DSV4-specific visibility and block
|
||||
// planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4.
|
||||
// FIXME: currently the cache only supports non-unified mode even if unified flag is passed
|
||||
// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819
|
||||
|
||||
class llama_kv_cache_dsv4 : public llama_memory_i {
|
||||
public:
|
||||
@@ -146,7 +148,7 @@ private:
|
||||
std::unique_ptr<llama_dsv4_comp_state> hca_state;
|
||||
std::unique_ptr<llama_dsv4_comp_state> lid_state;
|
||||
|
||||
void clear_compressed(bool data);
|
||||
void clear_compressed(llama_seq_id seq_id, bool data);
|
||||
};
|
||||
|
||||
// DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept
|
||||
|
||||
+30
-31
@@ -313,8 +313,7 @@ llama_model * llama_model_create(llm_arch arch, const llama_model_params & param
|
||||
|
||||
if (model != nullptr) {
|
||||
model->arch = arch;
|
||||
auto & devices = model->devices;
|
||||
if (!devices.empty() && devices[0].is_meta && !llm_arch_supports_sm_tensor(arch)) {
|
||||
if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR && !llm_arch_supports_sm_tensor(arch)) {
|
||||
throw std::runtime_error(std::string("LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '") + llm_arch_name(arch) + "'");
|
||||
}
|
||||
}
|
||||
@@ -336,38 +335,38 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
const llama_hparams & hparams = ud->model->hparams;
|
||||
const std::string tensor_name = tensor->name;
|
||||
|
||||
const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
|
||||
const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight");
|
||||
const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight");
|
||||
const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias");
|
||||
const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias");
|
||||
const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias");
|
||||
const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
|
||||
const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*");
|
||||
const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight");
|
||||
const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
|
||||
const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias");
|
||||
const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
|
||||
static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
|
||||
static const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight");
|
||||
static const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight");
|
||||
static const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias");
|
||||
static const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias");
|
||||
static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias");
|
||||
static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
|
||||
static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*");
|
||||
static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight");
|
||||
static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
|
||||
static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias");
|
||||
static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
|
||||
|
||||
const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias");
|
||||
const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a");
|
||||
const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight");
|
||||
const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight");
|
||||
const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight");
|
||||
const std::regex pattern_r_cache ("cache_r_l\\d*");
|
||||
const std::regex pattern_s_cache ("cache_s_l\\d*");
|
||||
const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight");
|
||||
const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight");
|
||||
static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias");
|
||||
static const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a");
|
||||
static const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight");
|
||||
static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight");
|
||||
static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight");
|
||||
static const std::regex pattern_r_cache ("cache_r_l\\d*");
|
||||
static const std::regex pattern_s_cache ("cache_s_l\\d*");
|
||||
static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight");
|
||||
static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight");
|
||||
|
||||
const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight");
|
||||
const std::regex pattern_ffn_up_gate_bias ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias");
|
||||
const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
|
||||
const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
|
||||
const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
|
||||
const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
|
||||
static const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight");
|
||||
static const std::regex pattern_ffn_up_gate_bias ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias");
|
||||
static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
|
||||
static const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
|
||||
static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
|
||||
static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
|
||||
|
||||
const std::regex pattern_output_weight("output\\.weight");
|
||||
const std::regex pattern_output_bias ("output\\.bias");
|
||||
static const std::regex pattern_output_weight("output\\.weight");
|
||||
static const std::regex pattern_output_bias ("output\\.bias");
|
||||
|
||||
struct tensor_config {
|
||||
ggml_backend_meta_split_axis axis;
|
||||
|
||||
+37
-30
@@ -301,43 +301,50 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
|
||||
indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);
|
||||
indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
|
||||
|
||||
// calculate indexer kq
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "indexer_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "indexer_k", il);
|
||||
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "indexer_kq", il);
|
||||
|
||||
// ReLU requires contiguous tensors
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "indexer_kq", il);
|
||||
|
||||
// apply ReLU
|
||||
ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// pre-scale weights to avoid scaling operations on huge indexer_score tensor
|
||||
indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));
|
||||
cb(indexer_weights, "indexer_weights", il);
|
||||
|
||||
// multiply scores by indexer weights
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
ggml_tensor * indexer_score = nullptr;
|
||||
if (cparams.fused_lid) {
|
||||
indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
|
||||
} else {
|
||||
// calculate indexer kq
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "indexer_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "indexer_k", il);
|
||||
|
||||
// sum by q n_indexer_head dimension
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "indexer_kq", il);
|
||||
|
||||
// permute result to match KQ mask
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
// ReLU requires contiguous tensors
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "indexer_kq", il);
|
||||
|
||||
// mask indexer scores
|
||||
ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
|
||||
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
// apply ReLU
|
||||
indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// multiply scores by indexer weights
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// sum by q n_indexer_head dimension
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// permute result to match KQ mask
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// mask indexer scores
|
||||
ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
|
||||
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
}
|
||||
|
||||
// get indices of top k indexer scores
|
||||
uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;
|
||||
|
||||
+22
-15
@@ -556,25 +556,32 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(
|
||||
indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream,
|
||||
indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
|
||||
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "lid_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "lid_k", il);
|
||||
ggml_tensor * indexer_score = nullptr;
|
||||
if (cparams.fused_lid) {
|
||||
indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask);
|
||||
cb(indexer_score, "lid_score_masked", il);
|
||||
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
|
||||
} else {
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "lid_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "lid_k", il);
|
||||
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
|
||||
ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "lid_score", il);
|
||||
indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "lid_score", il);
|
||||
|
||||
indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
|
||||
cb(indexer_score, "lid_score_masked", il);
|
||||
indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
|
||||
cb(indexer_score, "lid_score_masked", il);
|
||||
}
|
||||
|
||||
const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k;
|
||||
ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));
|
||||
|
||||
@@ -7097,6 +7097,67 @@ struct test_diag : public test_case {
|
||||
}
|
||||
};
|
||||
|
||||
// GGML_OP_LIGHTNING_INDEXER
|
||||
struct test_lightning_indexer : public test_case {
|
||||
const int64_t hsk; // indexer K head size
|
||||
const int64_t nh; // num indexer heads
|
||||
const int64_t kv; // kv size
|
||||
const int64_t nb; // batch size
|
||||
const int64_t ns; // num streams
|
||||
const int64_t nm; // ne[3] of mask
|
||||
|
||||
const ggml_type type_K;
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR7(hsk, nh, kv, nb, ns, nm, type_K);
|
||||
}
|
||||
|
||||
double max_nmse_err() override {
|
||||
return 1e-6;
|
||||
}
|
||||
|
||||
uint64_t op_flops(ggml_tensor * t) override {
|
||||
GGML_UNUSED(t);
|
||||
return ((2 * hsk + 2) * nh + 1) * kv * nb * ns;
|
||||
}
|
||||
|
||||
test_lightning_indexer(int64_t hsk = 128, int64_t nh = 64, int64_t kv = 256, int64_t nb = 128, int64_t ns = 1, int64_t nm = 1, ggml_type type_K = GGML_TYPE_F16)
|
||||
: hsk(hsk), nh(nh), kv(kv), nb(nb), ns(ns), nm(nm), type_K(type_K) {}
|
||||
|
||||
ggml_tensor * build_graph(ggml_context * ctx) override {
|
||||
ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hsk, nh, nb, ns);
|
||||
ggml_set_param(q);
|
||||
ggml_set_name(q, "q");
|
||||
|
||||
ggml_tensor * k = ggml_new_tensor_4d(ctx, type_K, hsk, 1, kv, ns);
|
||||
ggml_set_param(k);
|
||||
ggml_set_name(k, "k");
|
||||
|
||||
ggml_tensor * w = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, nh, nb, 1, ns);
|
||||
ggml_set_param(w);
|
||||
ggml_set_name(w, "w");
|
||||
|
||||
ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nm);
|
||||
ggml_set_param(m);
|
||||
ggml_set_name(m, "m");
|
||||
|
||||
ggml_tensor * out = ggml_lightning_indexer(ctx, q, k, w, m);
|
||||
ggml_set_name(out, "out");
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
void initialize_tensors(ggml_context * ctx) override {
|
||||
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
|
||||
if (strcmp(t->name, "m") == 0) {
|
||||
init_tensor_kq_mask(t);
|
||||
} else {
|
||||
init_tensor_uniform(t);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Deserializable generic test case
|
||||
struct input_tensor {
|
||||
ggml_type type;
|
||||
@@ -9393,6 +9454,19 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_falcon(2));
|
||||
#endif
|
||||
|
||||
// lightning_indexer
|
||||
for (int kv : { 256 }) {
|
||||
for (int bs : { 1, 512 }) {
|
||||
for (int nh : { 32, 64 }) {
|
||||
for (auto [ns, nm] : { std::pair{1, 1}, std::pair{4, 4}, std::pair{4, 1} }) {
|
||||
for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) {
|
||||
test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, nm, type_K));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return test_cases;
|
||||
}
|
||||
#ifdef _MSC_VER
|
||||
@@ -9722,6 +9796,19 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 1024, 1)); // 4h PP-1024
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 64, 1, 1, false, true)); // KDA PP-64
|
||||
|
||||
// lightning_indexer
|
||||
for (int kv : { 256, 4096, 65536 }) {
|
||||
for (int bs : { 1, 512, 2048 }) {
|
||||
for (int nh : { 32, 64 }) {
|
||||
for (int ns : { 1, 4 }) {
|
||||
for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) {
|
||||
test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, ns, type_K));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return test_cases;
|
||||
}
|
||||
|
||||
|
||||
@@ -5911,6 +5911,71 @@ static void test_developer_role_to_system_workaround() {
|
||||
}
|
||||
}
|
||||
|
||||
static void test_reasoning_budget_tokens_per_request() {
|
||||
LOG_DBG("%s\n", __func__);
|
||||
// Use Qwen3 template which has <think>...</think> reasoning markers.
|
||||
// The autoparser detects them and sets thinking_start/end_tag, which enables
|
||||
// the reasoning-budget code path in oaicompat_chat_params_parse.
|
||||
auto tmpls = read_templates("models/templates/Qwen-Qwen3-0.6B.jinja");
|
||||
|
||||
server_chat_params opt;
|
||||
opt.tmpls = std::move(tmpls);
|
||||
opt.use_jinja = true;
|
||||
opt.enable_thinking = true;
|
||||
opt.reasoning_budget = -1;
|
||||
opt.reasoning_format = COMMON_REASONING_FORMAT_NONE;
|
||||
|
||||
// Body with per-request reasoning_budget_tokens=0 (suppress thinking).
|
||||
json body = {
|
||||
{"messages", json::array({json{{"role", "user"}, {"content", "hello"}}})},
|
||||
{"reasoning_budget_tokens", 0},
|
||||
};
|
||||
std::vector<raw_buffer> out_files;
|
||||
auto llama_params = oaicompat_chat_params_parse(body, opt, out_files);
|
||||
|
||||
// The per-request value must win over the server default (-1).
|
||||
if (!llama_params.contains("reasoning_budget_tokens")) {
|
||||
throw std::runtime_error("reasoning_budget_tokens missing from llama_params (thinking_end_tag may be empty for this template)");
|
||||
}
|
||||
int got = llama_params["reasoning_budget_tokens"].get<int>();
|
||||
if (got != 0) {
|
||||
throw std::runtime_error(std::string("Expected reasoning_budget_tokens=0, got ") + std::to_string(got));
|
||||
}
|
||||
}
|
||||
|
||||
static void test_reasoning_budget_message_per_request() {
|
||||
LOG_DBG("%s\n", __func__);
|
||||
// Same code path as test_reasoning_budget_tokens_per_request: the Qwen3 template's
|
||||
// <think>...</think> markers enable the reasoning-budget block in oaicompat_chat_params_parse.
|
||||
auto tmpls = read_templates("models/templates/Qwen-Qwen3-0.6B.jinja");
|
||||
|
||||
server_chat_params opt;
|
||||
opt.tmpls = std::move(tmpls);
|
||||
opt.use_jinja = true;
|
||||
opt.enable_thinking = true;
|
||||
opt.reasoning_budget = -1;
|
||||
opt.reasoning_format = COMMON_REASONING_FORMAT_NONE;
|
||||
opt.reasoning_budget_message = "server default";
|
||||
|
||||
// Body with a per-request reasoning_budget_message override.
|
||||
const std::string per_request_message = "per-request message";
|
||||
json body = {
|
||||
{"messages", json::array({json{{"role", "user"}, {"content", "hello"}}})},
|
||||
{"reasoning_budget_message", per_request_message},
|
||||
};
|
||||
std::vector<raw_buffer> out_files;
|
||||
auto llama_params = oaicompat_chat_params_parse(body, opt, out_files);
|
||||
|
||||
// The per-request value must win over the server default.
|
||||
if (!llama_params.contains("reasoning_budget_message")) {
|
||||
throw std::runtime_error("reasoning_budget_message missing from llama_params (thinking_end_tag may be empty for this template)");
|
||||
}
|
||||
std::string got = llama_params["reasoning_budget_message"].get<std::string>();
|
||||
if (got != per_request_message) {
|
||||
throw std::runtime_error("Expected reasoning_budget_message='" + per_request_message + "', got '" + got + "'");
|
||||
}
|
||||
}
|
||||
|
||||
static void test_msg_diffs_compute() {
|
||||
LOG_DBG("%s\n", __func__);
|
||||
{
|
||||
@@ -6068,6 +6133,8 @@ int main(int argc, char ** argv) {
|
||||
test_convert_responses_to_chatcmpl();
|
||||
test_developer_role_to_system_workaround();
|
||||
test_template_generation_prompt();
|
||||
test_reasoning_budget_tokens_per_request();
|
||||
test_reasoning_budget_message_per_request();
|
||||
test_template_output_peg_parsers(detailed_debug);
|
||||
std::cout << "\n[chat] All tests passed!" << '\n';
|
||||
}
|
||||
|
||||
+6
-1
@@ -26,6 +26,7 @@ enum handcrafted_file_type {
|
||||
HANDCRAFTED_HEADER_EMPTY = 800,
|
||||
|
||||
HANDCRAFTED_KV_BAD_KEY_SIZE = 10 + offset_has_kv,
|
||||
HANDCRAFTED_KV_EMPTY_KEY = 15 + offset_has_kv,
|
||||
HANDCRAFTED_KV_BAD_TYPE = 20 + offset_has_kv,
|
||||
// HANDCRAFTED_KV_BAD_VALUE_SIZE = 30 + offset_has_kv, // removed because it can result in allocations > 1 TB (default sanitizer limit)
|
||||
HANDCRAFTED_KV_DUPLICATE_KEY = 40 + offset_has_kv,
|
||||
@@ -64,6 +65,7 @@ static std::string handcrafted_file_type_name(const enum handcrafted_file_type h
|
||||
case HANDCRAFTED_HEADER_EMPTY: return "HEADER_EMPTY";
|
||||
|
||||
case HANDCRAFTED_KV_BAD_KEY_SIZE: return "KV_BAD_KEY_SIZE";
|
||||
case HANDCRAFTED_KV_EMPTY_KEY: return "KV_EMPTY_KEY";
|
||||
case HANDCRAFTED_KV_BAD_TYPE: return "KV_BAD_TYPE";
|
||||
case HANDCRAFTED_KV_DUPLICATE_KEY: return "KV_DUPLICATE_KEY";
|
||||
case HANDCRAFTED_KV_BAD_ALIGN: return "KV_BAD_ALIGN";
|
||||
@@ -284,7 +286,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
|
||||
const enum gguf_type type = gguf_type(hft == HANDCRAFTED_KV_BAD_TYPE ? GGUF_TYPE_COUNT : kv_types[i].first);
|
||||
const enum gguf_type type_arr = gguf_type(hft == HANDCRAFTED_KV_BAD_TYPE ? GGUF_TYPE_COUNT : kv_types[i].second);
|
||||
|
||||
const std::string key = "my_key_" + std::to_string((hft == HANDCRAFTED_KV_DUPLICATE_KEY ? i/2 : i));
|
||||
const std::string key = hft == HANDCRAFTED_KV_EMPTY_KEY
|
||||
? ""
|
||||
: "my_key_" + std::to_string((hft == HANDCRAFTED_KV_DUPLICATE_KEY ? i/2 : i));
|
||||
|
||||
if (hft == HANDCRAFTED_KV_BAD_KEY_SIZE) {
|
||||
const uint64_t n = -1;
|
||||
@@ -732,6 +736,7 @@ static std::pair<int, int> test_handcrafted_file(const unsigned int seed) {
|
||||
HANDCRAFTED_HEADER_EMPTY,
|
||||
|
||||
HANDCRAFTED_KV_BAD_KEY_SIZE,
|
||||
HANDCRAFTED_KV_EMPTY_KEY,
|
||||
HANDCRAFTED_KV_BAD_TYPE,
|
||||
HANDCRAFTED_KV_DUPLICATE_KEY,
|
||||
HANDCRAFTED_KV_BAD_ALIGN,
|
||||
|
||||
@@ -78,7 +78,84 @@ static llama_tokens test_baseline(struct llama_model * model, const struct commo
|
||||
}
|
||||
|
||||
|
||||
// Test 2: state load
|
||||
// Test 2: sequence removal isolation
|
||||
// - decode the same prefix into two sequences
|
||||
// - remove sequence 0
|
||||
// - verify that sequence 1 remains unchanged
|
||||
static bool test_seq_rm_isolated(
|
||||
struct llama_model * model,
|
||||
const struct common_params & params,
|
||||
const llama_tokens & tokens) {
|
||||
auto params_ctx = common_context_params_to_llama(params);
|
||||
params_ctx.n_ctx = 256;
|
||||
params_ctx.n_seq_max = 2;
|
||||
params_ctx.kv_unified = true;
|
||||
|
||||
auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
|
||||
if (!ctx) {
|
||||
LOG_ERR("%s: failed to create context\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
LOG("\n=== Test 2: sequence removal isolation ===\n");
|
||||
|
||||
const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128;
|
||||
for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
|
||||
llama_batch_ptr batch(n_tokens, 0, 1);
|
||||
for (size_t i = 0; i < n_tokens; ++i) {
|
||||
common_batch_add(batch.get(), tokens[i], i, { seq_id }, false);
|
||||
}
|
||||
|
||||
if (llama_decode(ctx.get(), batch.get())) {
|
||||
LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) {
|
||||
const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id);
|
||||
if (state_size == 0) {
|
||||
LOG_ERR("%s: sequence state is empty\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
state.resize(state_size);
|
||||
const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id);
|
||||
if (ncopy != state.size()) {
|
||||
LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
|
||||
__func__, ncopy, state.size());
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
};
|
||||
|
||||
std::vector<uint8_t> state_before;
|
||||
if (!get_seq_state(1, state_before)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
|
||||
LOG_ERR("%s: failed to remove sequence 0\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> state_after;
|
||||
if (!get_seq_state(1, state_after)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (state_before != state_after) {
|
||||
LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
LOG("PASS\n");
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
// Test 3: state load
|
||||
// - create a new context
|
||||
// - load state from file
|
||||
// - replay the last prompt token
|
||||
@@ -90,7 +167,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para
|
||||
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
|
||||
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
|
||||
|
||||
LOG("\n=== Test 2: state load ===\n");
|
||||
LOG("\n=== Test 3: state load ===\n");
|
||||
|
||||
// Load state from file
|
||||
llama_tokens unused_sts(tokens.size());
|
||||
@@ -126,7 +203,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para
|
||||
}
|
||||
|
||||
|
||||
// Test 3: seq copy (host)
|
||||
// Test 4: seq copy (host)
|
||||
// - create a multi-seq context
|
||||
// - load state from file
|
||||
// - replay the last prompt token
|
||||
@@ -141,7 +218,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par
|
||||
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
|
||||
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
|
||||
|
||||
LOG("\n=== Test 3: seq copy (host) ===\n");
|
||||
LOG("\n=== Test 4: seq copy (host) ===\n");
|
||||
|
||||
// Load state from file
|
||||
llama_tokens unused_sts(tokens.size());
|
||||
@@ -198,7 +275,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par
|
||||
}
|
||||
|
||||
|
||||
// Test 4: seq copy (device)
|
||||
// Test 5: seq copy (device)
|
||||
// - create a multi-seq context
|
||||
// - load state from file
|
||||
// - replay the last prompt token
|
||||
@@ -213,7 +290,7 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p
|
||||
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
|
||||
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
|
||||
|
||||
LOG("\n=== Test 4: seq copy (device) ===\n");
|
||||
LOG("\n=== Test 5: seq copy (device) ===\n");
|
||||
|
||||
// Load state from file
|
||||
llama_tokens unused_sts(tokens.size());
|
||||
@@ -337,17 +414,22 @@ int main(int argc, char ** argv) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Test 2: state load
|
||||
// Test 2: sequence removal isolation
|
||||
if (!test_seq_rm_isolated(model, params, tokens)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Test 3: state load
|
||||
if (!test_state_load(model, params, tokens, result_baseline)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Test 3: seq copy (host)
|
||||
// Test 4: seq copy (host)
|
||||
if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Test 4: seq copy (device)
|
||||
// Test 5: seq copy (device)
|
||||
if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
@@ -250,7 +250,8 @@ static int eval_message(mtmd_cli_context & ctx, common_chat_msg & msg) {
|
||||
LOG_DBG("formatted_chat.prompt: %s\n", formatted_chat.c_str());
|
||||
|
||||
mtmd_input_text text;
|
||||
text.text = formatted_chat.c_str();
|
||||
text.text = formatted_chat.data();
|
||||
text.text_len = formatted_chat.size();
|
||||
text.add_special = add_bos;
|
||||
text.parse_special = true;
|
||||
|
||||
|
||||
+3
-2
@@ -809,7 +809,7 @@ void mtmd_free(mtmd_context * ctx) {
|
||||
struct mtmd_tokenizer {
|
||||
mtmd_context * ctx;
|
||||
|
||||
std::string input_text;
|
||||
std::string input_text; // note: can contain null bytes; do not use c_str()
|
||||
bool add_special;
|
||||
bool parse_special;
|
||||
const llama_vocab * vocab;
|
||||
@@ -839,9 +839,10 @@ struct mtmd_tokenizer {
|
||||
size_t n_bitmaps) : ctx(ctx) {
|
||||
add_special = text->add_special;
|
||||
parse_special = text->parse_special;
|
||||
input_text = text->text;
|
||||
vocab = ctx->vocab;
|
||||
|
||||
input_text.assign(text->text, text->text_len);
|
||||
|
||||
std::vector<const mtmd_bitmap *> bitmaps(bmps, bmps + n_bitmaps);
|
||||
auto parts_str = split_text(input_text, ctx->media_marker);
|
||||
size_t i_bm = 0;
|
||||
|
||||
@@ -67,6 +67,7 @@ struct mtmd_batch;
|
||||
|
||||
struct mtmd_input_text {
|
||||
const char * text;
|
||||
size_t text_len;
|
||||
bool add_special;
|
||||
bool parse_special;
|
||||
};
|
||||
|
||||
@@ -126,15 +126,15 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`.
|
||||
|
||||
The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:
|
||||
|
||||
- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session.
|
||||
- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` sets the flag the producer polls. One conv maps to at most one live session.
|
||||
- `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`.
|
||||
- `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation.
|
||||
|
||||
The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`.
|
||||
The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, the `server_res_spipe` response base, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager, consumer and the `server_stream_create_spipe` factory stay in the `.cpp`.
|
||||
|
||||
Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
|
||||
Producer side: `server_res_generator` extends `server_res_spipe`, which keeps all spipe logic out of the generic `server_http_res`. `set_req` attaches a producer when the header is present, and the wrapped `next` tees each chunk into the ring before the socket, so a chunk lost to a dead wire is already buffered. While attached, `should_stop` ignores peer disconnect: only a `DELETE` stops generation. On an early peer drop, `on_complete` drains the tail into the ring on the http worker.
|
||||
|
||||
Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free.
|
||||
Lifetime safety: the session holds no back reference to the response, so `spipe` is a plain `unique_ptr` touched only by the http worker. `cancel` raises an atomic the producer polls; the producer finalizes the session from its destructor, which also runs `~server_response_reader::stop()` to cancel the generation at the queue level. A `DELETE` stops work by raising the flag and letting the worker unwind.
|
||||
|
||||
Consumer side: `GET /v1/stream/<conv_id>?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400.
|
||||
|
||||
@@ -235,6 +235,29 @@ That requires `JSON.stringify` when formatted to message content:
|
||||
}
|
||||
```
|
||||
|
||||
Set `stream: true` in the request body to stream a tool's output as it runs, instead of waiting for it to finish. Only certain tools accept this (for ex. `exec_shell_command`);
|
||||
returns 404 if tool doesn't support it.
|
||||
|
||||
Response is SSE stream, one `data: <json>` line per chunk:
|
||||
|
||||
```json
|
||||
{"chunk": "hello\n"}
|
||||
```
|
||||
|
||||
followed by a final event once the tool returns:
|
||||
|
||||
```json
|
||||
{"done": true}
|
||||
```
|
||||
|
||||
or, if `invoke()` threw:
|
||||
|
||||
```json
|
||||
{"done": true, "error": "..."}
|
||||
```
|
||||
|
||||
There is no `[DONE]` sentinel (unlike `/chat/completions`), the stream ends after the `done`
|
||||
|
||||
### Router mode: how child <--> router communicates
|
||||
|
||||
Upon spawning a new child process using `subprocess`, both child and router listen to the stdout/stderr (combined)
|
||||
|
||||
@@ -431,22 +431,70 @@ json server_chat_convert_anthropic_to_oai(const json & body) {
|
||||
std::string tool_use_id = json_value(block, "tool_use_id", std::string());
|
||||
|
||||
auto result_content = json_value(block, "content", json());
|
||||
std::string result_text;
|
||||
if (result_content.is_string()) {
|
||||
result_text = result_content.get<std::string>();
|
||||
tool_results.push_back({
|
||||
{"role", "tool"},
|
||||
{"tool_call_id", tool_use_id},
|
||||
{"content", result_content.get<std::string>()}
|
||||
});
|
||||
} else if (result_content.is_array()) {
|
||||
// Single-pass: build both text and content_parts, decide format at the end
|
||||
std::string result_text;
|
||||
json content_parts = json::array();
|
||||
bool has_images = false;
|
||||
|
||||
for (const auto & c : result_content) {
|
||||
if (json_value(c, "type", std::string()) == "text") {
|
||||
result_text += json_value(c, "text", std::string());
|
||||
std::string c_type = json_value(c, "type", std::string());
|
||||
if (c_type == "text") {
|
||||
std::string text = json_value(c, "text", std::string());
|
||||
result_text += text;
|
||||
content_parts.push_back({
|
||||
{"type", "text"},
|
||||
{"text", text}
|
||||
});
|
||||
} else if (c_type == "image") {
|
||||
has_images = true;
|
||||
json source = json_value(c, "source", json::object());
|
||||
std::string source_type = json_value(source, "type", std::string());
|
||||
if (source_type == "base64") {
|
||||
std::string media_type = json_value(source, "media_type", std::string("image/jpeg"));
|
||||
std::string data = json_value(source, "data", std::string());
|
||||
std::string url = "data:" + media_type + ";base64," + data;
|
||||
content_parts.push_back({
|
||||
{"type", "image_url"},
|
||||
{"image_url", {{"url", url}}}
|
||||
});
|
||||
} else if (source_type == "url") {
|
||||
content_parts.push_back({
|
||||
{"type", "image_url"},
|
||||
{"image_url", {{"url", json_value(source, "url", std::string())}}}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tool_results.push_back({
|
||||
{"role", "tool"},
|
||||
{"tool_call_id", tool_use_id},
|
||||
{"content", result_text}
|
||||
});
|
||||
if (!has_images) {
|
||||
// Text-only: collapse to a plain string for maximum compatibility
|
||||
tool_results.push_back({
|
||||
{"role", "tool"},
|
||||
{"tool_call_id", tool_use_id},
|
||||
{"content", result_text}
|
||||
});
|
||||
} else {
|
||||
// Mixed or image-only: use array content parts (OpenAI multimodal tool format)
|
||||
tool_results.push_back({
|
||||
{"role", "tool"},
|
||||
{"tool_call_id", tool_use_id},
|
||||
{"content", content_parts}
|
||||
});
|
||||
}
|
||||
} else {
|
||||
tool_results.push_back({
|
||||
{"role", "tool"},
|
||||
{"tool_call_id", tool_use_id},
|
||||
{"content", ""}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -705,7 +705,8 @@ server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & promp
|
||||
std::vector<server_tokens> inputs;
|
||||
// multimodal
|
||||
mtmd_input_text inp_txt = {
|
||||
prompt.c_str(),
|
||||
prompt.data(),
|
||||
prompt.size(),
|
||||
/* add_special */ true,
|
||||
/* parse_special */ true,
|
||||
};
|
||||
@@ -1116,7 +1117,8 @@ json oaicompat_chat_params_parse(
|
||||
|
||||
// Reasoning budget: pass parameters through to sampling layer
|
||||
{
|
||||
int reasoning_budget = json_value(body, "thinking_budget_tokens", -1);
|
||||
int reasoning_budget = json_value(body, "reasoning_budget_tokens",
|
||||
json_value(body, "thinking_budget_tokens", -1));
|
||||
if (reasoning_budget == -1) {
|
||||
reasoning_budget = opt.reasoning_budget;
|
||||
}
|
||||
@@ -1125,7 +1127,7 @@ json oaicompat_chat_params_parse(
|
||||
llama_params["reasoning_budget_tokens"] = reasoning_budget;
|
||||
llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag;
|
||||
llama_params["reasoning_budget_end_tag"] = chat_params.thinking_end_tag;
|
||||
llama_params["reasoning_budget_message"] = opt.reasoning_budget_message;
|
||||
llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message);
|
||||
llama_params["reasoning_control"] = json_value(body, "reasoning_control", false);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2290,6 +2290,24 @@ private:
|
||||
|
||||
// n_tokens_cur: the number of tokens added to the batch for the current slot
|
||||
void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) {
|
||||
const int id_task = slot.task->id;
|
||||
|
||||
// evict checkpoints within min-step of a previous checkpoint, unless they were
|
||||
// created by the current task
|
||||
int64_t last = -1;
|
||||
for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end(); ) {
|
||||
if (it->id_task != id_task && last >= 0 && it->n_tokens <= last + params_base.checkpoint_min_step) {
|
||||
SLT_TRC(slot, "erasing context checkpoint too close to an earlier one (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
|
||||
it->pos_min, it->pos_max, it->n_tokens, (float) it->size() / 1024 / 1024);
|
||||
|
||||
it = slot.prompt.checkpoints.erase(it);
|
||||
continue;
|
||||
}
|
||||
|
||||
last = it->n_tokens;
|
||||
++it;
|
||||
}
|
||||
|
||||
while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) {
|
||||
// make room for the new checkpoint, if needed
|
||||
const auto & cur = slot.prompt.checkpoints.front();
|
||||
@@ -2302,6 +2320,8 @@ private:
|
||||
|
||||
auto & cur = slot.prompt.checkpoints.emplace_back();
|
||||
|
||||
cur.id_task = id_task;
|
||||
|
||||
// [TAG_CHECKPOINTS_FIX_POS_MIN]
|
||||
// TODO: here we incorrectly deterimne that the saved checkpoint data covers the [pos_min, pos_max] range
|
||||
// this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225
|
||||
@@ -3511,7 +3531,10 @@ private:
|
||||
do_checkpoint = do_checkpoint && !has_mtmd;
|
||||
|
||||
// no need to create checkpoints that are too close together, unless it's the last user message
|
||||
do_checkpoint = do_checkpoint && (slot.prompt.checkpoints.empty() || is_last_user_message || n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step);
|
||||
do_checkpoint = do_checkpoint && (
|
||||
slot.prompt.checkpoints.empty() ||
|
||||
is_last_user_message || near_prompt_end ||
|
||||
n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step);
|
||||
SLT_DBG(slot, "main/do_checkpoint = %s, pos_min = %d, pos_max = %d\n", do_checkpoint ? "yes" : "no", pos_min, pos_max);
|
||||
|
||||
// note: we create the checkpoint before calling llama_decode(), so the current batch is not
|
||||
@@ -3979,11 +4002,9 @@ server_context_meta server_context::get_meta() const {
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
|
||||
// generator-like API for HTTP response generation
|
||||
// may have bypass_sleep = true if the task does not use ctx_server
|
||||
struct server_res_generator : server_http_res {
|
||||
struct server_res_generator : server_res_spipe {
|
||||
server_response_reader rd;
|
||||
server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false)
|
||||
: rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) {
|
||||
@@ -3993,15 +4014,6 @@ struct server_res_generator : server_http_res {
|
||||
queue_tasks.wait_until_no_sleep();
|
||||
}
|
||||
}
|
||||
~server_res_generator() override {
|
||||
// cleanup() must run while rd is still alive (rd is destroyed after this body returns)
|
||||
if (spipe) {
|
||||
spipe->cleanup();
|
||||
}
|
||||
}
|
||||
void stop() override {
|
||||
rd.stop();
|
||||
}
|
||||
void ok(const json & response_data) {
|
||||
status = 200;
|
||||
data = safe_json_to_str(response_data);
|
||||
@@ -4039,6 +4051,8 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
|
||||
auto & rd = res->rd;
|
||||
auto & params = this->params;
|
||||
|
||||
res->set_req(&req); // will also set spipe if needed
|
||||
|
||||
int32_t sse_ping_interval = params.sse_ping_interval;
|
||||
|
||||
try {
|
||||
@@ -4181,7 +4195,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
|
||||
}
|
||||
res->status = 200;
|
||||
res->content_type = "text/event-stream";
|
||||
res->next = [res_this = res.get(), res_type, sse_ping_interval, &req](std::string & output) -> bool {
|
||||
res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool {
|
||||
static auto format_error = [](task_response_type res_type, const json & res_json) {
|
||||
if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) {
|
||||
return format_anthropic_sse({
|
||||
@@ -4193,7 +4207,9 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
|
||||
}
|
||||
};
|
||||
|
||||
auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop);
|
||||
auto effective_should_stop = [&res_this]() {
|
||||
return res_this->should_stop();
|
||||
};
|
||||
|
||||
try {
|
||||
if (effective_should_stop()) {
|
||||
@@ -4284,13 +4300,9 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
|
||||
// terminate on exception
|
||||
return false;
|
||||
}
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
// attach a producer pipe to the response when X-Conversation-Id is present.
|
||||
// the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
|
||||
server_stream_session_attach_pipe(*res, req.headers);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
#include "common.h"
|
||||
#include "http.h"
|
||||
#include "server-http.h"
|
||||
#include "server-stream.h"
|
||||
#include "server-common.h"
|
||||
#include "ui.h"
|
||||
|
||||
@@ -530,33 +529,20 @@ static void process_handler_response(server_http_req_ptr && request, server_http
|
||||
std::string chunk;
|
||||
const bool has_next = response->next(chunk);
|
||||
if (!chunk.empty()) {
|
||||
// mirror into the ring buffer first, the session must reflect every SSE chunk
|
||||
// whether or not the wire write below succeeds
|
||||
if (response->spipe) {
|
||||
response->spipe->write(chunk.data(), chunk.size());
|
||||
}
|
||||
if (!sink.write(chunk.data(), chunk.size())) {
|
||||
// peer is gone, stop the wire path here
|
||||
return false;
|
||||
}
|
||||
SRV_DBG("http: streamed chunk: %s\n", chunk.c_str());
|
||||
}
|
||||
if (!has_next) {
|
||||
// producer reached its natural end on the wire, a later close() skips the drain
|
||||
if (response->spipe) {
|
||||
response->spipe->done();
|
||||
}
|
||||
sink.done();
|
||||
SRV_DBG("%s", "http: stream ended\n");
|
||||
}
|
||||
return has_next;
|
||||
};
|
||||
const auto on_complete = [request = q_ptr, response = r_ptr](bool) mutable {
|
||||
// on a dropped peer, close() drains the rest of the generation into the ring buffer
|
||||
if (response->spipe) {
|
||||
response->spipe->close();
|
||||
}
|
||||
response.reset(); // spipe destructor finalizes the session if attached
|
||||
response->on_complete();
|
||||
response.reset();
|
||||
request.reset();
|
||||
};
|
||||
res.set_chunked_content_provider(content_type, chunked_content_provider, on_complete);
|
||||
@@ -564,6 +550,7 @@ static void process_handler_response(server_http_req_ptr && request, server_http
|
||||
res.status = response->status;
|
||||
set_headers(res, response->headers);
|
||||
res.set_content(response->data, response->content_type);
|
||||
response->on_complete();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
#include <unordered_map>
|
||||
|
||||
struct common_params;
|
||||
struct stream_pipe_producer; // defined in server-stream.h
|
||||
|
||||
// generator-like API for HTTP response generation
|
||||
// this object response with one of the 2 modes:
|
||||
@@ -25,19 +24,13 @@ struct server_http_res {
|
||||
std::string data;
|
||||
std::map<std::string, std::string> headers;
|
||||
|
||||
// if set, the stream survives a client disconnect: the producer pipe keeps draining into the
|
||||
// ring buffer and finalizes the session on destruction, so no explicit on_stream_end is needed.
|
||||
// shared_ptr (not unique_ptr) so the forward-declared type is safe to delete here.
|
||||
std::shared_ptr<stream_pipe_producer> spipe;
|
||||
|
||||
std::function<bool(std::string &)> next = nullptr;
|
||||
bool is_stream() const {
|
||||
return next != nullptr;
|
||||
}
|
||||
|
||||
// called when the session is cancelled (e.g. DELETE /v1/stream/<conv_id>).
|
||||
// server_res_generator overrides this to stop its reader; the default is a no-op.
|
||||
virtual void stop() {}
|
||||
// fired before req and res are destroyed
|
||||
virtual void on_complete() {}
|
||||
|
||||
virtual ~server_http_res() = default;
|
||||
};
|
||||
|
||||
@@ -219,13 +219,14 @@ void server_model_meta::update_caps() {
|
||||
"LLAMA_ARG_MODEL_URL",
|
||||
"LLAMA_ARG_MMPROJ",
|
||||
"LLAMA_ARG_MMPROJ_URL",
|
||||
"LLAMA_ARG_MMPROJ_AUTO",
|
||||
"LLAMA_ARG_HF_REPO",
|
||||
"LLAMA_ARG_HF_REPO_FILE",
|
||||
});
|
||||
params.offline = true;
|
||||
common_models_handler handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
|
||||
common_models_handler_apply(handler, params); // note: this won't download the model because offline=true
|
||||
if (params.mmproj.path.empty()) {
|
||||
if (params.no_mmproj || params.mmproj.path.empty()) {
|
||||
multimodal = { false, false };
|
||||
} else {
|
||||
multimodal = mtmd_get_cap_from_file(params.mmproj.path.c_str());
|
||||
|
||||
@@ -568,10 +568,16 @@ static void handle_with_catch(const char * name, std::function<void()> func) {
|
||||
}
|
||||
}
|
||||
|
||||
// treat a null value as absent so clients can send null to request the server default
|
||||
static bool has_value(const json & data, const char * n) {
|
||||
auto it = data.find(n);
|
||||
return it != data.end() && !it->is_null();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void field_num<T>::eval(field_eval_context & ctx, const json & data) {
|
||||
for (const auto & n : name) {
|
||||
if (data.contains(n)) {
|
||||
if (has_value(data, n)) {
|
||||
handle_with_catch(n, [&]() {
|
||||
if (custom_handler) {
|
||||
custom_handler(ctx, data);
|
||||
@@ -593,7 +599,7 @@ void field_num<T>::eval(field_eval_context & ctx, const json & data) {
|
||||
void field_str::eval(field_eval_context & ctx, const json & data) {
|
||||
GGML_ASSERT(custom_handler);
|
||||
for (const auto & n : name) {
|
||||
if (data.contains(n)) {
|
||||
if (has_value(data, n)) {
|
||||
handle_with_catch(n, [&]() {
|
||||
custom_handler(ctx, data);
|
||||
});
|
||||
@@ -604,7 +610,7 @@ void field_str::eval(field_eval_context & ctx, const json & data) {
|
||||
|
||||
void field_bool::eval(field_eval_context & ctx, const json & data) {
|
||||
for (const auto & n : name) {
|
||||
if (data.contains(n)) {
|
||||
if (has_value(data, n)) {
|
||||
handle_with_catch(n, [&]() {
|
||||
if (custom_handler) {
|
||||
custom_handler(ctx, data);
|
||||
@@ -620,7 +626,7 @@ void field_bool::eval(field_eval_context & ctx, const json & data) {
|
||||
void field_json::eval(field_eval_context & ctx, const json & data) {
|
||||
GGML_ASSERT(custom_handler);
|
||||
for (const auto & n : name) {
|
||||
if (data.contains(n)) {
|
||||
if (has_value(data, n)) {
|
||||
handle_with_catch(n, [&]() {
|
||||
custom_handler(ctx, data);
|
||||
});
|
||||
|
||||
@@ -96,8 +96,6 @@ struct stream_session {
|
||||
size_t dropped_prefix() const; // bytes evicted from the front due to cap
|
||||
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
|
||||
|
||||
void set_stop_producer(std::function<void()> fn);
|
||||
|
||||
void cancel();
|
||||
|
||||
private:
|
||||
@@ -109,7 +107,6 @@ private:
|
||||
bool done;
|
||||
std::atomic<bool> cancelled; // polled lock-free by the should_stop closure, no mu
|
||||
int64_t completed_ts;
|
||||
std::function<void()> stop_producer;
|
||||
};
|
||||
stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
|
||||
: conversation_id(std::move(conversation_id_))
|
||||
@@ -217,26 +214,10 @@ int64_t stream_session::completed_at() const {
|
||||
return completed_ts;
|
||||
}
|
||||
|
||||
void stream_session::set_stop_producer(std::function<void()> fn) {
|
||||
std::lock_guard<std::mutex> lock(mu);
|
||||
stop_producer = std::move(fn);
|
||||
}
|
||||
|
||||
void stream_session::cancel() {
|
||||
// flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the
|
||||
// recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
|
||||
// posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
|
||||
// the should_stop closure on both the producer and any HTTP reader polls is_cancelled()
|
||||
// so flipping this is the only signal needed to unwind both sides
|
||||
cancelled.store(true, std::memory_order_release);
|
||||
// copy the hook under the lock then invoke outside, the producer side may grab queue locks
|
||||
// and we do not want to hold our mu across that path
|
||||
std::function<void()> fn;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mu);
|
||||
fn = stop_producer;
|
||||
}
|
||||
if (fn) {
|
||||
fn();
|
||||
}
|
||||
}
|
||||
|
||||
bool stream_session::is_cancelled() const {
|
||||
@@ -325,8 +306,10 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i
|
||||
s = it->second;
|
||||
sessions.erase(it);
|
||||
}
|
||||
// signal the producer side first so the inference is cancelled at the queue level,
|
||||
// then finalize, which wakes any pending HTTP reader and lets the drain exit naturally
|
||||
// cancel first so the producer's on_complete() drain loop and any pending HTTP reader
|
||||
// observe is_cancelled() and stop pulling further output, then finalize to wake readers
|
||||
// blocked in read_from(). note: this does not interrupt the underlying generation itself,
|
||||
// which keeps running to its own natural stop condition (EOS/max_tokens)
|
||||
s->cancel();
|
||||
s->finalize();
|
||||
}
|
||||
@@ -431,65 +414,15 @@ stream_pipe_producer::stream_pipe_producer(stream_session_ptr session)
|
||||
}
|
||||
|
||||
stream_pipe_producer::~stream_pipe_producer() {
|
||||
cleanup();
|
||||
session_->finalize();
|
||||
}
|
||||
|
||||
void stream_pipe_producer::cleanup() {
|
||||
if (!alive_) {
|
||||
return;
|
||||
}
|
||||
alive_->store(false, std::memory_order_release);
|
||||
session_->set_stop_producer(nullptr);
|
||||
alive_.reset();
|
||||
}
|
||||
|
||||
bool stream_pipe_producer::write(const char * data, size_t len) {
|
||||
return session_->append(data, len);
|
||||
}
|
||||
|
||||
void stream_pipe_producer::done() {
|
||||
done_ = true;
|
||||
}
|
||||
|
||||
void stream_pipe_producer::close() {
|
||||
// httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest
|
||||
// of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short
|
||||
if (done_ || session_->is_cancelled()) {
|
||||
SRV_TRC("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n",
|
||||
done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str());
|
||||
return;
|
||||
}
|
||||
SRV_TRC("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str());
|
||||
size_t drained = 0;
|
||||
std::string chunk;
|
||||
while (true) {
|
||||
chunk.clear();
|
||||
bool has_next = res_->next(chunk);
|
||||
if (!chunk.empty()) {
|
||||
write(chunk.data(), chunk.size());
|
||||
drained += chunk.size();
|
||||
}
|
||||
if (!has_next) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
SRV_TRC("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained);
|
||||
}
|
||||
|
||||
std::shared_ptr<stream_pipe_producer> stream_pipe_producer::create(stream_session_ptr session,
|
||||
server_http_res & res) {
|
||||
auto alive = std::make_shared<std::atomic<bool>>(true);
|
||||
auto * res_ptr = &res;
|
||||
session->set_stop_producer([alive, res_ptr]() {
|
||||
if (alive->load(std::memory_order_acquire)) {
|
||||
res_ptr->stop();
|
||||
}
|
||||
});
|
||||
auto pipe = std::shared_ptr<stream_pipe_producer>(new stream_pipe_producer(std::move(session)));
|
||||
pipe->alive_ = std::move(alive);
|
||||
pipe->res_ = res_ptr;
|
||||
return pipe;
|
||||
stream_pipe_producer * stream_pipe_producer::create(stream_session_ptr session) {
|
||||
return new stream_pipe_producer(std::move(session));
|
||||
}
|
||||
|
||||
// stream_pipe_consumer
|
||||
@@ -661,21 +594,68 @@ std::string server_stream_conv_id_from_headers(const std::map<std::string, std::
|
||||
return std::string();
|
||||
}
|
||||
|
||||
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
|
||||
static stream_pipe_producer * server_stream_create_spipe(const std::map<std::string, std::string> & headers) {
|
||||
std::string conversation_id = server_stream_conv_id_from_headers(headers);
|
||||
SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
|
||||
if (conversation_id.empty()) {
|
||||
return;
|
||||
return nullptr;
|
||||
}
|
||||
auto session = g_stream_sessions.create_or_replace(conversation_id);
|
||||
res.spipe = stream_pipe_producer::create(session, res);
|
||||
return stream_pipe_producer::create(session);
|
||||
}
|
||||
|
||||
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
|
||||
return [res, fallback = std::move(fallback)]() -> bool {
|
||||
if (res->spipe) {
|
||||
return res->spipe->is_cancelled();
|
||||
//
|
||||
// server_res_spipe
|
||||
//
|
||||
|
||||
void server_res_spipe::set_req(const server_http_req * req) {
|
||||
this->req = req;
|
||||
// optionally attach spipe to the response when X-Conversation-Id is present
|
||||
spipe.reset(server_stream_create_spipe(req->headers));
|
||||
}
|
||||
|
||||
bool server_res_spipe::conn_alive() {
|
||||
GGML_ASSERT(req != nullptr);
|
||||
return !req->should_stop();
|
||||
}
|
||||
|
||||
bool server_res_spipe::should_stop() {
|
||||
if (spipe) {
|
||||
// note: if DELETE /v1/stream/<conv_id> is called, is_cancelled() will be true
|
||||
return spipe->is_cancelled();
|
||||
} else {
|
||||
return !conn_alive();
|
||||
}
|
||||
}
|
||||
|
||||
void server_res_spipe::on_complete() {
|
||||
if (!spipe || next_finished) {
|
||||
return;
|
||||
}
|
||||
std::string chunk;
|
||||
while (!spipe->is_cancelled()) {
|
||||
chunk.clear();
|
||||
bool has_next = next_orig(chunk);
|
||||
if (!chunk.empty()) {
|
||||
spipe->write(chunk.data(), chunk.size());
|
||||
}
|
||||
return fallback();
|
||||
if (!has_next) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void server_res_spipe::set_next(std::function<bool(std::string &)> next_fn) {
|
||||
next_orig = std::move(next_fn);
|
||||
next = [this](std::string & out) {
|
||||
bool has_next = next_orig(out);
|
||||
if (spipe) {
|
||||
// if spipe is set, tee-style pipe input to both HTTP and spipe
|
||||
spipe->write(out.data(), out.size());
|
||||
}
|
||||
if (!has_next) {
|
||||
next_finished = true;
|
||||
}
|
||||
return has_next;
|
||||
};
|
||||
}
|
||||
|
||||
@@ -30,36 +30,15 @@ protected:
|
||||
|
||||
// producer end: writes chunks into the ring buffer and owns the session lifetime, finalizing it
|
||||
// on destruction.
|
||||
//
|
||||
// lifetime safety: holds a shared_ptr<atomic<bool>> alive also captured by the session's
|
||||
// stop_producer hook. cleanup() sets alive=false and clears the hook; it must run while the
|
||||
// response the hook calls stop() on is still alive. ~server_res_generator() does this explicitly.
|
||||
struct stream_pipe_producer : stream_pipe {
|
||||
~stream_pipe_producer() override;
|
||||
|
||||
bool write(const char * data, size_t len);
|
||||
|
||||
// mark the natural end on the wire so a later close() is a no-op
|
||||
void done();
|
||||
|
||||
// on a peer drop, pump the response next() into the ring buffer until done. runs on the http
|
||||
// worker from on_complete, no-op after done() or cancel
|
||||
void close();
|
||||
|
||||
// disarm the stop hook and drop the alive guard, must run while the response the hook
|
||||
// references is still alive. idempotent, the destructor calls it too
|
||||
void cleanup();
|
||||
|
||||
// res.stop() is invoked when the session is cancelled, the alive guard ensures stop() is not
|
||||
// called after cleanup() has run
|
||||
static std::shared_ptr<stream_pipe_producer> create(stream_session_ptr session, server_http_res & res);
|
||||
static stream_pipe_producer * create(stream_session_ptr session);
|
||||
|
||||
private:
|
||||
explicit stream_pipe_producer(stream_session_ptr session);
|
||||
|
||||
bool done_ = false;
|
||||
std::shared_ptr<std::atomic<bool>> alive_;
|
||||
server_http_res * res_ = nullptr;
|
||||
};
|
||||
|
||||
void server_stream_session_manager_start();
|
||||
@@ -73,10 +52,22 @@ server_http_context::handler_t server_stream_make_delete_handler();
|
||||
// extract the X-Conversation-Id header value (case-insensitive), empty when absent
|
||||
std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
|
||||
|
||||
// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res
|
||||
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
|
||||
// implement tee-style pipe (spipe) for "stream replay" functionality
|
||||
struct server_res_spipe : server_http_res {
|
||||
private:
|
||||
// if set, the stream survives a client disconnect:
|
||||
// connection kept alive, output is forwarded to spipe and reuse later
|
||||
std::unique_ptr<stream_pipe_producer> spipe;
|
||||
// if spipe is set, use this next_orig to implement tee-style pipe
|
||||
std::function<bool(std::string &)> next_orig;
|
||||
const server_http_req * req = nullptr;
|
||||
// set once next_orig reports no more data, so on_complete() doesn't re-drain a finished stream
|
||||
bool next_finished = false;
|
||||
|
||||
// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit
|
||||
// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it
|
||||
// delegates to fallback, the legacy non-resumable flow
|
||||
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
|
||||
public:
|
||||
void set_req(const server_http_req * req);
|
||||
bool conn_alive();
|
||||
bool should_stop();
|
||||
void on_complete() override;
|
||||
void set_next(std::function<bool(std::string &)> next_fn);
|
||||
};
|
||||
|
||||
+149
-23
@@ -12,6 +12,7 @@
|
||||
#include <climits>
|
||||
#include <algorithm>
|
||||
#include <unordered_set>
|
||||
#include <functional>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
@@ -51,7 +52,13 @@ public:
|
||||
virtual bool write_file(const std::string & path, const std::string & content) const = 0;
|
||||
// paths relative to `base`, '/'-separated; sets `err` if `base` isn't a directory
|
||||
virtual std::vector<std::string> list_files(const std::string & base, std::string & err) const = 0;
|
||||
virtual exec_result run(const std::vector<std::string> & args, size_t max_output, int timeout_secs) const = 0;
|
||||
// on_chunk, if set, is called with each chunk of output as it is read (before truncation cuts in);
|
||||
// returning false terminates the process early (e.g. the client disconnected)
|
||||
virtual exec_result run(
|
||||
const std::vector<std::string> & args,
|
||||
size_t max_output,
|
||||
int timeout_secs,
|
||||
const std::function<bool(const std::string &)> & on_chunk = nullptr) const = 0;
|
||||
};
|
||||
|
||||
class tools_io_basic : public tools_io {
|
||||
@@ -123,7 +130,11 @@ public:
|
||||
return list_files_fallback(base);
|
||||
}
|
||||
|
||||
exec_result run(const std::vector<std::string> & args, size_t max_output, int timeout_secs) const override {
|
||||
exec_result run(
|
||||
const std::vector<std::string> & args,
|
||||
size_t max_output,
|
||||
int timeout_secs,
|
||||
const std::function<bool(const std::string &)> & on_chunk = nullptr) const override {
|
||||
exec_result res;
|
||||
|
||||
subprocess_s proc;
|
||||
@@ -164,8 +175,14 @@ public:
|
||||
size_t len = strlen(buf);
|
||||
if (output.size() + len <= max_output) {
|
||||
output.append(buf, len);
|
||||
if (on_chunk && !on_chunk(std::string(buf, len))) {
|
||||
subprocess_terminate(&proc);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
output.append(buf, max_output - output.size());
|
||||
size_t remaining = max_output - output.size();
|
||||
output.append(buf, remaining);
|
||||
if (on_chunk && remaining > 0) on_chunk(std::string(buf, remaining));
|
||||
truncated = true;
|
||||
}
|
||||
}
|
||||
@@ -287,7 +304,7 @@ struct server_tool_read_file : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params) const override {
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string path = params.at("path").get<std::string>();
|
||||
int start_line = json_value(params, "start_line", 1);
|
||||
int end_line = json_value(params, "end_line", -1); // -1 = no limit
|
||||
@@ -376,7 +393,7 @@ struct server_tool_file_glob_search : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params) const override {
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string base = params.at("path").get<std::string>();
|
||||
std::string include = json_value(params, "include", std::string("**"));
|
||||
std::string exclude = json_value(params, "exclude", std::string(""));
|
||||
@@ -457,7 +474,7 @@ struct server_tool_grep_search : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params) const override {
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string path = params.at("path").get<std::string>();
|
||||
std::string pat_str = params.at("pattern").get<std::string>();
|
||||
std::string include = json_value(params, "include", std::string("**"));
|
||||
@@ -577,6 +594,7 @@ struct server_tool_exec_shell_command : server_tool {
|
||||
name = "exec_shell_command";
|
||||
display_name = "Execute shell command";
|
||||
permission_write = true;
|
||||
support_stream = true;
|
||||
}
|
||||
|
||||
json get_definition() const override {
|
||||
@@ -598,7 +616,7 @@ struct server_tool_exec_shell_command : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params) const override {
|
||||
json invoke(json params, server_tool::stream * st) const override {
|
||||
std::string command = params.at("command").get<std::string>();
|
||||
int timeout = json_value(params, "timeout", 10);
|
||||
size_t max_output = (size_t) json_value(params, "max_output_size", (int) SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_OUTPUT_SIZE);
|
||||
@@ -612,7 +630,24 @@ struct server_tool_exec_shell_command : server_tool {
|
||||
std::vector<std::string> args = {"sh", "-c", command};
|
||||
#endif
|
||||
|
||||
auto io = make_tools_io(params);
|
||||
auto io = make_tools_io(params);
|
||||
|
||||
if (st) {
|
||||
auto res = io->run(args, max_output, timeout, [st](const std::string & chunk) {
|
||||
st->push(chunk);
|
||||
return !st->alive || st->alive();
|
||||
});
|
||||
if (st->alive && !st->alive()) {
|
||||
return json();
|
||||
}
|
||||
std::string tail = string_format("\n[exit code: %d]", res.exit_code);
|
||||
if (res.timed_out) {
|
||||
tail += " [exit due to timed out]";
|
||||
}
|
||||
st->push(tail);
|
||||
return json();
|
||||
}
|
||||
|
||||
auto res = io->run(args, max_output, timeout);
|
||||
|
||||
std::string text_output = res.output;
|
||||
@@ -654,7 +689,7 @@ struct server_tool_write_file : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params) const override {
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string path = params.at("path").get<std::string>();
|
||||
std::string content = params.at("content").get<std::string>();
|
||||
|
||||
@@ -710,7 +745,7 @@ struct server_tool_edit_file : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params) const override {
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string path = params.at("path").get<std::string>();
|
||||
const json & edits_json = params.at("edits");
|
||||
|
||||
@@ -1018,7 +1053,7 @@ struct server_tool_get_datetime : server_tool {
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json) const override {
|
||||
json invoke(json, server_tool::stream *) const override {
|
||||
auto now = std::chrono::system_clock::now();
|
||||
auto time = std::chrono::system_clock::to_time_t(now);
|
||||
|
||||
@@ -1026,6 +1061,59 @@ struct server_tool_get_datetime : server_tool {
|
||||
}
|
||||
};
|
||||
|
||||
struct server_tool_stream_result : server_task_result {
|
||||
std::string chunk;
|
||||
bool done = false;
|
||||
std::string error_msg;
|
||||
|
||||
json to_json() override {
|
||||
if (!done) {
|
||||
return {{"chunk", chunk}};
|
||||
} else {
|
||||
json result = {{"done", true}};
|
||||
if (!error_msg.empty()) {
|
||||
result["error"] = error_msg;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void server_tool::stream::push(const std::string & chunk) {
|
||||
if (chunk.empty()) return;
|
||||
auto r = std::make_unique<server_tool_stream_result>();
|
||||
r->id = id;
|
||||
r->chunk = chunk;
|
||||
qr.send(std::move(r));
|
||||
}
|
||||
|
||||
struct server_tools_res : server_http_res {
|
||||
std::thread worker;
|
||||
server_response * qr = nullptr; // set only for streaming responses
|
||||
int id = -1;
|
||||
|
||||
~server_tools_res() override {
|
||||
if (worker.joinable()) {
|
||||
worker.join();
|
||||
}
|
||||
if (qr) {
|
||||
qr->remove_waiting_task_id(id);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
static server_tool & find_tool(std::vector<std::unique_ptr<server_tool>> & tools, const std::string & name, bool require_stream) {
|
||||
for (auto & t : tools) {
|
||||
if (t->name == name) {
|
||||
if (require_stream && !t->support_stream) {
|
||||
throw std::invalid_argument(string_format("tool \"%s\" does not support stream = true", name.c_str()));
|
||||
}
|
||||
return *t;
|
||||
}
|
||||
}
|
||||
throw std::invalid_argument(string_format("unknown tool \"%s\"", name.c_str()));
|
||||
}
|
||||
|
||||
//
|
||||
// public API
|
||||
//
|
||||
@@ -1090,16 +1178,63 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
|
||||
};
|
||||
|
||||
handle_post = [this](const server_http_req & req) -> server_http_res_ptr {
|
||||
auto res = std::make_unique<server_http_res>();
|
||||
auto res = std::make_unique<server_tools_res>();
|
||||
try {
|
||||
json body = json::parse(req.body);
|
||||
std::string tool_name = body.at("tool").get<std::string>();
|
||||
json params = body.value("params", json::object());
|
||||
json result = invoke(tool_name, params);
|
||||
res->data = safe_json_to_str(result);
|
||||
bool stream = body.value("stream", false);
|
||||
|
||||
server_tool & tool = find_tool(tools, tool_name, stream);
|
||||
|
||||
if (stream) {
|
||||
int id = res_id.fetch_add(1);
|
||||
queue_res.add_waiting_task_id(id);
|
||||
res->qr = &queue_res;
|
||||
res->id = id;
|
||||
|
||||
res->worker = std::thread([this, id, &req, &tool, params]() mutable {
|
||||
server_tool::stream st{queue_res, id, [&req]() {
|
||||
return !req.should_stop();
|
||||
}};
|
||||
|
||||
auto done = std::make_unique<server_tool_stream_result>();
|
||||
try {
|
||||
tool.invoke(params, &st);
|
||||
} catch (const std::exception & e) {
|
||||
done->error_msg = e.what();
|
||||
} catch (...) {
|
||||
done->error_msg = "An unknown error occurred";
|
||||
}
|
||||
done->id = st.id;
|
||||
done->done = true;
|
||||
st.qr.send(std::move(done));
|
||||
});
|
||||
|
||||
res->content_type = "text/event-stream";
|
||||
res->status = 200;
|
||||
res->next = [this, id](std::string & output) -> bool {
|
||||
auto result = queue_res.recv(id);
|
||||
auto * r = dynamic_cast<server_tool_stream_result *>(result.get());
|
||||
GGML_ASSERT(r != nullptr);
|
||||
output = "data: " + safe_json_to_str(r->to_json()) + "\n\n";
|
||||
if (r->done) {
|
||||
queue_res.remove_waiting_task_id(id);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
} else {
|
||||
json result = tool.invoke(params, nullptr);
|
||||
res->status = 200;
|
||||
res->data = safe_json_to_str(result);
|
||||
}
|
||||
} catch (const json::exception & e) {
|
||||
res->status = 400;
|
||||
res->data = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
|
||||
} catch (const std::invalid_argument & e) {
|
||||
res->status = 404;
|
||||
res->data = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
|
||||
} catch (const std::exception & e) {
|
||||
SRV_ERR("got exception: %s\n", e.what());
|
||||
res->status = 500;
|
||||
@@ -1108,12 +1243,3 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
|
||||
return res;
|
||||
};
|
||||
}
|
||||
|
||||
json server_tools::invoke(const std::string & name, const json & params) {
|
||||
for (auto & t : tools) {
|
||||
if (t->name == name) {
|
||||
return t->invoke(params);
|
||||
}
|
||||
}
|
||||
return {{"error", "unknown tool: " + name}};
|
||||
}
|
||||
|
||||
@@ -2,15 +2,27 @@
|
||||
|
||||
#include "server-common.h"
|
||||
#include "server-http.h"
|
||||
#include "server-queue.h"
|
||||
|
||||
#include <atomic>
|
||||
#include <functional>
|
||||
|
||||
struct server_tool {
|
||||
std::string name;
|
||||
std::string display_name;
|
||||
bool permission_write = false;
|
||||
bool support_stream = false; // if true, output can be streamed
|
||||
|
||||
virtual ~server_tool() = default;
|
||||
virtual json get_definition() const = 0;
|
||||
virtual json invoke(json params) const = 0;
|
||||
|
||||
struct stream {
|
||||
server_response & qr;
|
||||
int id;
|
||||
std::function<bool()> alive;
|
||||
void push(const std::string & chunk);
|
||||
};
|
||||
virtual json invoke(json params, stream * st = nullptr) const = 0;
|
||||
|
||||
json to_json() const;
|
||||
};
|
||||
@@ -18,8 +30,11 @@ struct server_tool {
|
||||
struct server_tools {
|
||||
std::vector<std::unique_ptr<server_tool>> tools;
|
||||
|
||||
// for streaming
|
||||
server_response queue_res;
|
||||
std::atomic<int> res_id{0};
|
||||
|
||||
void setup(const std::vector<std::string> & enabled_tools);
|
||||
json invoke(const std::string & name, const json & params);
|
||||
|
||||
server_http_context::handler_t handle_get;
|
||||
server_http_context::handler_t handle_post;
|
||||
|
||||
@@ -402,6 +402,65 @@ def test_anthropic_tool_result_with_text():
|
||||
assert len(res.body["content"]) > 0
|
||||
|
||||
|
||||
def test_anthropic_tool_result_with_image():
|
||||
"""Test tool result containing mixed text and image blocks
|
||||
|
||||
Verifies that image blocks inside Anthropic tool_result content are
|
||||
properly converted to OpenAI image_url format rather than being
|
||||
silently dropped. With a non-multimodal model, the converted image
|
||||
triggers a clear error message instead of being ignored.
|
||||
"""
|
||||
server.jinja = True
|
||||
server.start()
|
||||
|
||||
# Small 1x1 red PNG image in base64 (same as vision tests)
|
||||
red_pixel_png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
|
||||
|
||||
res = server.make_request("POST", "/v1/messages", data={
|
||||
"model": "test",
|
||||
"max_tokens": 100,
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is in this image?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "tool_1",
|
||||
"name": "read",
|
||||
"input": {"file": "test.png"}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tool_1",
|
||||
"content": [
|
||||
{"type": "text", "text": "File: test.png"},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": red_pixel_png
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
})
|
||||
|
||||
# Without the fix, image block would cause "unsupported content[].type"
|
||||
# With the fix, image is converted to image_url but tinyllama doesn't support images
|
||||
assert res.status_code == 500
|
||||
assert "image input is not supported" in res.body.get("error", {}).get("message", "").lower()
|
||||
|
||||
|
||||
def test_anthropic_tool_result_error():
|
||||
"""Test tool result with error flag"""
|
||||
server.jinja = True
|
||||
|
||||
@@ -105,6 +105,24 @@ def test_tools_builtin_edit_file_rejects_non_unique_old_text():
|
||||
os.remove(log_path)
|
||||
|
||||
|
||||
def test_tools_builtin_exec_shell_command_stream():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
events = list(server.make_stream_request("POST", "/tools", data={
|
||||
"tool": "exec_shell_command",
|
||||
"params": {"command": "echo hello"},
|
||||
"stream": True,
|
||||
}))
|
||||
|
||||
assert len(events) >= 2
|
||||
assert events[-1]["done"] is True
|
||||
assert not events[-1].get("error")
|
||||
chunks = "".join(e["chunk"] for e in events[:-1])
|
||||
assert "hello" in chunks
|
||||
assert "[exit code: 0]" in chunks
|
||||
|
||||
|
||||
def test_tools_builtin_edit_file_rejects_overlapping_edits():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
+1
-1
@@ -17,7 +17,7 @@
|
||||
let { onMcpSettingsClick }: Props = $props();
|
||||
|
||||
let mcpSearchQuery = $state('');
|
||||
let allMcpServers = $derived(mcpStore.getServersSorted());
|
||||
let allMcpServers = $derived(mcpStore.getServers());
|
||||
let mcpServers = $derived(mcpStore.visibleMcpServers);
|
||||
let hasMcpServers = $derived(mcpServers.length > 0);
|
||||
// let hasAnyMcpServers = $derived(allMcpServers.length > 0);
|
||||
|
||||
+1
-1
@@ -10,7 +10,7 @@
|
||||
import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
|
||||
|
||||
const toolsPanel = useToolsPanel();
|
||||
const hasMcpServersAvailable = $derived(mcpStore.getServersSorted().length > 0);
|
||||
const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0);
|
||||
</script>
|
||||
|
||||
<DropdownMenu.Sub onOpenChange={(open) => open && toolsPanel.handleOpen()}>
|
||||
|
||||
+1
-1
@@ -322,7 +322,7 @@
|
||||
}
|
||||
|
||||
let filteredPrompts = $derived.by(() => {
|
||||
const sortedServers = mcpStore.getServersSorted();
|
||||
const sortedServers = mcpStore.getServers();
|
||||
const serverOrderMap = new Map(sortedServers.map((server, index) => [server.id, index]));
|
||||
|
||||
const sortedPrompts = [...prompts].sort((a, b) => {
|
||||
|
||||
+1
-1
@@ -138,7 +138,7 @@
|
||||
}
|
||||
|
||||
let filteredResources = $derived.by(() => {
|
||||
const sortedServers = mcpStore.getServersSorted();
|
||||
const sortedServers = mcpStore.getServers();
|
||||
const serverOrderMap = new Map(sortedServers.map((server, index) => [server.id, index]));
|
||||
|
||||
const sortedResources = [...resources].sort((a, b) => {
|
||||
|
||||
@@ -1,210 +0,0 @@
|
||||
<script lang="ts">
|
||||
import { Button } from '$lib/components/ui/button';
|
||||
import * as Card from '$lib/components/ui/card';
|
||||
import * as Dialog from '$lib/components/ui/dialog';
|
||||
import { fly } from 'svelte/transition';
|
||||
import { McpServerCardCompact, McpServerForm } from '$lib/components/app/mcp';
|
||||
import { RECOMMENDED_MCP_SERVERS, SETTINGS_KEYS } from '$lib/constants';
|
||||
import { conversationsStore } from '$lib/stores/conversations.svelte';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
import { settingsStore } from '$lib/stores/settings.svelte';
|
||||
import { uuid } from '$lib/utils';
|
||||
import { MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, MCP_SERVER_ID_PREFIX } from '$lib/constants';
|
||||
import type { MCPServerSettingsEntry } from '$lib/types';
|
||||
import { Plus } from '@lucide/svelte';
|
||||
|
||||
interface Props {
|
||||
open: boolean;
|
||||
onOpenChange?: (open: boolean) => void;
|
||||
}
|
||||
|
||||
let { open = $bindable(), onOpenChange }: Props = $props();
|
||||
|
||||
let selected = $state<Record<string, boolean>>(
|
||||
Object.fromEntries(RECOMMENDED_MCP_SERVERS.map((server) => [server.id, false]))
|
||||
);
|
||||
|
||||
let addedServers = $state<MCPServerSettingsEntry[]>([]);
|
||||
let didAddAny = $state(false);
|
||||
|
||||
let selectedRecommendedCount = $derived.by(
|
||||
() => RECOMMENDED_MCP_SERVERS.filter((server) => selected[server.id]).length
|
||||
);
|
||||
|
||||
let footerLabel = $derived.by(() => {
|
||||
const recommended = selectedRecommendedCount;
|
||||
const custom = addedServers.length;
|
||||
const total = recommended + custom;
|
||||
|
||||
if (total === 0) return 'Continue';
|
||||
if (recommended === 0) return custom === 1 ? 'Add server' : `Add ${custom} servers`;
|
||||
if (custom === 0) return recommended === 1 ? 'Add server' : `Add ${recommended} servers`;
|
||||
return `Add ${recommended} servers and ${custom} custom`;
|
||||
});
|
||||
|
||||
let showAddForm = $state(false);
|
||||
let newServerUrl = $state('');
|
||||
let newServerHeaders = $state('');
|
||||
let newServerUrlError = $derived.by(() => {
|
||||
if (!newServerUrl.trim()) return 'URL is required';
|
||||
try {
|
||||
new URL(newServerUrl);
|
||||
|
||||
return null;
|
||||
} catch {
|
||||
return 'Invalid URL format';
|
||||
}
|
||||
});
|
||||
|
||||
function handleOpenChange(value: boolean) {
|
||||
if (!value) {
|
||||
showAddForm = false;
|
||||
newServerUrl = '';
|
||||
newServerHeaders = '';
|
||||
|
||||
if (!didAddAny) {
|
||||
settingsStore.updateConfig(SETTINGS_KEYS.MCP_SERVERS, []);
|
||||
}
|
||||
|
||||
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
|
||||
addedServers = [];
|
||||
didAddAny = false;
|
||||
}
|
||||
open = value;
|
||||
onOpenChange?.(value);
|
||||
}
|
||||
|
||||
function resetAddForm() {
|
||||
showAddForm = false;
|
||||
newServerUrl = '';
|
||||
newServerHeaders = '';
|
||||
}
|
||||
|
||||
function enableSelected() {
|
||||
didAddAny = true;
|
||||
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
|
||||
|
||||
for (const server of RECOMMENDED_MCP_SERVERS) {
|
||||
if (selected[server.id]) {
|
||||
const existing = mcpStore.getServerById(server.id);
|
||||
if (existing) {
|
||||
mcpStore.updateServer(server.id, { enabled: true });
|
||||
} else {
|
||||
mcpStore.addServer({
|
||||
id: server.id,
|
||||
enabled: true,
|
||||
url: server.url,
|
||||
name: server.name
|
||||
});
|
||||
}
|
||||
conversationsStore.setMcpServerOverride(server.id, true);
|
||||
}
|
||||
}
|
||||
handleOpenChange(false);
|
||||
}
|
||||
|
||||
function saveNewServer() {
|
||||
if (newServerUrlError) return;
|
||||
|
||||
didAddAny = true;
|
||||
|
||||
const newServerId = uuid() ?? `${MCP_SERVER_ID_PREFIX}-${Date.now()}`;
|
||||
|
||||
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
|
||||
|
||||
const newServer = mcpStore.addServer({
|
||||
id: newServerId,
|
||||
enabled: true,
|
||||
url: newServerUrl.trim(),
|
||||
headers: newServerHeaders.trim() || undefined
|
||||
});
|
||||
|
||||
conversationsStore.setMcpServerOverride(newServerId, true);
|
||||
|
||||
if (newServer) {
|
||||
addedServers = [...addedServers, newServer];
|
||||
}
|
||||
|
||||
resetAddForm();
|
||||
}
|
||||
</script>
|
||||
|
||||
<Dialog.Root bind:open onOpenChange={handleOpenChange}>
|
||||
<Dialog.Content class="sm:max-w-lg">
|
||||
<Dialog.Header>
|
||||
<Dialog.Title>Do more with MCP</Dialog.Title>
|
||||
<Dialog.Description>
|
||||
Power-up your experience by adding tools, resources and more capabilities provided by MCP
|
||||
servers.
|
||||
</Dialog.Description>
|
||||
</Dialog.Header>
|
||||
|
||||
<div class="max-h-[60vh] space-y-4 overflow-y-auto py-4" in:fly={{ y: 16, duration: 300 }}>
|
||||
<h3 class="text-sm font-semibold">Quickly get started with</h3>
|
||||
|
||||
{#each RECOMMENDED_MCP_SERVERS as server (server.id)}
|
||||
<McpServerCardCompact
|
||||
{server}
|
||||
enabled={selected[server.id]}
|
||||
onToggle={(enabled) => (selected[server.id] = enabled)}
|
||||
/>
|
||||
{/each}
|
||||
|
||||
{#if addedServers.length > 0}
|
||||
{#each addedServers as server (server.id)}
|
||||
<McpServerCardCompact {server} enabled={true} />
|
||||
{/each}
|
||||
{/if}
|
||||
|
||||
{#if showAddForm}
|
||||
<Card.Root class="gap-3! bg-muted/30 p-4">
|
||||
<McpServerForm
|
||||
url={newServerUrl}
|
||||
headers={newServerHeaders}
|
||||
onUrlChange={(v) => (newServerUrl = v)}
|
||||
onHeadersChange={(v) => (newServerHeaders = v)}
|
||||
urlError={newServerUrl ? newServerUrlError : null}
|
||||
id="recommendation-new-server"
|
||||
/>
|
||||
|
||||
<div class="flex justify-end gap-2 pt-2">
|
||||
<Button variant="secondary" size="sm" onclick={resetAddForm}>Cancel</Button>
|
||||
|
||||
<Button
|
||||
variant="default"
|
||||
size="sm"
|
||||
onclick={saveNewServer}
|
||||
disabled={!!newServerUrlError}
|
||||
aria-label="Save"
|
||||
>
|
||||
Add
|
||||
</Button>
|
||||
</div>
|
||||
</Card.Root>
|
||||
{:else}
|
||||
<Card.Root class="gap-0 border-dashed bg-muted/30 p-0 transition-colors hover:bg-muted/50">
|
||||
<button
|
||||
type="button"
|
||||
class="flex w-full items-center justify-center gap-2 rounded-lg p-6 text-sm text-muted-foreground transition-colors hover:text-foreground"
|
||||
onclick={() => (showAddForm = true)}
|
||||
aria-label="Add your own MCP server"
|
||||
>
|
||||
<Plus class="h-4 w-4" />
|
||||
<span>Add your own server</span>
|
||||
</button>
|
||||
</Card.Root>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
<Dialog.Footer>
|
||||
<Button variant="secondary" size="sm" onclick={() => handleOpenChange(false)}>Not now</Button>
|
||||
|
||||
<Button
|
||||
variant="default"
|
||||
size="sm"
|
||||
onclick={enableSelected}
|
||||
disabled={footerLabel === 'Continue'}>{footerLabel}</Button
|
||||
>
|
||||
</Dialog.Footer>
|
||||
</Dialog.Content>
|
||||
</Dialog.Root>
|
||||
@@ -18,15 +18,6 @@
|
||||
*/
|
||||
export { default as DialogMcpServerAddNew } from './DialogMcpServerAddNew.svelte';
|
||||
|
||||
/**
|
||||
* **DialogMcpServerRecommendations** - Suggested MCP servers opt-in dialog
|
||||
*
|
||||
* Prompts the user to enable pre-defined recommended MCP servers on first launch.
|
||||
* Shows one switch per suggested server and persists the choice as a per-chat
|
||||
* override so the selected servers become available in conversations.
|
||||
*/
|
||||
export { default as DialogMcpServerRecommendations } from './DialogMcpServerRecommendations.svelte';
|
||||
|
||||
/**
|
||||
* **DialogExportSettings** - Settings export dialog with sensitive data warning
|
||||
*
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
|
||||
let { class: className = '', onclick }: Props = $props();
|
||||
|
||||
let mcpServers = $derived(mcpStore.getServersSorted().filter((s) => s.enabled));
|
||||
let mcpServers = $derived(mcpStore.getServers().filter((s) => s.enabled));
|
||||
let enabledMcpServersForChat = $derived(
|
||||
mcpServers.filter((s) => conversationsStore.isMcpServerEnabledForChat(s.id) && s.url.trim())
|
||||
);
|
||||
|
||||
@@ -1,156 +0,0 @@
|
||||
<script lang="ts">
|
||||
import * as Card from '$lib/components/ui/card';
|
||||
import { Badge } from '$lib/components/ui/badge';
|
||||
import { Skeleton } from '$lib/components/ui/skeleton';
|
||||
import { Switch } from '$lib/components/ui/switch';
|
||||
import * as Tooltip from '$lib/components/ui/tooltip';
|
||||
import { McpServerIdentity } from '$lib/components/app/mcp';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
import { HealthCheckStatus } from '$lib/enums';
|
||||
import type { MCPServerDisplayInfo, HealthCheckState, MCPServerSettingsEntry } from '$lib/types';
|
||||
import { onMount } from 'svelte';
|
||||
import { MCP_CARD_VISIBLE_TOOL_LIMIT, NEWLINE } from '$lib/constants';
|
||||
|
||||
interface Props {
|
||||
server: MCPServerDisplayInfo & { description?: string };
|
||||
enabled?: boolean;
|
||||
onToggle?: (enabled: boolean) => void;
|
||||
}
|
||||
|
||||
let { server, enabled = false, onToggle }: Props = $props();
|
||||
|
||||
onMount(() => {
|
||||
const state = mcpStore.getHealthCheckState(server.id);
|
||||
|
||||
if (state.status === HealthCheckStatus.IDLE) {
|
||||
mcpStore.runHealthCheck(server as MCPServerSettingsEntry).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
let healthState = $derived<HealthCheckState>(mcpStore.getHealthCheckState(server.id));
|
||||
let displayName = $derived(mcpStore.getServerLabel(server));
|
||||
let faviconUrl = $derived(mcpStore.getServerFavicon(server.id));
|
||||
let isIdle = $derived(healthState.status === HealthCheckStatus.IDLE);
|
||||
let isHealthChecking = $derived(healthState.status === HealthCheckStatus.CONNECTING);
|
||||
let isError = $derived(healthState.status === HealthCheckStatus.ERROR);
|
||||
let errorMessage = $derived(
|
||||
healthState.status === HealthCheckStatus.ERROR ? healthState.message : undefined
|
||||
);
|
||||
let serverInfo = $derived(
|
||||
healthState.status === HealthCheckStatus.SUCCESS ? healthState.serverInfo : undefined
|
||||
);
|
||||
let tools = $derived(healthState.status === HealthCheckStatus.SUCCESS ? healthState.tools : []);
|
||||
let instructions = $derived(
|
||||
healthState.status === HealthCheckStatus.SUCCESS ? healthState.instructions : undefined
|
||||
);
|
||||
let showSkeleton = $derived(isIdle || isHealthChecking);
|
||||
|
||||
// Curated descriptions get two lines; instructions fallback is one line so the
|
||||
// compact card stays scannable.
|
||||
let description = $derived.by(() => {
|
||||
if (server.description) {
|
||||
return { text: server.description, lines: 2 };
|
||||
}
|
||||
if (!instructions) return null;
|
||||
const firstLine = instructions.split(NEWLINE).find((line: string) => line.trim().length > 0);
|
||||
const trimmed = firstLine?.trim();
|
||||
return trimmed ? { text: trimmed, lines: 1 } : null;
|
||||
});
|
||||
|
||||
let visibleTools = $derived(tools.slice(0, MCP_CARD_VISIBLE_TOOL_LIMIT));
|
||||
let hiddenTools = $derived(tools.slice(MCP_CARD_VISIBLE_TOOL_LIMIT));
|
||||
let hiddenToolCount = $derived(hiddenTools.length);
|
||||
|
||||
function handleToggle(checked: boolean) {
|
||||
onToggle?.(checked);
|
||||
}
|
||||
</script>
|
||||
|
||||
<Card.Root class="!gap-3 bg-muted/30 p-4">
|
||||
<div class="flex items-start justify-between gap-3">
|
||||
<div class="min-w-0 flex-1">
|
||||
{#if showSkeleton}
|
||||
<span class="flex min-w-0 items-center gap-1.5">
|
||||
<Skeleton class="h-5 w-5 rounded" />
|
||||
<Skeleton class="h-4 w-32" />
|
||||
</span>
|
||||
{:else}
|
||||
<McpServerIdentity
|
||||
{displayName}
|
||||
{faviconUrl}
|
||||
{serverInfo}
|
||||
iconClass="h-5 w-5"
|
||||
iconRounded="rounded"
|
||||
nameClass="font-medium"
|
||||
/>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
<Switch checked={enabled} disabled={isError || showSkeleton} onCheckedChange={handleToggle} />
|
||||
</div>
|
||||
|
||||
{#if isError && errorMessage}
|
||||
<p class="text-xs text-destructive">{errorMessage}</p>
|
||||
{/if}
|
||||
|
||||
{#if showSkeleton}
|
||||
<div class="space-y-1.5">
|
||||
<Skeleton class="h-3 w-full max-w-md" />
|
||||
</div>
|
||||
|
||||
<div class="flex flex-wrap items-center gap-1.5">
|
||||
<Skeleton class="h-5 w-16 rounded-full" />
|
||||
<Skeleton class="h-5 w-20 rounded-full" />
|
||||
<Skeleton class="h-5 w-24 rounded-full" />
|
||||
<Skeleton class="h-5 w-14 rounded-full" />
|
||||
</div>
|
||||
{:else}
|
||||
{#if description}
|
||||
{#if description.lines === 2}
|
||||
<p class="line-clamp-2 text-xs text-muted-foreground" title={description.text}>
|
||||
{description.text}
|
||||
</p>
|
||||
{:else}
|
||||
<p class="line-clamp-1 truncate text-xs text-muted-foreground" title={description.text}>
|
||||
{description.text}
|
||||
</p>
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
{#if tools.length > 0}
|
||||
<div class="flex flex-wrap items-center gap-1.5">
|
||||
{#each visibleTools as tool (tool.name)}
|
||||
<Tooltip.Root>
|
||||
<Tooltip.Trigger>
|
||||
<Badge variant="secondary" class="h-5 max-w-40 px-2 text-[11px]">
|
||||
<span class="block min-w-0 flex-1 truncate">{tool.name}</span>
|
||||
</Badge>
|
||||
</Tooltip.Trigger>
|
||||
|
||||
<Tooltip.Content>
|
||||
<p class="max-w-xs text-xs">
|
||||
{tool.description ?? 'No description'}
|
||||
</p>
|
||||
</Tooltip.Content>
|
||||
</Tooltip.Root>
|
||||
{/each}
|
||||
|
||||
{#if hiddenToolCount > 0}
|
||||
<Tooltip.Root>
|
||||
<Tooltip.Trigger>
|
||||
<Badge variant="secondary" class="h-5 px-2 text-[11px] text-muted-foreground">
|
||||
+ {hiddenToolCount} more tools
|
||||
</Badge>
|
||||
</Tooltip.Trigger>
|
||||
|
||||
<Tooltip.Content class="max-w-md">
|
||||
<p class="text-xs">
|
||||
{hiddenTools.map((tool) => tool.name).join(', ')}
|
||||
</p>
|
||||
</Tooltip.Content>
|
||||
</Tooltip.Root>
|
||||
{/if}
|
||||
</div>
|
||||
{/if}
|
||||
{/if}
|
||||
</Card.Root>
|
||||
@@ -180,16 +180,6 @@ export { default as McpServerCardDeleteDialog } from './McpServerCard/McpServerC
|
||||
/** Skeleton loading state for server card during health checks. */
|
||||
export { default as McpServerCardSkeleton } from './McpServerCardSkeleton.svelte';
|
||||
|
||||
/**
|
||||
* **McpServerCardCompact** - Condensed MCP server card
|
||||
*
|
||||
* Compact alternative to McpServerCard tailored for picker-style UIs.
|
||||
* Shows the server identity, status, and a flex-wrapped list of available tools.
|
||||
* Tool names are rendered as badges; hovering a badge shows its description in a tooltip.
|
||||
* Does not show connection logs or server instructions.
|
||||
*/
|
||||
export { default as McpServerCardCompact } from './McpServerCard/McpServerCardCompact.svelte';
|
||||
|
||||
/**
|
||||
* **McpServerIdentity** - Server identity display (icon, name, version)
|
||||
*
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
<script lang="ts">
|
||||
import { X, Plus } from '@lucide/svelte';
|
||||
import { Button } from '$lib/components/ui/button';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
import { conversationsStore } from '$lib/stores/conversations.svelte';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { Button } from '$lib/components/ui/button';
|
||||
import * as Empty from '$lib/components/ui/empty';
|
||||
import { ActionIcon, McpServerCard, McpServerCardSkeleton } from '$lib/components/app';
|
||||
import { DialogMcpServerAddNew } from '$lib/components/app/dialogs';
|
||||
import { HealthCheckStatus } from '$lib/enums';
|
||||
@@ -23,7 +24,6 @@
|
||||
|
||||
let servers = $derived(mcpStore.visibleMcpServers);
|
||||
|
||||
let initialLoadComplete = $state(false);
|
||||
let isAddingServer = $state(false);
|
||||
|
||||
let previousRouteId = $state<string | null>(null);
|
||||
@@ -55,26 +55,16 @@
|
||||
}
|
||||
});
|
||||
|
||||
$effect(() => {
|
||||
if (initialLoadComplete) return;
|
||||
|
||||
const allChecked =
|
||||
servers.length > 0 &&
|
||||
servers.every((server) => {
|
||||
const state = mcpStore.getHealthCheckState(server.id);
|
||||
|
||||
return (
|
||||
state.status === HealthCheckStatus.SUCCESS || state.status === HealthCheckStatus.ERROR
|
||||
);
|
||||
});
|
||||
|
||||
if (allChecked) {
|
||||
initialLoadComplete = true;
|
||||
}
|
||||
});
|
||||
// Each card decides for itself whether to render based on its own
|
||||
// health-check state, so adding a server only flashes the new card
|
||||
// (not every other already-loaded card) until its health check resolves.
|
||||
function isServerPending(serverId: string): boolean {
|
||||
const status = mcpStore.getHealthCheckState(serverId).status;
|
||||
return status === HealthCheckStatus.IDLE || status === HealthCheckStatus.CONNECTING;
|
||||
}
|
||||
</script>
|
||||
|
||||
<div in:fade={{ duration: 150 }}>
|
||||
<div in:fade={{ duration: 150 }} class="flex min-h-[calc(100dvh-4rem)] flex-col">
|
||||
<div class="fixed top-4.5 right-4 z-50 md:hidden">
|
||||
<ActionIcon icon={X} tooltip="Close" onclick={handleClose} />
|
||||
</div>
|
||||
@@ -87,53 +77,78 @@
|
||||
|
||||
<h1 class="text-lg font-semibold md:text-2xl">MCP Servers</h1>
|
||||
</div>
|
||||
|
||||
<Button
|
||||
variant="outline"
|
||||
size="lg"
|
||||
class="shrink-0 fixed md:static bottom-6 right-6"
|
||||
onclick={() => (isAddingServer = true)}
|
||||
>
|
||||
<Plus class="h-4 w-4" />
|
||||
|
||||
Add New Server
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<DialogMcpServerAddNew bind:open={isAddingServer} />
|
||||
|
||||
<div class="grid gap-5 md:space-y-4 {className}">
|
||||
{#if servers.length === 0 && !isAddingServer}
|
||||
<div class="rounded-md border border-dashed p-4 text-sm text-muted-foreground">
|
||||
No MCP Servers configured yet. Add one to enable agentic features.
|
||||
</div>
|
||||
{/if}
|
||||
{#if servers.length === 0}
|
||||
<div class="flex flex-1 items-center justify-center py-16">
|
||||
<Empty.Root class="max-w-md">
|
||||
<Empty.Header>
|
||||
<Empty.Media variant="icon">
|
||||
<Plus />
|
||||
</Empty.Media>
|
||||
|
||||
{#if servers.length > 0}
|
||||
<div
|
||||
class="grid gap-3"
|
||||
style="grid-template-columns: repeat(auto-fill, minmax(min(32rem, calc(100dvw - 2rem)), 1fr));"
|
||||
>
|
||||
{#each servers as server (server.id)}
|
||||
{#if !initialLoadComplete}
|
||||
<McpServerCardSkeleton />
|
||||
{:else}
|
||||
<McpServerCard
|
||||
{server}
|
||||
enabled={conversationsStore.isMcpServerEnabledForChat(server.id)}
|
||||
onToggle={async () => {
|
||||
const wasEnabled = conversationsStore.isMcpServerEnabledForChat(server.id);
|
||||
await conversationsStore.toggleMcpServerForChat(server.id);
|
||||
if (!wasEnabled) {
|
||||
toolsStore.enableAllToolsForServer(server.id);
|
||||
}
|
||||
}}
|
||||
onUpdate={(updates) => mcpStore.updateServer(server.id, updates)}
|
||||
onDelete={() => mcpStore.removeServer(server.id)}
|
||||
/>
|
||||
{/if}
|
||||
{/each}
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
<Empty.Title>Add your first MCP server</Empty.Title>
|
||||
|
||||
<Empty.Description>Connect a remote MCP server by URL.</Empty.Description>
|
||||
</Empty.Header>
|
||||
|
||||
<Empty.Content>
|
||||
<Button size="sm" onclick={() => (isAddingServer = true)}>
|
||||
<Plus />
|
||||
|
||||
Add New Server
|
||||
</Button>
|
||||
</Empty.Content>
|
||||
</Empty.Root>
|
||||
</div>
|
||||
{:else}
|
||||
<div
|
||||
class="grid gap-3 {className}"
|
||||
style="grid-template-columns: repeat(auto-fill, minmax(min(32rem, calc(100dvw - 2rem)), 1fr));"
|
||||
>
|
||||
{#each servers as server (server.id)}
|
||||
{#if isServerPending(server.id)}
|
||||
<McpServerCardSkeleton />
|
||||
{:else}
|
||||
<McpServerCard
|
||||
{server}
|
||||
enabled={conversationsStore.isMcpServerEnabledForChat(server.id)}
|
||||
onToggle={async () => {
|
||||
const wasEnabled = conversationsStore.isMcpServerEnabledForChat(server.id);
|
||||
await conversationsStore.toggleMcpServerForChat(server.id);
|
||||
if (!wasEnabled) {
|
||||
toolsStore.enableAllToolsForServer(server.id);
|
||||
}
|
||||
}}
|
||||
onUpdate={(updates) => mcpStore.updateServer(server.id, updates)}
|
||||
onDelete={() => mcpStore.removeServer(server.id)}
|
||||
/>
|
||||
{/if}
|
||||
{/each}
|
||||
|
||||
{#if !isAddingServer}
|
||||
<Empty.Root class="border">
|
||||
<Empty.Header>
|
||||
<Empty.Media variant="icon">
|
||||
<Plus />
|
||||
</Empty.Media>
|
||||
|
||||
<Empty.Title>Add another MCP server</Empty.Title>
|
||||
|
||||
<Empty.Description>Connect a remote MCP server by URL.</Empty.Description>
|
||||
</Empty.Header>
|
||||
|
||||
<Empty.Content>
|
||||
<Button size="sm" onclick={() => (isAddingServer = true)}>
|
||||
<Plus />
|
||||
|
||||
Add New Server
|
||||
</Button>
|
||||
</Empty.Content>
|
||||
</Empty.Root>
|
||||
{/if}
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
<script lang="ts">
|
||||
import { cn, type WithElementRef } from '$lib/components/ui/utils.js';
|
||||
import type { HTMLAttributes } from 'svelte/elements';
|
||||
|
||||
let {
|
||||
ref = $bindable(null),
|
||||
class: className,
|
||||
children,
|
||||
...restProps
|
||||
}: WithElementRef<HTMLAttributes<HTMLDivElement>> = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={ref}
|
||||
data-slot="empty-content"
|
||||
class={cn(
|
||||
'gap-2.5 text-sm flex w-full max-w-sm min-w-0 flex-col items-center text-balance',
|
||||
className
|
||||
)}
|
||||
{...restProps}
|
||||
>
|
||||
{@render children?.()}
|
||||
</div>
|
||||
@@ -0,0 +1,23 @@
|
||||
<script lang="ts">
|
||||
import { cn, type WithElementRef } from '$lib/components/ui/utils.js';
|
||||
import type { HTMLAttributes } from 'svelte/elements';
|
||||
|
||||
let {
|
||||
ref = $bindable(null),
|
||||
class: className,
|
||||
children,
|
||||
...restProps
|
||||
}: WithElementRef<HTMLAttributes<HTMLDivElement>> = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={ref}
|
||||
data-slot="empty-description"
|
||||
class={cn(
|
||||
'text-sm/relaxed text-muted-foreground [&>a:hover]:text-primary text-sm/relaxed [&>a]:underline [&>a]:underline-offset-4',
|
||||
className
|
||||
)}
|
||||
{...restProps}
|
||||
>
|
||||
{@render children?.()}
|
||||
</div>
|
||||
@@ -0,0 +1,20 @@
|
||||
<script lang="ts">
|
||||
import { cn, type WithElementRef } from '$lib/components/ui/utils.js';
|
||||
import type { HTMLAttributes } from 'svelte/elements';
|
||||
|
||||
let {
|
||||
ref = $bindable(null),
|
||||
class: className,
|
||||
children,
|
||||
...restProps
|
||||
}: WithElementRef<HTMLAttributes<HTMLDivElement>> = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={ref}
|
||||
data-slot="empty-header"
|
||||
class={cn('gap-2 flex max-w-sm flex-col items-center', className)}
|
||||
{...restProps}
|
||||
>
|
||||
{@render children?.()}
|
||||
</div>
|
||||
@@ -0,0 +1,41 @@
|
||||
<script lang="ts" module>
|
||||
import { tv, type VariantProps } from 'tailwind-variants';
|
||||
|
||||
export const emptyMediaVariants = tv({
|
||||
base: 'mb-2 flex shrink-0 items-center justify-center [&_svg]:pointer-events-none [&_svg]:shrink-0',
|
||||
variants: {
|
||||
variant: {
|
||||
default: 'bg-transparent',
|
||||
icon: "bg-muted text-foreground flex size-8 shrink-0 items-center justify-center rounded-lg [&_svg:not([class*='size-'])]:size-4"
|
||||
}
|
||||
},
|
||||
defaultVariants: {
|
||||
variant: 'default'
|
||||
}
|
||||
});
|
||||
|
||||
export type EmptyMediaVariant = VariantProps<typeof emptyMediaVariants>['variant'];
|
||||
</script>
|
||||
|
||||
<script lang="ts">
|
||||
import { cn, type WithElementRef } from '$lib/components/ui/utils.js';
|
||||
import type { HTMLAttributes } from 'svelte/elements';
|
||||
|
||||
let {
|
||||
ref = $bindable(null),
|
||||
class: className,
|
||||
children,
|
||||
variant = 'default',
|
||||
...restProps
|
||||
}: WithElementRef<HTMLAttributes<HTMLDivElement>> & { variant?: EmptyMediaVariant } = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={ref}
|
||||
data-slot="empty-icon"
|
||||
data-variant={variant}
|
||||
class={cn(emptyMediaVariants({ variant }), className)}
|
||||
{...restProps}
|
||||
>
|
||||
{@render children?.()}
|
||||
</div>
|
||||
@@ -0,0 +1,20 @@
|
||||
<script lang="ts">
|
||||
import { cn, type WithElementRef } from '$lib/components/ui/utils.js';
|
||||
import type { HTMLAttributes } from 'svelte/elements';
|
||||
|
||||
let {
|
||||
ref = $bindable(null),
|
||||
class: className,
|
||||
children,
|
||||
...restProps
|
||||
}: WithElementRef<HTMLAttributes<HTMLDivElement>> = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={ref}
|
||||
data-slot="empty-title"
|
||||
class={cn('text-sm font-medium tracking-tight', className)}
|
||||
{...restProps}
|
||||
>
|
||||
{@render children?.()}
|
||||
</div>
|
||||
@@ -0,0 +1,23 @@
|
||||
<script lang="ts">
|
||||
import { cn, type WithElementRef } from '$lib/components/ui/utils.js';
|
||||
import type { HTMLAttributes } from 'svelte/elements';
|
||||
|
||||
let {
|
||||
ref = $bindable(null),
|
||||
class: className,
|
||||
children,
|
||||
...restProps
|
||||
}: WithElementRef<HTMLAttributes<HTMLDivElement>> = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={ref}
|
||||
data-slot="empty"
|
||||
class={cn(
|
||||
'gap-4 rounded-xl border-dashed p-6 flex w-full min-w-0 flex-1 flex-col items-center justify-center text-center text-balance',
|
||||
className
|
||||
)}
|
||||
{...restProps}
|
||||
>
|
||||
{@render children?.()}
|
||||
</div>
|
||||
@@ -0,0 +1,22 @@
|
||||
import Root from './empty.svelte';
|
||||
import Header from './empty-header.svelte';
|
||||
import Media from './empty-media.svelte';
|
||||
import Title from './empty-title.svelte';
|
||||
import Description from './empty-description.svelte';
|
||||
import Content from './empty-content.svelte';
|
||||
|
||||
export {
|
||||
Root,
|
||||
Header,
|
||||
Media,
|
||||
Title,
|
||||
Description,
|
||||
Content,
|
||||
//
|
||||
Root as Empty,
|
||||
Header as EmptyHeader,
|
||||
Media as EmptyMedia,
|
||||
Title as EmptyTitle,
|
||||
Description as EmptyDescription,
|
||||
Content as EmptyContent
|
||||
};
|
||||
@@ -8,7 +8,6 @@ export * from './attachment-labels';
|
||||
export * from './database';
|
||||
export * from './reasoning-effort';
|
||||
export * from './reasoning-effort-tokens';
|
||||
export * from './recommended-mcp-servers';
|
||||
export * from './storage';
|
||||
export * from './attachment-menu';
|
||||
export * from './auto-scroll';
|
||||
|
||||
@@ -1,4 +1,2 @@
|
||||
export const MCP_SERVER_URL_PLACEHOLDER = 'https://mcp.example.com/sse';
|
||||
export const MIN_AUTOCOMPLETE_INPUT_LENGTH = 1;
|
||||
/** Number of tools shown on the compact MCP server card before collapsing to a "+ N more" badge */
|
||||
export const MCP_CARD_VISIBLE_TOOL_LIMIT = 4;
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
import { DEFAULT_MCP_CONFIG } from './mcp';
|
||||
import type { RecommendedMCPServer } from '$lib/types';
|
||||
|
||||
/**
|
||||
* Pre-defined recommended MCP servers.
|
||||
*
|
||||
* Servers are enabled by default, but they are not turned on for individual
|
||||
* conversations until the user explicitly enables them (so their tools are
|
||||
* disabled by default).
|
||||
*/
|
||||
export const RECOMMENDED_MCP_SERVERS: RecommendedMCPServer[] = [
|
||||
{
|
||||
id: 'exa-web-search',
|
||||
name: 'Exa Web Search',
|
||||
description: 'Search the web and retrieve relevant content.',
|
||||
url: 'https://mcp.exa.ai/mcp',
|
||||
enabled: true,
|
||||
requestTimeoutSeconds: DEFAULT_MCP_CONFIG.requestTimeoutSeconds
|
||||
},
|
||||
{
|
||||
id: 'huggingface-mcp',
|
||||
name: 'Hugging Face',
|
||||
description:
|
||||
'Browse models, datasets, spaces and machine learning papers from the Hugging Face hub.',
|
||||
url: 'https://huggingface.co/mcp',
|
||||
enabled: true,
|
||||
requestTimeoutSeconds: DEFAULT_MCP_CONFIG.requestTimeoutSeconds
|
||||
}
|
||||
];
|
||||
|
||||
export const RECOMMENDED_MCP_SERVER_IDS = new Set(
|
||||
RECOMMENDED_MCP_SERVERS.map((server) => server.id)
|
||||
);
|
||||
|
||||
export const RECOMMENDED_MCP_SERVERS_OPTIN_DIALOG_DELAY = 1000;
|
||||
@@ -58,7 +58,6 @@ export const SETTINGS_KEYS = {
|
||||
// MCP
|
||||
MCP_SERVERS: 'mcpServers',
|
||||
MCP_REQUEST_TIMEOUT_SECONDS: 'mcpRequestTimeoutSeconds',
|
||||
MCP_DEFAULT_SERVER_OVERRIDES: 'mcpDefaultServerOverrides',
|
||||
AGENTIC_MAX_TURNS: 'agenticMaxTurns',
|
||||
AGENTIC_MAX_TOOL_PREVIEW_LINES: 'agenticMaxToolPreviewLines',
|
||||
SHOW_TOOL_CALL_IN_PROGRESS: 'showToolCallInProgress',
|
||||
|
||||
@@ -28,7 +28,6 @@ import McpLogo from '$lib/components/app/mcp/McpLogo.svelte';
|
||||
import { SETTINGS_KEYS } from './settings-keys';
|
||||
import { ROUTES, SETTINGS_SECTION_SLUGS } from './routes';
|
||||
import { TITLE_GENERATION } from './title-generation';
|
||||
import { RECOMMENDED_MCP_SERVERS } from './recommended-mcp-servers';
|
||||
|
||||
export const SETTINGS_SECTION_TITLES = {
|
||||
GENERAL: 'General',
|
||||
@@ -775,16 +774,9 @@ const NON_UI_SETTINGS: SettingsEntry[] = [
|
||||
key: SETTINGS_KEYS.MCP_SERVERS,
|
||||
label: 'MCP servers',
|
||||
help: 'Configure MCP servers as a JSON list. Use the form in the MCP Client settings section to edit.',
|
||||
defaultValue: JSON.stringify(RECOMMENDED_MCP_SERVERS),
|
||||
defaultValue: '[]',
|
||||
type: SettingsFieldType.INPUT,
|
||||
sync: { serverKey: SETTINGS_KEYS.MCP_SERVERS, paramType: SyncableParameterType.STRING }
|
||||
},
|
||||
{
|
||||
key: SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES,
|
||||
label: 'MCP default server overrides',
|
||||
help: 'Per-server enable/disable defaults inherited by new chats. JSON-serialized list of {serverId, enabled} entries.',
|
||||
defaultValue: '[]',
|
||||
type: SettingsFieldType.INPUT
|
||||
}
|
||||
// {
|
||||
// key: SETTINGS_KEYS.PY_INTERPRETER_ENABLED,
|
||||
|
||||
@@ -22,8 +22,6 @@ export const DISABLED_TOOLS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledTool
|
||||
export const DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledToolKeys`;
|
||||
export const FAVORITE_MODELS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.favoriteModels`;
|
||||
export const REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.reasoningEffortDefault`;
|
||||
/** Set when user has interacted with the MCP server recommendations dialog (checked servers, added custom server, or dismissed) */
|
||||
export const MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.mcpServersSetupDone`;
|
||||
export const USER_OVERRIDES_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.userOverrides`;
|
||||
|
||||
/** Key prefix for per-conversation resumable stream state, conversationId is appended */
|
||||
|
||||
@@ -1,80 +0,0 @@
|
||||
import { browser } from '$app/environment';
|
||||
import {
|
||||
MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY,
|
||||
RECOMMENDED_MCP_SERVER_IDS,
|
||||
RECOMMENDED_MCP_SERVERS_OPTIN_DIALOG_DELAY
|
||||
} from '$lib/constants';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
|
||||
/**
|
||||
* First-run opt-in dialog for the recommended MCP servers.
|
||||
*
|
||||
* Owns the dismissed / open / trigger-timeout state and the effect that
|
||||
* schedules the dialog. Reads opt-in status and the configured server list
|
||||
* from `mcpStore`, so callers don't need to recompute on their side.
|
||||
*/
|
||||
export function useMcpRecommendations() {
|
||||
let dismissed = $state(
|
||||
browser && localStorage.getItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY) === 'true'
|
||||
);
|
||||
let open = $state(false);
|
||||
let checked = $state(false);
|
||||
let triggerTimeout: ReturnType<typeof setTimeout> | null = null;
|
||||
|
||||
function dismiss() {
|
||||
if (browser) {
|
||||
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
|
||||
}
|
||||
dismissed = true;
|
||||
open = false;
|
||||
if (triggerTimeout) {
|
||||
clearTimeout(triggerTimeout);
|
||||
triggerTimeout = null;
|
||||
}
|
||||
}
|
||||
|
||||
function handleOpenChange(next: boolean) {
|
||||
open = next;
|
||||
if (!next) dismiss();
|
||||
}
|
||||
|
||||
$effect(() => {
|
||||
if (!browser) return;
|
||||
|
||||
if (open || dismissed) {
|
||||
if (triggerTimeout) {
|
||||
clearTimeout(triggerTimeout);
|
||||
triggerTimeout = null;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Already evaluated once this session; leave any pending trigger alone so
|
||||
// it can still fire later. Setting `checked = true` below re-runs this
|
||||
// effect, and we must not wipe the timeout that was just scheduled.
|
||||
if (checked) return;
|
||||
|
||||
const hasRecommendations = mcpStore
|
||||
.getServers()
|
||||
.some((server) => RECOMMENDED_MCP_SERVER_IDS.has(server.id));
|
||||
|
||||
if (hasRecommendations) {
|
||||
triggerTimeout = setTimeout(() => {
|
||||
open = true;
|
||||
}, RECOMMENDED_MCP_SERVERS_OPTIN_DIALOG_DELAY);
|
||||
}
|
||||
|
||||
checked = true;
|
||||
});
|
||||
|
||||
return {
|
||||
get open() {
|
||||
return open;
|
||||
},
|
||||
get dismissed() {
|
||||
return dismissed;
|
||||
},
|
||||
dismiss,
|
||||
handleOpenChange
|
||||
};
|
||||
}
|
||||
@@ -95,7 +95,7 @@ export function useToolsPanel(): UseToolsPanelReturn {
|
||||
if (toolsStore.builtinTools.length === 0 && !toolsStore.loading) {
|
||||
toolsStore.fetchBuiltinTools();
|
||||
}
|
||||
mcpStore.runHealthChecksForServers(mcpStore.getServersSorted().filter((s) => s.enabled));
|
||||
mcpStore.runHealthChecksForServers(mcpStore.getServers().filter((s) => s.enabled));
|
||||
}
|
||||
|
||||
return {
|
||||
|
||||
@@ -522,7 +522,7 @@ const mcpDefaultEnabledMigration: Migration = {
|
||||
const config = configRaw ? JSON.parse(configRaw) : {};
|
||||
|
||||
// Don't overwrite an existing config entry — current data wins.
|
||||
if (SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES in config) {
|
||||
if (MCP_DEFAULT_OVERRIDES_LEGACY_KEY in config) {
|
||||
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG)
|
||||
console.log('[Migration] MCP default enabled: config already has overrides, skipping');
|
||||
return;
|
||||
@@ -543,7 +543,7 @@ const mcpDefaultEnabledMigration: Migration = {
|
||||
return;
|
||||
}
|
||||
|
||||
config[SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES] = raw;
|
||||
config[MCP_DEFAULT_OVERRIDES_LEGACY_KEY] = raw;
|
||||
localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config));
|
||||
|
||||
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG)
|
||||
@@ -586,6 +586,83 @@ const configTypesMigration: Migration = {
|
||||
}
|
||||
};
|
||||
|
||||
const MCP_DEFAULT_OVERRIDES_LEGACY_KEY = `${STORAGE_APP_NAME}.mcpDefaultServerOverrides`;
|
||||
const MCP_DEFAULT_OVERRIDES_MERGE_MIGRATION_ID = 'mcp-default-overrides-merge-v1';
|
||||
|
||||
/**
|
||||
* Folds `mcpDefaultServerOverrides` (the legacy "default for new chats" list,
|
||||
* JSON-encoded as `[{ serverId, enabled }, ...]`) into `mcpServers[i].enabled`.
|
||||
* The legacy override key is intentionally left in the config so a downgrade
|
||||
* keeps reading it. Runs after `mcpDefaultEnabledMigration` so any legacy
|
||||
* standalone overrides are already inside the config.
|
||||
*/
|
||||
const mcpDefaultOverridesMergeMigration: Migration = {
|
||||
id: MCP_DEFAULT_OVERRIDES_MERGE_MIGRATION_ID,
|
||||
description:
|
||||
'Merge mcpDefaultServerOverrides entries onto mcpServers[i].enabled (preserves legacy key)',
|
||||
|
||||
async run(): Promise<void> {
|
||||
const configRaw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY);
|
||||
if (configRaw === null) return;
|
||||
|
||||
const config = JSON.parse(configRaw);
|
||||
const raw = config[MCP_DEFAULT_OVERRIDES_LEGACY_KEY];
|
||||
|
||||
if (typeof raw !== 'string' || raw.length === 0) {
|
||||
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG)
|
||||
console.log('[Migration] MCP default overrides merge: nothing to merge');
|
||||
return;
|
||||
}
|
||||
|
||||
let overrides: { serverId: string; enabled: boolean }[];
|
||||
try {
|
||||
const parsed = JSON.parse(raw);
|
||||
if (!Array.isArray(parsed)) return;
|
||||
overrides = parsed.filter(
|
||||
(o) =>
|
||||
typeof o === 'object' &&
|
||||
o !== null &&
|
||||
typeof (o as Record<string, unknown>).serverId === 'string' &&
|
||||
typeof (o as Record<string, unknown>).enabled === 'boolean'
|
||||
) as { serverId: string; enabled: boolean }[];
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
|
||||
const serversRaw = config[SETTINGS_KEYS.MCP_SERVERS];
|
||||
let servers: { id: string; enabled?: boolean }[];
|
||||
try {
|
||||
servers = typeof serversRaw === 'string' ? JSON.parse(serversRaw) : [];
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!Array.isArray(servers)) servers = [];
|
||||
|
||||
let serversChanged = false;
|
||||
const knownIds = new Set(servers.map((s) => s.id));
|
||||
for (const override of overrides) {
|
||||
if (!knownIds.has(override.serverId)) continue;
|
||||
const index = servers.findIndex((s) => s.id === override.serverId);
|
||||
|
||||
if (index >= 0 && servers[index].enabled !== override.enabled) {
|
||||
servers[index] = { ...servers[index], enabled: override.enabled };
|
||||
serversChanged = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (serversChanged) {
|
||||
config[SETTINGS_KEYS.MCP_SERVERS] = JSON.stringify(servers);
|
||||
localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config));
|
||||
}
|
||||
|
||||
if (import.meta.env.DEV && import.meta.env.VITE_DEBUG)
|
||||
console.log(
|
||||
`[Migration] MCP default overrides merge: applied=${overrides.length} serversChanged=${serversChanged} (legacy key preserved)`
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
const migrations: Migration[] = [
|
||||
localStorageMigration,
|
||||
idxdbMigration,
|
||||
@@ -593,6 +670,7 @@ const migrations: Migration[] = [
|
||||
themeMigration,
|
||||
customJsonKeyMigration,
|
||||
mcpDefaultEnabledMigration,
|
||||
mcpDefaultOverridesMergeMigration,
|
||||
configTypesMigration
|
||||
];
|
||||
|
||||
|
||||
@@ -23,7 +23,8 @@ import { browser } from '$app/environment';
|
||||
import { toast } from 'svelte-sonner';
|
||||
import { DatabaseService } from '$lib/services/database.service';
|
||||
import { MigrationService } from '$lib/services/migration.service';
|
||||
import { config, settingsStore } from '$lib/stores/settings.svelte';
|
||||
import { config } from '$lib/stores/settings.svelte';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
import { filterByLeafNodeId, findLeafNode, generateConversationTitle } from '$lib/utils';
|
||||
import type { McpServerOverride } from '$lib/types/database';
|
||||
import { zipSync, unzipSync, strToU8, strFromU8 } from 'fflate';
|
||||
@@ -46,7 +47,6 @@ import {
|
||||
ISO_TIME_SEPARATOR_REPLACEMENT,
|
||||
NON_ALPHANUMERIC_REGEX,
|
||||
MULTIPLE_UNDERSCORE_REGEX,
|
||||
SETTINGS_KEYS,
|
||||
REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY
|
||||
} from '$lib/constants';
|
||||
|
||||
@@ -80,9 +80,6 @@ class ConversationsStore {
|
||||
/** Whether the store has been initialized */
|
||||
isInitialized = $state(false);
|
||||
|
||||
/** Pending MCP server overrides for new conversations (before first message) */
|
||||
pendingMcpServerOverrides = $state<McpServerOverride[]>(ConversationsStore.loadMcpDefaults());
|
||||
|
||||
/** Global (non-conversation-specific) thinking toggle default, derived from reasoning effort */
|
||||
pendingThinkingEnabled = $state(false);
|
||||
|
||||
@@ -94,28 +91,6 @@ class ConversationsStore {
|
||||
/** Last non-off reasoning effort, restored when re-enabling thinking globally */
|
||||
private lastNonOffEffort: ReasoningEffort | null = null;
|
||||
|
||||
private static loadMcpDefaults(): McpServerOverride[] {
|
||||
const raw = config()[SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES];
|
||||
if (typeof raw !== 'string' || raw.length === 0) return [];
|
||||
try {
|
||||
const parsed = JSON.parse(raw);
|
||||
if (!Array.isArray(parsed)) return [];
|
||||
return parsed.filter(
|
||||
(o: unknown) => typeof o === 'object' && o !== null && 'serverId' in o && 'enabled' in o
|
||||
) as McpServerOverride[];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
private saveMcpDefaults(): void {
|
||||
const plain = this.pendingMcpServerOverrides.map((o) => ({
|
||||
serverId: o.serverId,
|
||||
enabled: o.enabled
|
||||
}));
|
||||
settingsStore.updateConfig(SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES, JSON.stringify(plain));
|
||||
}
|
||||
|
||||
/** Load reasoning effort default from localStorage */
|
||||
private static loadReasoningEffortDefault(): ReasoningEffort | ReasoningEffort.OFF {
|
||||
if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.OFF;
|
||||
@@ -162,11 +137,6 @@ class ConversationsStore {
|
||||
|
||||
try {
|
||||
await MigrationService.runAllMigrations();
|
||||
|
||||
// Re-read defaults after migrations: a migration may have populated
|
||||
// the settings config (e.g. moved legacy MCP overrides into it).
|
||||
this.pendingMcpServerOverrides = ConversationsStore.loadMcpDefaults();
|
||||
|
||||
await this.loadConversations();
|
||||
this.isInitialized = true;
|
||||
} catch (error) {
|
||||
@@ -273,18 +243,9 @@ class ConversationsStore {
|
||||
const conversationName = name || `Chat ${new Date().toLocaleString()}`;
|
||||
const conversation = await DatabaseService.createConversation(conversationName);
|
||||
|
||||
if (this.pendingMcpServerOverrides.length > 0) {
|
||||
// Deep clone to plain objects (Svelte 5 $state uses Proxies which can't be cloned to IndexedDB)
|
||||
const plainOverrides = this.pendingMcpServerOverrides.map((o) => ({
|
||||
serverId: o.serverId,
|
||||
enabled: o.enabled
|
||||
}));
|
||||
conversation.mcpServerOverrides = plainOverrides;
|
||||
await DatabaseService.updateConversation(conversation.id, {
|
||||
mcpServerOverrides: plainOverrides
|
||||
});
|
||||
this.pendingMcpServerOverrides = [];
|
||||
}
|
||||
// New conversations inherit per-server enabled defaults directly from
|
||||
// `mcpServers[i].enabled` (see #checkServerEnabled). No per-conversation
|
||||
// override list needs to be seeded.
|
||||
|
||||
// Inherit global thinking/reasoning defaults into the new conversation
|
||||
const thinkingEnabled = this.getThinkingEnabled();
|
||||
@@ -321,7 +282,6 @@ class ConversationsStore {
|
||||
return false;
|
||||
}
|
||||
|
||||
this.pendingMcpServerOverrides = [];
|
||||
this.activeConversation = conversation;
|
||||
|
||||
if (conversation.currNode) {
|
||||
@@ -351,7 +311,6 @@ class ConversationsStore {
|
||||
this.activeConversation = null;
|
||||
this.activeMessages = [];
|
||||
// reload defaults so new chats inherit persisted state
|
||||
this.pendingMcpServerOverrides = ConversationsStore.loadMcpDefaults();
|
||||
this.pendingReasoningEffort = ConversationsStore.loadReasoningEffortDefault();
|
||||
}
|
||||
|
||||
@@ -641,11 +600,30 @@ class ConversationsStore {
|
||||
*
|
||||
*/
|
||||
|
||||
/**
|
||||
/**
|
||||
* Resolve the per-server enabled value when no active conversation exists.
|
||||
* The default for new chats is the server's own `enabled` flag in `mcpServers`.
|
||||
*/
|
||||
#getDefaultOverrideForNoConversation(serverId: string): McpServerOverride | undefined {
|
||||
const server = mcpStore.getServers().find((s) => s.id === serverId);
|
||||
if (!server) return undefined;
|
||||
return { serverId, enabled: server.enabled };
|
||||
}
|
||||
|
||||
/**
|
||||
* Default overrides for new chats are derived from `mcpServers[i].enabled`,
|
||||
* so the global on/off state lives in one place.
|
||||
*/
|
||||
#getAllDefaultOverridesForNoConversation(): McpServerOverride[] {
|
||||
return mcpStore.getServers().map((s) => ({ serverId: s.id, enabled: s.enabled }));
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets MCP server override for a specific server in the active conversation.
|
||||
* Falls back to pending overrides if no active conversation exists.
|
||||
* Falls back to `mcpServers[i].enabled` if no active conversation exists.
|
||||
* @param serverId - The server ID to check
|
||||
* @returns The override if set, undefined if using global setting
|
||||
* @returns The override if set, undefined if no matching server
|
||||
*/
|
||||
getMcpServerOverride(serverId: string): McpServerOverride | undefined {
|
||||
if (this.activeConversation) {
|
||||
@@ -653,18 +631,18 @@ class ConversationsStore {
|
||||
(o: McpServerOverride) => o.serverId === serverId
|
||||
);
|
||||
}
|
||||
return this.pendingMcpServerOverrides.find((o) => o.serverId === serverId);
|
||||
return this.#getDefaultOverrideForNoConversation(serverId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all MCP server overrides for the current conversation.
|
||||
* Returns pending overrides if no active conversation.
|
||||
* When no active conversation, derives from `mcpServers[i].enabled`.
|
||||
*/
|
||||
getAllMcpServerOverrides(): McpServerOverride[] {
|
||||
if (this.activeConversation?.mcpServerOverrides) {
|
||||
return this.activeConversation.mcpServerOverrides;
|
||||
}
|
||||
return this.pendingMcpServerOverrides;
|
||||
return this.#getAllDefaultOverridesForNoConversation();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -679,13 +657,16 @@ class ConversationsStore {
|
||||
|
||||
/**
|
||||
* Sets or removes MCP server override for the active conversation.
|
||||
* If no conversation exists, stores as pending override.
|
||||
* If no conversation exists, persists `enabled` onto `mcpServers[i].enabled`
|
||||
* (the single source of truth for new-chat defaults).
|
||||
* @param serverId - The server ID to override
|
||||
* @param enabled - The enabled state, or undefined to remove override
|
||||
* @param enabled - The enabled state, or undefined to remove per-conversation override
|
||||
*/
|
||||
async setMcpServerOverride(serverId: string, enabled: boolean | undefined): Promise<void> {
|
||||
if (!this.activeConversation) {
|
||||
this.setPendingMcpServerOverride(serverId, enabled);
|
||||
if (enabled !== undefined) {
|
||||
mcpStore.updateServer(serverId, { enabled });
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -729,29 +710,6 @@ class ConversationsStore {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets or removes a pending MCP server override (for new conversations).
|
||||
*/
|
||||
private setPendingMcpServerOverride(serverId: string, enabled: boolean | undefined): void {
|
||||
if (enabled === undefined) {
|
||||
this.pendingMcpServerOverrides = this.pendingMcpServerOverrides.filter(
|
||||
(o) => o.serverId !== serverId
|
||||
);
|
||||
} else {
|
||||
const existingIndex = this.pendingMcpServerOverrides.findIndex(
|
||||
(o) => o.serverId === serverId
|
||||
);
|
||||
if (existingIndex >= 0) {
|
||||
const newOverrides = [...this.pendingMcpServerOverrides];
|
||||
newOverrides[existingIndex] = { serverId, enabled };
|
||||
this.pendingMcpServerOverrides = newOverrides;
|
||||
} else {
|
||||
this.pendingMcpServerOverrides = [...this.pendingMcpServerOverrides, { serverId, enabled }];
|
||||
}
|
||||
}
|
||||
this.saveMcpDefaults();
|
||||
}
|
||||
|
||||
/**
|
||||
* Toggles MCP server enabled state for the active conversation.
|
||||
* @param serverId - The server ID to toggle
|
||||
@@ -769,14 +727,6 @@ class ConversationsStore {
|
||||
await this.setMcpServerOverride(serverId, undefined);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clears all pending MCP server overrides.
|
||||
*/
|
||||
clearPendingMcpServerOverrides(): void {
|
||||
this.pendingMcpServerOverrides = [];
|
||||
this.saveMcpDefaults();
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the effective thinking-enabled state for the active conversation.
|
||||
* Returns the conversation override if set, otherwise the global default.
|
||||
|
||||
@@ -470,18 +470,12 @@ class MCPStore {
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: try favicon from root domain
|
||||
const fallbackUrl = this.#getServerFaviconFallback(server.url);
|
||||
if (fallbackUrl) {
|
||||
return fallbackUrl;
|
||||
}
|
||||
|
||||
return null;
|
||||
return this.#getServerFaviconFallback(server.url);
|
||||
}
|
||||
|
||||
/**
|
||||
* Construct a fallback favicon URL from the MCP server URL.
|
||||
* e.g. https://mcp.exa.ai/mcp -> https://exa.ai/favicon.ico
|
||||
* e.g. https://mcp.example.com/sse -> https://example.com/favicon.ico
|
||||
*/
|
||||
#getServerFaviconFallback(serverUrl: string): string | null {
|
||||
try {
|
||||
@@ -505,27 +499,6 @@ class MCPStore {
|
||||
return null;
|
||||
}
|
||||
|
||||
isAnyServerLoading(): boolean {
|
||||
return this.getServers().some((s) => {
|
||||
const state = this.getHealthCheckState(s.id);
|
||||
|
||||
return (
|
||||
state.status === HealthCheckStatus.IDLE || state.status === HealthCheckStatus.CONNECTING
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
getServersSorted(): MCPServerSettingsEntry[] {
|
||||
const servers = this.getServers();
|
||||
if (this.isAnyServerLoading()) {
|
||||
return servers;
|
||||
}
|
||||
|
||||
return [...servers].sort((a, b) =>
|
||||
this.getServerLabel(a).localeCompare(this.getServerLabel(b))
|
||||
);
|
||||
}
|
||||
|
||||
addServer(
|
||||
serverData: Omit<MCPServerSettingsEntry, 'id' | 'requestTimeoutSeconds'> & { id?: string }
|
||||
): MCPServerSettingsEntry {
|
||||
@@ -579,10 +552,11 @@ class MCPStore {
|
||||
}
|
||||
|
||||
/**
|
||||
* MCP servers selectable in chat-add UIs and the settings page.
|
||||
* MCP servers selectable in chat-add UIs and the settings page,
|
||||
* in the order they were added to the config.
|
||||
*/
|
||||
get visibleMcpServers(): MCPServerSettingsEntry[] {
|
||||
return this.getServersSorted().filter((server) => server.enabled);
|
||||
return this.getServers().filter((server) => server.enabled);
|
||||
}
|
||||
|
||||
async ensureInitialized(perChatOverrides?: McpServerOverride[]): Promise<boolean> {
|
||||
|
||||
@@ -128,7 +128,6 @@ export type {
|
||||
MCPClientConfig,
|
||||
MCPServerSettingsEntry,
|
||||
MCPServerDisplayInfo,
|
||||
RecommendedMCPServer,
|
||||
MCPToolCall,
|
||||
OpenAIToolDefinition,
|
||||
ServerStatus,
|
||||
|
||||
Vendored
-9
@@ -226,15 +226,6 @@ export type MCPServerSettingsEntry = MCPServerDisplayInfo & {
|
||||
useProxy?: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Pre-defined recommended MCP server shown to the user in onboarding/picker UIs.
|
||||
*/
|
||||
export interface RecommendedMCPServer extends MCPServerDisplayInfo {
|
||||
description: string;
|
||||
enabled: boolean;
|
||||
requestTimeoutSeconds: number;
|
||||
}
|
||||
|
||||
export interface MCPHostManagerConfig {
|
||||
servers: MCPClientConfig['servers'];
|
||||
clientInfo?: Implementation;
|
||||
|
||||
@@ -8,7 +8,6 @@
|
||||
import { onMount } from 'svelte';
|
||||
|
||||
import { SidebarNavigation, DialogConversationTitleUpdate } from '$lib/components/app';
|
||||
import { DialogMcpServerRecommendations } from '$lib/components/app/dialogs';
|
||||
import { PwaMetaTags, PwaRefreshAlert } from '$lib/components/pwa';
|
||||
import { pwaAssetsHead } from 'virtual:pwa-assets/head';
|
||||
|
||||
@@ -27,7 +26,6 @@
|
||||
import { FAVICON_PATHS, FAVICON_SELECTORS } from '$lib/constants/pwa';
|
||||
import { useKeyboardShortcuts } from '$lib/hooks/use-keyboard-shortcuts.svelte';
|
||||
import { usePwa } from '$lib/hooks/use-pwa.svelte';
|
||||
import { useMcpRecommendations } from '$lib/hooks/use-mcp-recommendations.svelte';
|
||||
import { conversations } from '$lib/stores/conversations.svelte';
|
||||
import { isMobile } from '$lib/stores/viewport.svelte';
|
||||
import { theme } from '$lib/stores/theme.svelte';
|
||||
@@ -39,8 +37,6 @@
|
||||
let innerHeight = $state<number | undefined>();
|
||||
let innerWidth = $state(browser ? window.innerWidth : 0);
|
||||
|
||||
const mcpRecommendations = useMcpRecommendations();
|
||||
|
||||
let chatSidebar:
|
||||
| {
|
||||
activateSearchMode?: () => void;
|
||||
@@ -239,7 +235,10 @@
|
||||
});
|
||||
|
||||
// Background MCP server health checks on app load
|
||||
// Fetch enabled servers from settings and run health checks in background
|
||||
// Fetch enabled servers from settings and run health checks in background.
|
||||
// Only IDLE servers are checked; already-resolved (SUCCESS / ERROR) servers
|
||||
// keep their existing state, so adding or removing a server does not flash
|
||||
// every other card back through skeleton state.
|
||||
$effect(() => {
|
||||
if (!browser) return;
|
||||
|
||||
@@ -251,7 +250,7 @@
|
||||
if (enabledServers.length > 0) {
|
||||
untrack(() => {
|
||||
// Run health checks in background (don't await)
|
||||
mcpStore.runHealthChecksForServers(enabledServers, false).catch((error) => {
|
||||
mcpStore.runHealthChecksForServers(enabledServers, true).catch((error) => {
|
||||
console.warn('[layout] MCP health checks failed:', error);
|
||||
});
|
||||
});
|
||||
@@ -325,11 +324,6 @@
|
||||
onConfirm={handleTitleUpdateConfirm}
|
||||
onCancel={handleTitleUpdateCancel}
|
||||
/>
|
||||
|
||||
<DialogMcpServerRecommendations
|
||||
open={mcpRecommendations.open}
|
||||
onOpenChange={mcpRecommendations.handleOpenChange}
|
||||
/>
|
||||
</Tooltip.Provider>
|
||||
|
||||
<!-- PWA update prompt + version -->
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
import { afterEach, beforeAll, beforeEach, describe, expect, it } from 'vitest';
|
||||
import { STORAGE_APP_NAME, CONFIG_LOCALSTORAGE_KEY } from '$lib/constants';
|
||||
|
||||
// node env unit project has no DOM, install a minimal localStorage backed by a Map
|
||||
beforeAll(() => {
|
||||
const store = new Map<string, string>();
|
||||
const polyfill: Storage = {
|
||||
get length() {
|
||||
return store.size;
|
||||
},
|
||||
clear: () => store.clear(),
|
||||
getItem: (k) => (store.has(k) ? store.get(k)! : null),
|
||||
key: (i) => Array.from(store.keys())[i] ?? null,
|
||||
removeItem: (k) => {
|
||||
store.delete(k);
|
||||
},
|
||||
setItem: (k, v) => {
|
||||
store.set(k, String(v));
|
||||
}
|
||||
};
|
||||
(globalThis as unknown as { localStorage: Storage }).localStorage = polyfill;
|
||||
});
|
||||
|
||||
/**
|
||||
* Migration `mcp-default-overrides-merge-v1` folds the values of the parallel
|
||||
* `mcpDefaultServerOverrides` config entry onto `mcpServers[i].enabled` (the
|
||||
* single source of truth for new-chat defaults). The legacy key is kept on
|
||||
* disk for downgrade compatibility.
|
||||
*/
|
||||
describe('mcp-default-overrides-merge-v1 migration', () => {
|
||||
const MIGRATION_STATE_KEY = `${STORAGE_APP_NAME}.migration-state`;
|
||||
const MCP_DEFAULT_OVERRIDES_KEY = `${STORAGE_APP_NAME}.mcpDefaultServerOverrides`;
|
||||
|
||||
beforeEach(async () => {
|
||||
localStorage.clear();
|
||||
// Reset the migration run counter so `runAllMigrations` is guaranteed to execute.
|
||||
await import('$lib/services/migration.service').then((mod) =>
|
||||
mod.MigrationService.resetState()
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
localStorage.clear();
|
||||
});
|
||||
|
||||
async function runMigrations() {
|
||||
const { MigrationService } = await import('$lib/services/migration.service');
|
||||
await MigrationService.runAllMigrations();
|
||||
}
|
||||
|
||||
function readConfig(): Record<string, unknown> {
|
||||
const raw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY);
|
||||
|
||||
return raw ? (JSON.parse(raw) as Record<string, unknown>) : {};
|
||||
}
|
||||
|
||||
function writeConfig(config: Record<string, unknown>) {
|
||||
localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config));
|
||||
}
|
||||
|
||||
it('applies matching overrides onto mcpServers[i].enabled and preserves the legacy key', async () => {
|
||||
writeConfig({
|
||||
mcpServers: JSON.stringify([
|
||||
{ id: 'exa', enabled: false, url: 'https://mcp.exa.ai/mcp' },
|
||||
{ id: 'hf', enabled: false, url: 'https://huggingface.co/mcp' }
|
||||
]),
|
||||
[MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([
|
||||
{ serverId: 'exa', enabled: true },
|
||||
{ serverId: 'hf', enabled: false }
|
||||
])
|
||||
});
|
||||
|
||||
await runMigrations();
|
||||
|
||||
const after = readConfig();
|
||||
const servers = JSON.parse(after.mcpServers as string) as Array<{
|
||||
id: string;
|
||||
enabled: boolean;
|
||||
}>;
|
||||
|
||||
expect(servers.find((s) => s.id === 'exa')?.enabled).toBe(true);
|
||||
expect(servers.find((s) => s.id === 'hf')?.enabled).toBe(false);
|
||||
expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(true);
|
||||
});
|
||||
|
||||
it('skips override ids that do not match any configured server', async () => {
|
||||
writeConfig({
|
||||
mcpServers: JSON.stringify([{ id: 'exa', enabled: false, url: 'https://mcp.exa.ai/mcp' }]),
|
||||
[MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([
|
||||
{ serverId: 'orphan', enabled: true },
|
||||
{ serverId: 'exa', enabled: true }
|
||||
])
|
||||
});
|
||||
|
||||
await runMigrations();
|
||||
|
||||
const after = readConfig();
|
||||
const servers = JSON.parse(after.mcpServers as string) as Array<{
|
||||
id: string;
|
||||
enabled: boolean;
|
||||
}>;
|
||||
|
||||
expect(servers).toHaveLength(1);
|
||||
expect(servers[0].enabled).toBe(true);
|
||||
expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(true);
|
||||
});
|
||||
|
||||
it('is a no-op when there are no legacy overrides', async () => {
|
||||
writeConfig({
|
||||
mcpServers: JSON.stringify([{ id: 'exa', enabled: true, url: 'https://mcp.exa.ai/mcp' }])
|
||||
});
|
||||
|
||||
await runMigrations();
|
||||
|
||||
const after = readConfig();
|
||||
const servers = JSON.parse(after.mcpServers as string) as Array<{
|
||||
id: string;
|
||||
enabled: boolean;
|
||||
}>;
|
||||
|
||||
expect(servers[0].enabled).toBe(true);
|
||||
expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(false);
|
||||
});
|
||||
|
||||
it('does not rewrite mcpServers when override.enabled already matches', async () => {
|
||||
const originalServers = JSON.stringify([
|
||||
{ id: 'exa', enabled: true, url: 'https://mcp.exa.ai/mcp' }
|
||||
]);
|
||||
|
||||
writeConfig({
|
||||
mcpServers: originalServers,
|
||||
[MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([{ serverId: 'exa', enabled: true }])
|
||||
});
|
||||
|
||||
await runMigrations();
|
||||
|
||||
const after = readConfig();
|
||||
expect(after.mcpServers).toBe(originalServers);
|
||||
expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(true);
|
||||
});
|
||||
|
||||
it('records itself as completed so subsequent loads do not re-run', async () => {
|
||||
writeConfig({
|
||||
mcpServers: JSON.stringify([{ id: 'exa', enabled: false, url: 'https://mcp.exa.ai/mcp' }]),
|
||||
[MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([{ serverId: 'exa', enabled: true }])
|
||||
});
|
||||
|
||||
const { MigrationService } = await import('$lib/services/migration.service');
|
||||
|
||||
await MigrationService.runAllMigrations();
|
||||
|
||||
const stateRaw = localStorage.getItem(MIGRATION_STATE_KEY);
|
||||
expect(stateRaw).not.toBeNull();
|
||||
const state = JSON.parse(stateRaw!) as { completed: string[]; failed: string[] };
|
||||
expect(state.completed).toContain('mcp-default-overrides-merge-v1');
|
||||
expect(state.failed).not.toContain('mcp-default-overrides-merge-v1');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,19 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { SETTINGS_KEYS } from '$lib/constants/settings-keys';
|
||||
|
||||
/**
|
||||
* Default-value policy for the `MCP_SERVERS` setting.
|
||||
*
|
||||
* Earlier versions of the UI preloaded a hard-coded list of suggested
|
||||
* MCP servers into this setting on first install. That caused silent
|
||||
* third-party HTTP requests at app load (see issue #25509) and a popup
|
||||
* "recommendation" dialog (see issue #25274). New users must now opt
|
||||
* in explicitly when adding a server, so the default is an empty list.
|
||||
*/
|
||||
describe('MCP_SERVERS default value', () => {
|
||||
it('does not preload any servers in the MCP_SERVERS setting default', async () => {
|
||||
const { SETTING_CONFIG_DEFAULT } = await import('$lib/constants/settings-registry');
|
||||
|
||||
expect(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.MCP_SERVERS]).toBe('[]');
|
||||
}, 15000);
|
||||
});
|
||||
@@ -5,11 +5,10 @@ import { DEFAULT_MCP_CONFIG, MCP_SERVER_ID_PREFIX } from '$lib/constants/mcp';
|
||||
/**
|
||||
* Tests for the mcpServers settings parser.
|
||||
*
|
||||
* The branch seeds the MCP servers setting with a default value of
|
||||
* `JSON.stringify(RECOMMENDED_MCP_SERVERS)`, so the parser has to be
|
||||
* resilient to anything that may live in the user's localStorage: malformed
|
||||
* JSON, wrong shapes, missing fields, falsy-but-not-zero numbers, and entry
|
||||
* arrays that have been mutated by the user via the settings form.
|
||||
* The parser has to be resilient to anything that may live in the
|
||||
* user's localStorage: malformed JSON, wrong shapes, missing fields,
|
||||
* falsy-but-not-zero numbers, and entry arrays that have been mutated
|
||||
* by the user via the settings form.
|
||||
*/
|
||||
describe('parseMcpServerSettings', () => {
|
||||
it('returns an empty array for falsy or whitespace-only input', () => {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user