Compare commits

..

13 Commits

Author SHA1 Message Date
kdkd e3546c7948 Fix conditional to display 'LLAMA_SPLIT_MODE_TENSOR not implemented for architecture' message (#24926) 2026-07-11 20:03:24 +02:00
Rohit Mahesh d72bfa38f7 gguf : reject empty metadata keys (#24917) 2026-07-11 20:02:44 +02:00
cphlipot 3cec3bcd16 cuda: Don't crash when querying memory on device with no free memory. (#25157)
If a Cuda device has no or limited available memory, the actual call
to cudaMemGetInfo() itself can cause a fatal crash due to a cuda out
of memory error (there is not enough memory to actually query memory)

This causes an issue because we query memory for all devices at
startup even if the user isn't trying to use the device for inference.

Fix this by making the error non-fatal and assigning zero total/free
memory to the device. This will have the downstream effect of the fit
algorithm not trying to put any layers on it, which is desired outcome
vs hard crashing.

this also prevents crashes in cuda enabled builds when user explicitly
passes '-dev none'
2026-07-11 19:13:43 +02:00
Aman Gupta 13f2b28b09 DeepseekV4: clear cache only for seq rather than full (#25521) 2026-07-11 23:35:45 +08:00
Xuan-Son Nguyen c92e806d1c server: allow stream for exec_shell_command (#25526)
* init stream

* add stream for shell tool

* add test

* nits

* update docs
2026-07-11 12:42:55 +02:00
Xuan-Son Nguyen ea1f7bbb5d server: refactor server_stream (#25541)
* server: refactoring, remove spipe from server_http_res

* wip

* remove non-thread-safe rd.stop() call

* move server_res_spipe

* nits

* improve server_stream_create_spipe

* server-stream: update dev docs for the improved API

---------

Co-authored-by: Pascal <admin@serveurperso.com>
2026-07-11 12:41:47 +02:00
fairydreaming 00f5442cc4 ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer (#24231)
* ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer

* ggml : remove scale parameters from lightning indexer OP, add f16 mask parameter

* tests : add GGML_OP_LIGHTNING_INDEXER tests

* ggml : bump RPC version

* chore : check if lightning indexer input tensors are not transposed

* tests : count flops instead of bandwidth in lightning indexer test

* chore : add missing const

* chore : whitespace

* ggml : renamed variables in CPU lightning indexer implementation

* ggml : fix lightning indexer mask broadcasting

* tests : tests for lightning indexer mask broadcasting

* chore : whitespace

* llama : use GGML_OP_LIGHTNING_INDEXER in DeepSeek V3.2 and DeepSeek V4 models

---------

Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
2026-07-11 11:39:07 +02:00
Raman Shinde 76f2798059 Vulkan: route large matmuls to medium tile on Adreno (#24877)
* [Vulkan] Fixes llama-cli breaking over longer promts sizes

The llama-cli was breaking for longer promts sizes for q4_0 quantized networks. Causing due to insufficient shared memory.

* Removed the un-used Adreno device

* Updated matmul for small pipeline.
2026-07-11 10:28:29 +02:00
Hongqiang Wang 1d1d9a9ed7 opencl: add int8 dp4 dense and MoE prefill optimization for Adreno GPUs (#25537)
* opencl: add int8 dp4 dense and moe GEMM

* opencl: refactor

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>
2026-07-10 23:05:58 -07:00
Pascal 4f37f51972 server: accept null sampling params (#25538)
* server: accept null sampling params

Extend the schema validation to treat a null value as absent, so
clients can send null on nullable params (temperature, top_p, ...)
to request the server default. This matches the OpenAI spec and the
json_value convention used elsewhere.

Add has_field() to skip null in the field eval guards.

* has_field -> has_value​
2026-07-10 22:07:29 +02:00
eduardopessin c749cb0417 llama : make tensor-split regex patterns static (#24710)
llama_meta_device_get_split_state() recompiled 29 std::regex on every call.
In -sm tensor mode the callback runs once per tensor per token, so this
dominated the decode thread in profiling. Mark them static const so they are
compiled once. Kept inside the function (local statics are thread-safe since
C++11). Patterns are literal and stateless, so behavior is unchanged.
2026-07-10 19:04:12 +02:00
Max Krasnyansky 67776eaee5 hexagon: improve ARGSORT performance for small tensors (#25512)
* hex-sort: add efficient bitomic sort in hvx regs up to 1024 elements

* hex-sort: fix inverted vrors

* hex-sort: specialize sort functions for the common cases

* hex-sort: add tracing and local context
2026-07-10 09:06:06 -07:00
Xuan-Son Nguyen 22b69b6e92 arg: prevent duplicate spec model downloads (#25527) 2026-07-10 16:53:26 +02:00
55 changed files with 5899 additions and 468 deletions
+17 -11
View File
@@ -488,12 +488,15 @@ void common_models_handler_apply(common_models_handler & handler, common_params
task.opts = opts;
tasks.push_back(task);
}
bool had_spec_url = false;
if (!params.speculative.draft.mparams.url.empty()) {
common_download_task task;
task.url = params.speculative.draft.mparams.url;
task.local_path = params.speculative.draft.mparams.path;
task.opts = opts;
tasks.push_back(task);
had_spec_url = true;
}
// handle hf_plan tasks
@@ -513,6 +516,18 @@ void common_models_handler_apply(common_models_handler & handler, common_params
});
}
};
// handle plan_spec (e.g. --spec-draft-hf)
if (!plan_spec.model_files.empty() && !had_spec_url) {
add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
had_spec_url = true;
}
// handle vocoder plan (e.g. --hf-repo-v)
if (!plan_voc.model_files.empty()) {
add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model);
}
if (!plan.model_files.empty()) {
add_tasks(plan.model_files, plan.primary, params.model);
}
@@ -521,7 +536,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
params.mmproj.path = hf_cache::finalize_file(plan.mmproj);
});
}
if (!plan.mtp.local_path.empty()) {
if (!plan.mtp.local_path.empty() && !had_spec_url) {
tasks.emplace_back(plan.mtp, opts, [&]() {
// only fall back to the discovered MTP head when no draft was explicitly provided
if (params.speculative.draft.mparams.empty()) {
@@ -540,16 +555,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
});
}
// handle plan_spec (e.g. --spec-draft-hf)
if (!plan_spec.model_files.empty()) {
add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
}
// handle vocoder plan (e.g. --hf-repo-v)
if (!plan_voc.model_files.empty()) {
add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model);
}
// run all tasks in parallel
if (!params.offline) {
// if duplicated files are found, only download once (but still call on_done for each task)
@@ -562,6 +567,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
}
std::vector<common_download_task> unique_tasks_vec;
for (auto & pair : unique_tasks) {
LOG_DBG("download task: %s -> %s\n", pair.second->url.c_str(), pair.second->local_path.c_str());
unique_tasks_vec.push_back(*pair.second);
}
common_download_run_tasks(unique_tasks_vec);
+2 -2
View File
@@ -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
+19
View File
@@ -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);
+11
View File
@@ -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;
}
+84
View File
@@ -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]);
}
}
}
}
+1
View File
@@ -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);
+8 -1
View File
@@ -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__)
+2 -2
View File
@@ -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
+239 -19
View File
@@ -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;
}
+16
View File
@@ -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
File diff suppressed because it is too large Load Diff
+118
View File
@@ -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);
}
+8
View File
@@ -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
View File
@@ -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) {
+4
View File
@@ -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);
+15
View File
@@ -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() {
+2
View File
@@ -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
View File
@@ -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;
+1
View File
@@ -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
View File
@@ -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);
}
//
+4 -2
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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));
+87
View File
@@ -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;
}
+6 -1
View File
@@ -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,
+91 -9
View File
@@ -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;
}
+27 -4
View File
@@ -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)
+8 -19
View File
@@ -3979,11 +3979,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 +3991,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 +4028,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 +4172,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 +4184,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 +4277,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;
}
+3 -16
View File
@@ -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();
}
}
+2 -9
View File
@@ -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;
};
+10 -4
View File
@@ -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);
});
+63 -83
View File
@@ -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;
};
}
+19 -28
View File
@@ -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
View File
@@ -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}};
}
+17 -2
View File
@@ -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;
@@ -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()