### GGML Version
set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 16)
+set(GGML_VERSION_MINOR 18)
set(GGML_VERSION_PATCH 0)
set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
option(GGML_HIP_GRAPHS "ggml: use HIP graph" ON)
option(GGML_HIP_RCCL "ggml: use ROCm Collective Comm. Library" OFF)
option(GGML_HIP_NO_VMM "ggml: do not try to use HIP VMM" ON)
-option(GGML_HIP_ROCWMMA_FATTN "ggml: enable rocWMMA for FlashAttention" OFF)
option(GGML_HIP_MMQ_MFMA "ggml: enable MFMA MMA for CDNA in MMQ" ON)
option(GGML_HIP_EXPORT_METRICS "ggml: enable kernel perf metrics output" OFF)
option(GGML_MUSA_GRAPHS "ggml: use MUSA graph, experimental, unstable" OFF)
"ggml: sycl device architecture")
option(GGML_OPENVINO "ggml: use OPENVINO" OFF)
+option(GGML_ET "ggml: use ET backend" OFF)
+option(GGML_ET_SYSEMU "ggml: use ET backend via sysemu" OFF)
option(GGML_OPENCL "ggml: use OpenCL" OFF)
option(GGML_OPENCL_PROFILING "ggml: use OpenCL profiling (increases overhead)" OFF)
GGML_BACKEND_API int ggml_cpu_has_sve (void);
GGML_BACKEND_API int ggml_cpu_get_sve_cnt (void); // sve vector length in bytes
GGML_BACKEND_API int ggml_cpu_has_sme (void);
+ GGML_BACKEND_API int ggml_cpu_has_sme2 (void);
// other
GGML_BACKEND_API int ggml_cpu_has_riscv_v (void);
GGML_BACKEND_API int ggml_cpu_get_rvv_vlen (void); // risc-v vector length in bytes
--- /dev/null
+#pragma once
+
+#include "ggml.h"
+#include "ggml-backend.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define GGML_ET_NAME "ET"
+
+// backend API
+GGML_BACKEND_API ggml_guid_t ggml_backend_et_guid(void);
+GGML_BACKEND_API ggml_backend_t ggml_backend_et_init(size_t devidx);
+
+GGML_BACKEND_API bool ggml_backend_is_et(ggml_backend_t backend);
+GGML_BACKEND_API int ggml_backend_et_get_device_count(void);
+GGML_BACKEND_API void ggml_backend_et_get_device_description(int devidx, char * description, size_t description_size);
+GGML_BACKEND_API void ggml_backend_et_get_device_memory(int devidx, size_t * free, size_t * total);
+
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_et_buffer_type(size_t dev_num);
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_et_host_buffer_type(void);
+
+GGML_BACKEND_API ggml_backend_reg_t ggml_backend_et_reg(void);
+
+#ifdef __cplusplus
+}
+#endif
extern "C" {
#endif
-#define RPC_PROTO_MAJOR_VERSION 4
+#define RPC_PROTO_MAJOR_VERSION 5
#define RPC_PROTO_MINOR_VERSION 0
-#define RPC_PROTO_PATCH_VERSION 1
+#define RPC_PROTO_PATCH_VERSION 0
#ifdef __cplusplus
-static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
#endif
#define GGML_RPC_MAX_SERVERS 16
GGML_OP_RWKV_WKV7,
GGML_OP_SOLVE_TRI,
GGML_OP_GATED_DELTA_NET,
+ GGML_OP_LIGHTNING_INDEXER,
+ GGML_OP_DSV4_HC_COMB,
+ GGML_OP_DSV4_HC_PRE,
+ GGML_OP_DSV4_HC_POST,
GGML_OP_UNARY,
GGML_API bool ggml_is_contiguous_1(const struct ggml_tensor * tensor); // contiguous for dims >= 1
GGML_API bool ggml_is_contiguous_2(const struct ggml_tensor * tensor); // contiguous for dims >= 2
+ GGML_API bool ggml_is_contiguous_to_1(const struct ggml_tensor * tensor); // contiguous for dims < 1
+ GGML_API bool ggml_is_contiguous_to_2(const struct ggml_tensor * tensor); // contiguous for dims < 2
+ GGML_API bool ggml_is_contiguous_to_3(const struct ggml_tensor * tensor); // contiguous for dims < 3
+
// returns whether the tensor elements are allocated as one contiguous block of memory (no gaps, but permutation ok)
GGML_API bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor);
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);
+
+ // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
+ // In short these operations are replacements for the original residual connection (x = transformer(x) + x)
+ // using a richer representation through streams.
+ //
+ // hc_comb: mixes [(2 + hc)*hc, n_tokens], scale [3], base [(2 + hc)*hc]
+ // -> [dst_hc, src_hc, n_tokens]
+ // logits[dst, src, t] = mixes[2*hc + dst + hc*src, t]*scale[2]
+ // + base[2*hc + dst + hc*src]
+ // Softmax over dst, add eps, normalize over src, then repeat normalization
+ // over dst followed by src for iterations 1 through n_iter - 1.
+ GGML_API struct ggml_tensor * ggml_dsv4_hc_comb(
+ struct ggml_context * ctx,
+ struct ggml_tensor * mixes,
+ struct ggml_tensor * scale,
+ struct ggml_tensor * base,
+ float eps,
+ int32_t n_iter);
+
+ // hc_pre: x [n_embd, hc, n_tokens], weights [hc, n_tokens] -> [n_embd, n_tokens]
+ // result[i, t] = sum_h x[i, h, t]*weights[h, t]
+ //
+ GGML_API struct ggml_tensor * ggml_dsv4_hc_pre(
+ struct ggml_context * ctx,
+ struct ggml_tensor * x,
+ struct ggml_tensor * weights);
+
+ // hc_post: x [n_embd, n_tokens], residual [n_embd, hc, n_tokens],
+ // post [hc, n_tokens], comb [dst_hc, src_hc, n_tokens]
+ // -> [n_embd, hc, n_tokens]
+ // result[i, dst, t] = x[i, t]*post[dst, t]
+ // + sum_src residual[i, src, t]*comb[dst, src, t]
+ //
+ GGML_API struct ggml_tensor * ggml_dsv4_hc_post(
+ struct ggml_context * ctx,
+ struct ggml_tensor * x,
+ struct ggml_tensor * residual,
+ struct ggml_tensor * post,
+ struct ggml_tensor * comb);
+
// custom operators
typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);
// get ith C string from array with given key_id
GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i);
- GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx);
- GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found
- GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id);
- GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id);
- GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id);
- GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id);
+ GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx);
+ GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found
+ GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id);
+ GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id);
+ GGML_API const int64_t * gguf_get_tensor_ne (const struct gguf_context * ctx, int64_t tensor_id); // returns ne, an array of GGML_MAX_DIMS elements; ne[dim] is 1 for dim >= n_dims
+ GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id);
+ GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id);
// removes key if it exists, returns id that the key had prior to removal (-1 if it didn't exist)
GGML_API int64_t gguf_remove_key(struct gguf_context * ctx, const char * key);
message(FATAL_ERROR "Unsupported ARM target OS: ${CMAKE_SYSTEM_NAME}")
endif()
elseif (GGML_SYSTEM_ARCH STREQUAL "PowerPC")
- if (CMAKE_SYSTEM_NAME MATCHES "Linux")
+ if (CMAKE_SYSTEM_NAME MATCHES "Linux|AIX")
ggml_add_cpu_backend_variant(power0)
ggml_add_cpu_backend_variant(power7_1 POWER7)
ggml_add_cpu_backend_variant(power7_2 POWER7 VSX)
ggml_add_backend(BLAS)
ggml_add_backend(CANN)
ggml_add_backend(CUDA)
+ggml_add_backend(ET)
ggml_add_backend(HIP)
ggml_add_backend(METAL)
ggml_add_backend(MUSA)
case GGML_OP_GATED_DELTA_NET: {
split_state = handle_gated_delta_net(src_ss);
} break;
+ case GGML_OP_DSV4_HC_COMB:
+ case GGML_OP_DSV4_HC_PRE:
+ case GGML_OP_DSV4_HC_POST: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
case GGML_OP_UNARY: {
split_state = handle_generic(src_ss, /*scalar_only =*/ false);
} break;
#include "ggml-openvino.h"
#endif
+#ifdef GGML_USE_ET
+#include "ggml-et.h"
+#endif
+
namespace fs = std::filesystem;
static std::string path_str(const fs::path & path) {
#ifdef GGML_USE_OPENVINO
register_backend(ggml_backend_openvino_reg());
#endif
+#ifdef GGML_USE_ET
+ register_backend(ggml_backend_et_reg());
+#endif
#ifdef GGML_USE_CPU
register_backend(ggml_backend_cpu_reg());
#endif
}
// operations with weights are preferably run on the same backend as the weights
- for (int i = 0; i < GGML_MAX_SRC; i++) {
- const struct ggml_tensor * src = tensor->src[i];
- if (src == NULL) {
- continue;
- }
- // skip ROPE since the rope freqs tensor is too small to choose a backend based on it
- // not an ideal solution
- if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
- int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor);
- // check if a backend with higher prio wants to offload the op
- if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) {
- for (int b = 0; b < src_backend_id; b++) {
- if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) {
- SET_CAUSE(tensor, "1.off");
- return b;
+ // TODO: there are exceptions (see below) - not an ideal solution
+ bool allow = true;
+
+ // skip ROPE since the rope freqs tensor is too small to choose a backend based on it
+ allow = allow && tensor->op != GGML_OP_ROPE;
+
+ // skip FLASH_ATTN_EXT since the sinks tensor is too small to choose a based based on it
+ allow = allow && tensor->op != GGML_OP_FLASH_ATTN_EXT;
+
+ if (allow) {
+ for (int i = 0; i < GGML_MAX_SRC; i++) {
+ const struct ggml_tensor * src = tensor->src[i];
+ if (src == NULL) {
+ continue;
+ }
+ if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
+ int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor);
+ // check if a backend with higher prio wants to offload the op
+ if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) {
+ for (int b = 0; b < src_backend_id; b++) {
+ if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) {
+ SET_CAUSE(tensor, "1.off");
+ return b;
+ }
}
}
+ SET_CAUSE(tensor, "1.wgt%d", i);
+ return src_backend_id;
}
- SET_CAUSE(tensor, "1.wgt%d", i);
- return src_backend_id;
}
}
+#include "ggml.h"
#include "ggml-impl.h"
#include "ggml-blas.h"
#include "ggml-backend-impl.h"
// TODO: find the optimal value
const int64_t min_batch = 32;
+ // default back to CPU fast path
+ // see: https://github.com/ggml-org/llama.cpp/issues/25565
+ if (ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
+ return false;
+ }
+
return ggml_is_contiguous(src0) &&
ggml_is_contiguous(src1) &&
src1->type == GGML_TYPE_F32 &&
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/)
set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}")
endif()
if (NOT SME_ENABLED MATCHES -1)
- list(APPEND GGML_KLEIDIAI_SOURCES
+ list(APPEND GGML_KLEIDIAI_SME_SOURCES
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa_asm.S
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot_asm.S
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa_asm.S)
+ set_source_files_properties(${GGML_KLEIDIAI_SME_SOURCES}
+ PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme")
+ list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME_SOURCES})
+
+ list(APPEND GGML_KLEIDIAI_SME2_SOURCES
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa_asm.S
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa_asm.S
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa_asm.S
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa_asm.S
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme_asm.S
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.c
+ ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme_asm.S
${KLEIDIAI_SRC}/kai/kai_common_sme_asm.S)
- set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}+sve+sve2+sme2+fp16")
+ set_source_files_properties(${GGML_KLEIDIAI_SME2_SOURCES}
+ PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme2+fp16")
+ list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME2_SOURCES})
+ set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}")
endif()
if (NOT SVE_ENABLED MATCHES -1)
bool has_sve2 = false;
bool has_i8mm = false;
bool has_sme = false;
+ bool has_sme2 = false;
aarch64_features() {
#if defined(__linux__)
has_sme = static_cast<bool>(oldp);
}
+ if (sysctlbyname("hw.optional.arm.FEAT_SME2", &oldp, &size, NULL, 0) == 0) {
+ has_sme2 = static_cast<bool>(oldp);
+ }
+
// Apple apparently does not implement SVE yet
#endif
}
{
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_DSV4_HC_COMB:
+ {
+ ggml_compute_forward_dsv4_hc_comb(params, tensor);
+ } break;
+ case GGML_OP_DSV4_HC_PRE:
+ {
+ ggml_compute_forward_dsv4_hc_pre(params, tensor);
+ } break;
+ case GGML_OP_DSV4_HC_POST:
+ {
+ ggml_compute_forward_dsv4_hc_post(params, tensor);
+ } break;
case GGML_OP_MAP_CUSTOM1:
{
ggml_compute_forward_map_custom1(params, tensor);
case GGML_OP_COUNT_EQUAL:
case GGML_OP_SOLVE_TRI:
case GGML_OP_GATED_DELTA_NET:
+ case GGML_OP_DSV4_HC_COMB:
+ case GGML_OP_DSV4_HC_PRE:
+ case GGML_OP_DSV4_HC_POST:
{
n_tasks = n_threads;
} break;
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;
} break;
case GGML_OP_OUT_PROD:
{
- if (ggml_is_quantized(node->src[0]->type)) {
+ if (ggml_is_quantized(node->src[0]->type) ||
+ node->src[0]->type == GGML_TYPE_F16) {
+ cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
+ }
+ } break;
+ case GGML_OP_SET_ROWS:
+ {
+ if (node->src[0]->type == GGML_TYPE_F16 && node->type != GGML_TYPE_F16) {
cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
}
} break;
{
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;
}
#endif
}
+int ggml_cpu_has_sme2(void) {
+#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SME2)
+ return 1;
+#else
+ return 0;
+#endif
+}
+
void ggml_cpu_init(void) {
// needed to initialize ggml_time
{
return max_bias == 0.0f;
}
case GGML_OP_IM2COL_BACK:
- return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32;
+ return src0->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
case GGML_OP_GET_ROWS_BACK:
return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16;
case GGML_OP_OUT_PROD:
- return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
+ return (src0->type == GGML_TYPE_F32 ||
+ ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
default:
return true;
if (ggml_cpu_has_sme()) {
features.push_back({ "SME", "1" });
}
+ if (ggml_cpu_has_sme2()) {
+ features.push_back({ "SME2", "1" });
+ }
if (ggml_cpu_has_riscv_v()) {
features.push_back({ "RISCV_V", "1" });
}
#include "kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.h"
#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
+#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.h"
+#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.h"
#include "kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.h"
#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.h"
#include "kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.h"
#include "kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.h"
#include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
#include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
+#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
+#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
#include "kai_lhs_pack_bf16p2vlx2_f32_sme.h"
+#include "kai_lhs_pack_f32p2vlx1_f32_sme.h"
#include "kai_lhs_quant_pack_qsi8d32p_f32.h"
#include "kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.h"
#include "kai_lhs_quant_pack_qsi8d32p_f32_neon.h"
#include "kai_lhs_quant_pack_qai8dxp_f32.h"
#include "kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.h"
+#include "kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
#include "kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.h"
#include "kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.h"
#include "kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h"
/* .packed_stride_ex = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
/* .pack_func_ex = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
},
- /* .required_cpu = */ CPU_FEATURE_SME,
+ /* .required_cpu = */ CPU_FEATURE_SME2,
/* .lhs_type = */ GGML_TYPE_F32,
/* .rhs_type = */ GGML_TYPE_Q4_0,
/* .op_type = */ GGML_TYPE_F32,
/* .packed_stride_ex = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme>,
/* .pack_func_ex = */ &rhs_pack_fn13<kai_run_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme>,
},
- /* .required_cpu = */ CPU_FEATURE_SME,
+ /* .required_cpu = */ CPU_FEATURE_SME2,
/* .lhs_type = */ GGML_TYPE_F32,
/* .rhs_type = */ GGML_TYPE_F16,
/* .op_type = */ GGML_TYPE_F32,
/* .packed_stride_ex = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
/* .pack_func_ex = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
},
+ /* .required_cpu = */ CPU_FEATURE_SME2,
+ /* .lhs_type = */ GGML_TYPE_F32,
+ /* .rhs_type = */ GGML_TYPE_Q8_0,
+ /* .op_type = */ GGML_TYPE_F32,
+ },
+ {
+ /* SME GEMM (pure SME, no SME2 required) */
+ {
+ /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_mr = */ kai_get_mr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_nr = */ kai_get_nr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_kr = */ kai_get_kr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_sr = */ kai_get_sr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+ /* .get_lhs_offset_ex = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa>,
+ /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa>,
+ /* .run_kernel_ex = */ &kernel_run_float_fn10<kai_run_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa>,
+ },
+ /* .gemm_lhs_info = */ {
+ /* .get_offset = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+ /* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+ /* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+ /* .pack_func_ex = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+ },
+ /* SME GEMV (pure SME, no SME2 required) */
+ {
+ /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_mr = */ kai_get_mr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_nr = */ kai_get_nr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_kr = */ kai_get_kr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_sr = */ kai_get_sr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+ /* .get_lhs_offset_ex = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot>,
+ /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot>,
+ /* .run_kernel_ex = */ &kernel_run_float_fn10<kai_run_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot>,
+ },
+ /* .gemv_lhs_info = */ {
+ /* .get_offset = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+ /* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+ /* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+ /* .pack_func_ex = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+ },
+ /* .rhs_info = */ {
+ /* .packed_stride = */ kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon,
+ /* .to_float = */ dequantize_row_qsi8cxp,
+ /* .packed_size_ex = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+ /* .packed_stride_ex = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+ /* .pack_func_ex = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+ },
/* .required_cpu = */ CPU_FEATURE_SME,
/* .lhs_type = */ GGML_TYPE_F32,
/* .rhs_type = */ GGML_TYPE_Q8_0,
{ /* Sentinel */ }
};
+static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
+#if defined(__ARM_FEATURE_SME)
+ {
+ /* SME2 GEMM */
+ {
+ /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_mr = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_lhs_offset_ex = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa>,
+ /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa>,
+ /* .run_kernel_ex = */ &kernel_run_fn10<kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa>,
+ },
+ /* .gemm_lhs_info = */ {
+ /* .get_offset = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+ /* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .pack_func_ex = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+ },
+ /* SME GEMV */
+ {
+ /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_mr = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+ /* .get_lhs_offset_ex = */ nullptr,
+ /* .get_rhs_packed_offset_ex = */ nullptr,
+ /* .run_kernel_ex = */ nullptr,
+ },
+ /* .gemv_lhs_info = */ {
+ /* .get_offset = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+ /* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .pack_func_ex = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+ },
+ /* .rhs_info = */ {
+ /* .packed_stride = */ nullptr,
+ /* .to_float = */ nullptr,
+ /* .packed_size_ex = */ &rhs_ps_fn2<kai_get_rhs_packed_size_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+ /* .packed_stride_ex = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+ /* .pack_func_ex = */ &rhs_pack_fn13<kai_run_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+ },
+ /* .required_cpu = */ CPU_FEATURE_SME2,
+ /* .lhs_type = */ GGML_TYPE_F32,
+ /* .rhs_type = */ GGML_TYPE_F32,
+ /* .op_type = */ GGML_TYPE_F32,
+ },
+ {
+ /* SME GEMM */
+ {
+ /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_mr = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_lhs_offset_ex = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa>,
+ /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa>,
+ /* .run_kernel_ex = */ &kernel_run_fn10<kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa>,
+ },
+ /* .gemm_lhs_info = */ {
+ /* .get_offset = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+ /* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .pack_func_ex = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+ },
+ /* SME GEMV */
+ {
+ /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_mr = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+ /* .get_lhs_offset_ex = */ nullptr,
+ /* .get_rhs_packed_offset_ex = */ nullptr,
+ /* .run_kernel_ex = */ nullptr,
+ },
+ /* .gemv_lhs_info = */ {
+ /* .get_offset = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+ /* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+ /* .pack_func_ex = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+ },
+ /* .rhs_info = */ {
+ /* .packed_stride = */ nullptr,
+ /* .to_float = */ nullptr,
+ /* .packed_size_ex = */ &rhs_ps_fn2<kai_get_rhs_packed_size_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+ /* .packed_stride_ex = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+ /* .pack_func_ex = */ &rhs_pack_fn13<kai_run_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+ },
+ /* .required_cpu = */ CPU_FEATURE_SME,
+ /* .lhs_type = */ GGML_TYPE_F32,
+ /* .rhs_type = */ GGML_TYPE_F32,
+ /* .op_type = */ GGML_TYPE_F32,
+ },
+#endif
+ { /* Sentinel */ }
+};
+
ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor) {
ggml_kleidiai_kernels * kernel = nullptr;
if (tensor->src[0]->type == GGML_TYPE_Q8_0) {
try_table(gemm_gemv_kernels_q8);
+ } else if (tensor->src[0]->type == GGML_TYPE_F32) {
+ try_table(ggml_kleidiai_kernels_f32);
} else {
try_table(gemm_gemv_kernels);
}
#else
GGML_UNUSED(gemm_gemv_kernels);
GGML_UNUSED(gemm_gemv_kernels_q8);
+ GGML_UNUSED(ggml_kleidiai_kernels_f32);
GGML_UNUSED(cpu_features);
#endif
}
return kernels;
}
+
+ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features) {
+ ggml_kleidiai_kernels * kernels = nullptr;
+
+#if defined(__ARM_FEATURE_SME)
+ for (size_t i = 0; i < NELEMS(ggml_kleidiai_kernels_f32) - 1; ++i) {
+ if ((features & ggml_kleidiai_kernels_f32[i].required_cpu) == ggml_kleidiai_kernels_f32[i].required_cpu) {
+ kernels = &ggml_kleidiai_kernels_f32[i];
+ break;
+ }
+ }
+#else
+ GGML_UNUSED(features);
+#endif
+
+ return kernels;
+}
CPU_FEATURE_DOTPROD = 1,
CPU_FEATURE_I8MM = 2,
CPU_FEATURE_SVE = 4,
- CPU_FEATURE_SME = 8
+ CPU_FEATURE_SME = 8,
+ CPU_FEATURE_SME2 = 16
};
inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {
size_t m_idx_start, const void * lhs, size_t lhs_stride, void * lhs_packed);
};
+enum rhs_repack_mode {
+ RHS_REPACK_PER_KERNEL,
+ RHS_REPACK_SHARED,
+ RHS_REPACK_SINGLE_ONLY,
+};
+
struct rhs_packing_info {
size_t (*packed_stride)(size_t k, size_t nr, size_t kr, size_t bl);
void (*pack_func_ex)(size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t bl,
size_t rhs_stride, const void * rhs, const void * bias, const void * scale, void * rhs_packed, size_t extra_bytes, const void * params);
+
+ rhs_repack_mode repack_mode = RHS_REPACK_PER_KERNEL;
};
struct ggml_kleidiai_kernels {
ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor);
ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features);
ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features);
+ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features);
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
+#ifndef HWCAP2_SME2
+#define HWCAP2_SME2 (1UL << 37)
+#endif
#elif defined(__APPLE__)
#include <string_view>
#include <sys/sysctl.h>
cpu_feature features;
ggml_kleidiai_kernels * kernels_q4;
ggml_kleidiai_kernels * kernels_q8;
+ ggml_kleidiai_kernels * kernels_f32;
int sme_thread_cap; // <= 0 means “SME disabled/unknown”;
int thread_hint; // <= 0 means “no hint”
int chunk_multiplier;
-} static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, 0, -1, 4 };
+} static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, nullptr, 0, -1, 4 };
+
+static inline bool is_sme_family(cpu_feature f) {
+ return (f & (CPU_FEATURE_SME | CPU_FEATURE_SME2)) != CPU_FEATURE_NONE;
+}
static const char* cpu_feature_to_string(cpu_feature f) {
if (f == CPU_FEATURE_NONE) {
return "NONE";
+ } else if ((f & CPU_FEATURE_SME2) == CPU_FEATURE_SME2) {
+ return "SME2";
} else if ((f & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
return "SME";
} else if ((f & CPU_FEATURE_SVE) == CPU_FEATURE_SVE) {
}
}
}
- return 1;
+ return 0;
#else
- return 1;
+ return 0;
#endif
}
const char *env_threads = getenv("GGML_TOTAL_THREADS");
const char *env_chunk_mult = getenv("GGML_KLEIDIAI_CHUNK_MULTIPLIER");
- const bool cpu_has_sme = ggml_cpu_has_sme();
size_t detected_smcus = 0;
ctx.features = (ggml_cpu_has_dotprod() ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
}
// SME policy:
- // - If CPU doesn't support SME: SME always off.
- // - Else:
- // - env unset => auto-detect cores; enable if detected > 0.
- // - env=0 => force off.
- // - env>0 => force N cores (skip detection).
+ // - env unset => auto-detect SMCUs; enable SME only if detected > 0.
+ // - env=0 => force off.
+ // - env>0 => force N cores, if the binary was built with SME.
int sme_cores = 0;
bool sme_env_ok = false;
bool sme_env_set = (env_sme != nullptr);
- if (!cpu_has_sme) {
- if (sme_env_set) {
- bool ok = false;
- int req = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
- if (ok && req > 0) {
- GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but SME is not supported on this CPU; disabling SME\n", req);
- }
- }
- sme_cores = 0;
- } else {
- if (sme_env_set) {
- bool ok = false;
- int v = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
- sme_env_ok = ok;
-
- if (!ok) {
- GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; falling back to runtime SME-core detection\n");
- detected_smcus = detect_num_smcus();
- sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
- } else if (v == 0) {
- sme_cores = 0;
- } else {
- sme_cores = v;
- }
- } else {
+ if (sme_env_set) {
+ bool ok = false;
+ int v = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
+ sme_env_ok = ok;
+
+ if (!ok) {
+ GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; falling back to runtime SME-core detection\n");
detected_smcus = detect_num_smcus();
sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
+ } else if (v == 0) {
+ sme_cores = 0;
+ } else if (!ggml_cpu_has_sme()) {
+ GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but the binary was not built with SME; disabling SME\n", v);
+ sme_cores = 0;
+ } else {
+ sme_cores = v;
}
+ } else {
+ detected_smcus = detect_num_smcus();
+ sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
+ }
- if (!sme_env_set && sme_cores == 0) {
- GGML_LOG_WARN("kleidiai: SME supported but runtime SME-core detection returned 0; falling back to NEON\n");
- }
+ if (!sme_env_set && ggml_cpu_has_sme() && sme_cores == 0) {
+ GGML_LOG_WARN("kleidiai: runtime SME-core detection returned 0; falling back to NEON\n");
+ }
- if (sme_cores > 0) {
- ctx.features |= CPU_FEATURE_SME;
+ if (sme_cores > 0) {
+ ctx.features |= CPU_FEATURE_SME;
+#if defined(__aarch64__) && defined(__linux__)
+ // ARM guarantees SME2 implies SME, so only check SME2 when SME is enabled.
+ if (getauxval(AT_HWCAP2) & HWCAP2_SME2) {
+ ctx.features |= CPU_FEATURE_SME2;
}
+#elif defined(__aarch64__) && defined(__APPLE__)
+ int feat_sme2 = 0;
+ size_t size = sizeof(feat_sme2);
+ if (sysctlbyname("hw.optional.arm.FEAT_SME2", &feat_sme2, &size, NULL, 0) == 0 && feat_sme2) {
+ ctx.features |= CPU_FEATURE_SME2;
+ }
+#endif
}
// Kernel selection
- ctx.kernels_q4 = ggml_kleidiai_select_kernels_q4_0(ctx.features);
- ctx.kernels_q8 = ggml_kleidiai_select_kernels_q8_0(ctx.features);
+ ctx.kernels_q4 = ggml_kleidiai_select_kernels_q4_0(ctx.features);
+ ctx.kernels_q8 = ggml_kleidiai_select_kernels_q8_0(ctx.features);
+ ctx.kernels_f32 = ggml_kleidiai_select_kernels_f32(ctx.features);
if (!ctx.kernels_q4) {
GGML_LOG_INFO("kleidiai: no compatible q4 kernels found for CPU features mask %d\n", (int)ctx.features);
GGML_LOG_INFO("kleidiai: primary q8 kernel feature %s\n", cpu_feature_to_string(ctx.kernels_q8->required_cpu));
}
+ if (!ctx.kernels_f32) {
+ GGML_LOG_INFO("kleidiai: no compatible f32 kernels found for CPU features mask %d\n", (int)ctx.features);
+ } else {
+ GGML_LOG_INFO("kleidiai: primary f32 kernel feature %s\n", cpu_feature_to_string(ctx.kernels_f32->required_cpu));
+ }
+
ctx.sme_thread_cap = (ctx.features & CPU_FEATURE_SME) ? sme_cores : 0;
if (ctx.features & CPU_FEATURE_SME) {
+ const bool has_sme2 = (ctx.features & CPU_FEATURE_SME2) != CPU_FEATURE_NONE;
if (sme_env_set && sme_env_ok && sme_cores > 0) {
- GGML_LOG_INFO("kleidiai: SME enabled (GGML_KLEIDIAI_SME=%d override)\n", sme_cores);
+ GGML_LOG_INFO("kleidiai: SME%s enabled (GGML_KLEIDIAI_SME=%d override)\n",
+ has_sme2 ? "2" : "", sme_cores);
} else {
- GGML_LOG_INFO("kleidiai: SME enabled (runtime-detected SME cores=%d)\n", sme_cores);
+ GGML_LOG_INFO("kleidiai: SME%s enabled (runtime-detected SME cores=%d)\n",
+ has_sme2 ? "2" : "", sme_cores);
}
} else {
GGML_LOG_INFO("kleidiai: SME disabled\n");
return b == 0 ? 0 : (a + b - 1) / b;
}
+static inline size_t kleidiai_chunk_cols(size_t n, int nth_total, bool disable_chunking, size_t n_step) {
+ const size_t multiplier = (nth_total == 1 || disable_chunking) ? 1 : std::max<size_t>(1, (size_t) ctx.chunk_multiplier);
+ const size_t divisor = std::max<size_t>(1, (size_t) nth_total * multiplier);
+ const size_t chunk_cols = align_up(std::max<size_t>(1, ceil_div_size(n, divisor)), n_step);
+ return chunk_cols ? chunk_cols : n_step;
+}
+
struct kleidiai_block_args {
size_t lhs_bl;
size_t rhs_bl;
return ctx.kernels_q8;
}
+static inline ggml_kleidiai_kernels * kleidiai_primary_kernel_f32() {
+ return ctx.kernels_f32;
+}
+
template <typename SelectFallback>
static int kleidiai_collect_kernel_chain_common(
ggml_kleidiai_kernels * primary,
}
out[count++] = primary;
- if ((primary->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
- const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~CPU_FEATURE_SME);
+ if (primary->rhs_info.repack_mode == RHS_REPACK_SINGLE_ONLY) {
+ return count;
+ }
+
+ if (is_sme_family(primary->required_cpu)) {
+ const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~CPU_FEATURE_SME & ~CPU_FEATURE_SME2);
if (fallback_mask != CPU_FEATURE_NONE) {
ggml_kleidiai_kernels * fallback = select_fallback(fallback_mask);
if (fallback && fallback != primary &&
+ fallback->rhs_info.repack_mode != RHS_REPACK_SINGLE_ONLY &&
fallback->lhs_type == primary->lhs_type &&
fallback->rhs_type == primary->rhs_type &&
fallback->op_type == primary->op_type) {
[&](cpu_feature mask) { return ggml_kleidiai_select_kernels_q8_0(mask); });
}
+static int kleidiai_collect_f32_chain(std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> & out) {
+ ggml_kleidiai_kernels * primary = kleidiai_primary_kernel_f32();
+ return kleidiai_collect_kernel_chain_common(primary, ctx.features, out,
+ [&](cpu_feature mask) { return ggml_kleidiai_select_kernels_f32(mask); });
+}
+
static inline int64_t ggml_ne(const ggml_tensor * tensor, int dim) {
GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS);
return tensor->ne[dim];
return true;
}
+ if (op->src[0]->type == GGML_TYPE_F32) {
+ size_t cursor = 0;
+ bool any_slot = false;
+
+ for (int slot = 0; slot < slot_count; ++slot) {
+ ggml_kleidiai_kernels * kernels = kernel_chain[slot];
+ lhs_packing_info * lhs_info = &kernels->gemm_lhs_info;
+ kernel_info * kernel = &kernels->gemm;
+
+ if (!lhs_info || !lhs_info->packed_size_ex || !kernel) {
+ return false;
+ }
+
+ const size_t mr = kernel->get_mr();
+ const size_t kr = kernel->get_kr();
+ const size_t sr = kernel->get_sr();
+
+ cursor = align_up(cursor, GGML_KLEIDIAI_PACK_ALIGN);
+ cursor += lhs_info->packed_size_ex(m, k, 0, mr, kr, sr);
+ any_slot = true;
+ }
+
+ if (!any_slot) {
+ return false;
+ }
+
+ size = cursor;
+ return true;
+ }
+
if (op->src[0]->type == GGML_TYPE_F16) {
const int64_t lhs_batch_size0 = op->src[1]->ne[2];
const int64_t rhs_batch_size0 = op->src[0]->ne[2];
if (dst->op == GGML_OP_MUL_MAT) {
if (dst->src[0]->type == GGML_TYPE_Q4_0 || dst->src[0]->type == GGML_TYPE_Q8_0) {
return compute_forward_qx(params, dst);
+ } else if (dst->src[0]->type == GGML_TYPE_F32) {
+ return compute_forward_f32(params, dst);
} else if (dst->src[0]->type == GGML_TYPE_F16) {
return compute_forward_fp16(params, dst);
}
return false;
}
+ bool compute_forward_f32(ggml_compute_params * params, struct ggml_tensor * dst) {
+ GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32);
+
+ const ggml_tensor * src0 = dst->src[0];
+ const ggml_tensor * src1 = dst->src[1];
+
+ GGML_TENSOR_BINARY_OP_LOCALS
+
+ if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ ggml_kleidiai_kernels * kernels = kleidiai_primary_kernel_f32();
+ if (!kernels) {
+ return false;
+ }
+
+ kernel_info * kernel = &kernels->gemm;
+ lhs_packing_info * lhs_info = &kernels->gemm_lhs_info;
+
+ if (!kernel || !lhs_info || !lhs_info->get_offset || !lhs_info->get_packed_offset_ex ||
+ !lhs_info->packed_size_ex || !lhs_info->pack_func_ex ||
+ !kernel->get_rhs_packed_offset_ex || !kernel->run_kernel_ex || !kernel->get_dst_offset) {
+ return false;
+ }
+
+ const kleidiai_weight_header * header = kleidiai_weight_header_from_ptr(src0->data);
+ const bool has_header = kleidiai_is_weight_header_valid(header);
+
+ const uint8_t * rhs_base = has_header ? kleidiai_weight_slot_ptr(header, 0)
+ : static_cast<const uint8_t *>(src0->data);
+ if (!rhs_base) {
+ return false;
+ }
+
+ const int nth = params->nth > 0 ? params->nth : 1;
+ const int ith = params->ith;
+
+ const size_t k = ne00;
+ const size_t m = ne11;
+ const size_t n = ne01;
+
+ const size_t mr = kernel->get_mr();
+ const size_t kr = kernel->get_kr();
+ const size_t sr = kernel->get_sr();
+
+ const size_t lhs_packed_size = lhs_info->packed_size_ex(m, k, 0, mr, kr, sr);
+ GGML_ASSERT(lhs_packed_size <= params->wsize);
+
+ uint8_t * lhs_packed = static_cast<uint8_t *>(params->wdata);
+ const size_t dst_stride = dst->nb[1];
+ const size_t n_step = kernel->get_n_step() ? kernel->get_n_step() : 1;
+ const bool disable_chunking = ggml_is_numa();
+ GGML_ASSERT(n <= (size_t) INT_MAX);
+
+ for (int64_t batch_idx = 0; batch_idx < ne12; ++batch_idx) {
+ const uint8_t * lhs_batch_base = static_cast<const uint8_t *>(src1->data) + batch_idx * src1->nb[2];
+ uint8_t * dst_batch_base = static_cast<uint8_t *>(dst->data) + batch_idx * dst->nb[2];
+
+ {
+ const int64_t m_roundup_mr = kai_roundup((int64_t)m, (int64_t)mr);
+ int64_t max_threads = mr ? (m_roundup_mr / (int64_t)mr) : nth;
+ max_threads = std::max<int64_t>(1, max_threads);
+ const int64_t use_threads = std::min<int64_t>(nth, max_threads);
+
+ if (ith < use_threads) {
+ const int64_t num_m_per_thread0 = round_down((size_t)(m_roundup_mr / use_threads), mr);
+ const int64_t num_m_per_threadN_1 = (int64_t)m - (use_threads - 1) * num_m_per_thread0;
+
+ const int64_t m_start = (int64_t)ith * num_m_per_thread0;
+ const int64_t m_count = (ith == use_threads - 1) ? num_m_per_threadN_1 : num_m_per_thread0;
+
+ const size_t base_packed_off = lhs_info->get_packed_offset_ex(m_start, k, 0, mr, kr, sr);
+ const size_t next_block_off = lhs_info->get_packed_offset_ex(m_start + mr, k, 0, mr, kr, sr);
+ const size_t row_stride_bytes = mr ? (next_block_off - base_packed_off) / mr : 0;
+
+ int64_t remaining = m_count;
+ int64_t cur = m_start;
+
+ while (remaining > 0) {
+ const int64_t take = std::min<int64_t>((int64_t)m - cur, remaining);
+ const size_t src_off = lhs_info->get_offset(cur, src1->nb[1]);
+ const void * src_ptr = lhs_batch_base + src_off;
+ const size_t dst_off = base_packed_off + (size_t)(cur - m_start) * row_stride_bytes;
+ void * dst_ptr = lhs_packed + dst_off;
+
+ lhs_info->pack_func_ex(take, k, 0, mr, kr, sr, 0, src_ptr, src1->nb[1], dst_ptr);
+
+ cur += take;
+ remaining -= take;
+ }
+ }
+ }
+
+ if (ith == 0) {
+ ggml_threadpool_chunk_set(params->threadpool, 0);
+ }
+
+ ggml_barrier(params->threadpool);
+
+ const size_t chunk_cols = kleidiai_chunk_cols(n, nth, disable_chunking, n_step);
+ GGML_ASSERT(chunk_cols <= (size_t) INT_MAX);
+
+ int current_col = ggml_threadpool_chunk_add(params->threadpool, (int) chunk_cols);
+ while ((size_t) current_col < n) {
+ const size_t n_start = (size_t) current_col;
+ const size_t n_to_process = std::min(chunk_cols, n - n_start);
+
+ if (n_to_process > 0) {
+ const size_t lhs_packed_offset = lhs_info->get_packed_offset_ex(0, k, 0, mr, kr, sr);
+ const size_t rhs_packed_offset = kernel->get_rhs_packed_offset_ex(n_start, k, 0);
+ const size_t dst_offset = kernel->get_dst_offset(0, n_start, dst_stride);
+
+ const void * lhs_ptr = lhs_packed + lhs_packed_offset;
+ const void * rhs_ptr = rhs_base + rhs_packed_offset;
+ float * dst_ptr = reinterpret_cast<float *>(dst_batch_base + dst_offset);
+
+ kernel->run_kernel_ex(m, n_to_process, k, 0,
+ lhs_ptr,
+ rhs_ptr,
+ dst_ptr,
+ dst_stride,
+ sizeof(float),
+ -FLT_MAX,
+ FLT_MAX);
+ }
+
+ current_col = ggml_threadpool_chunk_add(params->threadpool, (int) chunk_cols);
+ }
+
+ if (batch_idx != ne12 - 1) {
+ ggml_barrier(params->threadpool);
+ }
+ }
+
+ return true;
+ }
+
bool compute_forward_fp16(ggml_compute_params * params, struct ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
int sme_slot = -1;
for (int i = 0; i < runtime_count; ++i) {
- if ((runtime[i].kernels->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
+ if (is_sme_family(runtime[i].kernels->required_cpu)) {
sme_slot = i;
break;
}
}
int non_sme_slot = -1;
for (int i = 0; i < runtime_count; ++i) {
- if ((runtime[i].kernels->required_cpu & CPU_FEATURE_SME) != CPU_FEATURE_SME) {
+ if (!is_sme_family(runtime[i].kernels->required_cpu)) {
non_sme_slot = i;
break;
}
// Recompute SME slot based on the collapsed runtime[0]
sme_slot = -1;
if (runtime_count > 0 &&
- (runtime[0].kernels->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
+ is_sme_family(runtime[0].kernels->required_cpu)) {
sme_slot = 0;
}
}
public:
int repack(struct ggml_tensor * tensor, const void * data, size_t data_size) {
- GGML_ASSERT(tensor->type == GGML_TYPE_Q4_0 || tensor->type == GGML_TYPE_Q8_0);
+ GGML_ASSERT(tensor->type == GGML_TYPE_Q4_0 || tensor->type == GGML_TYPE_Q8_0 || tensor->type == GGML_TYPE_F32);
const size_t n = tensor->ne[1];
const size_t k = tensor->ne[0];
std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> kernel_chain;
const bool want_q8 = tensor->type == GGML_TYPE_Q8_0;
- const int slot_total = want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
- : kleidiai_collect_q4_chain(kernel_chain);
+ const bool want_f32 = tensor->type == GGML_TYPE_F32;
+ const int slot_total = want_f32 ? kleidiai_collect_f32_chain(kernel_chain)
+ : want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
+ : kleidiai_collect_q4_chain(kernel_chain);
const bool allow_fallback = kleidiai_pack_fallback_allowed();
std::vector<int8_t> qdata;
std::vector<float> scales;
+ std::vector<float> bias;
if (want_q8 && slot_total > 0) {
qdata.resize(n * k, 0);
}
}
+ if (want_f32 && slot_total > 0) {
+ bias.resize(n, 0.0f);
+ }
+
for (int slot = 0; slot < slot_total && slot < GGML_KLEIDIAI_MAX_KERNEL_SLOTS; ++slot) {
if (!allow_fallback && slot > 0) {
break;
const size_t sr = kernel->get_sr();
const ggml_type rhs_type = kernels->rhs_type;
const size_t block_len = rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
- rhs_type == GGML_TYPE_Q4_0 ? QK4_0 : 0;
- if (block_len == 0) {
+ rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
+ rhs_type == GGML_TYPE_F32 ? 0 : SIZE_MAX;
+ if (block_len == SIZE_MAX) {
continue;
}
rhs_info->pack_func_ex(1, n, k, nr, kr, sr, 0, 0,
qdata.data(), nullptr, scales.data(),
dst_ptr, 0, ¶ms);
+ } else if (rhs_type == GGML_TYPE_F32) {
+ rhs_info->pack_func_ex(1, n, k, nr, kr, sr, 0, tensor->nb[1],
+ data, bias.data(), nullptr,
+ dst_ptr, 0, nullptr);
} else {
continue;
}
static size_t ggml_backend_cpu_kleidiai_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor) {
GGML_UNUSED(buft);
- if (tensor->type != GGML_TYPE_Q4_0 && tensor->type != GGML_TYPE_Q8_0) {
+ if (tensor->type != GGML_TYPE_Q4_0 && tensor->type != GGML_TYPE_Q8_0 && tensor->type != GGML_TYPE_F32) {
return ggml_nbytes(tensor);
}
std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> kernel_chain;
const bool want_q8 = tensor->type == GGML_TYPE_Q8_0;
- const int slot_total = want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
- : kleidiai_collect_q4_chain(kernel_chain);
+ const bool want_f32 = tensor->type == GGML_TYPE_F32;
+ const int slot_total = want_f32 ? kleidiai_collect_f32_chain(kernel_chain)
+ : want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
+ : kleidiai_collect_q4_chain(kernel_chain);
const bool allow_fallback = kleidiai_pack_fallback_allowed();
size_t slot_count = 0;
const ggml_type rhs_type = kernels->rhs_type;
const size_t block_len = rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
- rhs_type == GGML_TYPE_Q8_0 ? QK8_0 : 0;
- if (block_len == 0) {
+ rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
+ rhs_type == GGML_TYPE_F32 ? 0 : SIZE_MAX;
+ if (block_len == SIZE_MAX) {
continue;
}
bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override {
std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> kernel_chain;
const int slot_total = kleidiai_collect_kernel_chain(op, kernel_chain);
- if ((op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_GET_ROWS) &&
- (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0) &&
+ const bool src0_is_kleidiai =
op->src[0]->buffer &&
(ggml_n_dims(op->src[0]) == 2) &&
op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type() &&
- slot_total > 0) {
+ slot_total > 0;
+
+ if ((op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_GET_ROWS) &&
+ (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0 || op->src[0]->type == GGML_TYPE_F32) &&
+ src0_is_kleidiai) {
if (op->src[0]->type == GGML_TYPE_Q4_0 && ctx.kernels_q4 == nullptr) {
return false;
}
if (op->src[0]->type == GGML_TYPE_Q8_0 && ctx.kernels_q8 == nullptr) {
return false;
}
+ if (op->src[0]->type == GGML_TYPE_F32 && ctx.kernels_f32 == nullptr) {
+ return false;
+ }
if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
return false;
}
- if ((op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) &&
- ggml_ne(op->src[1], 3) == 1) {
- return true;
+
+ if (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0) {
+ if ((op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) &&
+ ggml_ne(op->src[1], 3) == 1) {
+ return true;
+ }
+ return false;
}
+
+ if (op->op != GGML_OP_MUL_MAT || op->src[1]->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ return true;
}
+
return false;
}
if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type()) {
return (ggml::cpu::tensor_traits *) op->src[0]->extra;
} else {
+ // KleidiAI only has kernels for Q4_0 and Q8_0. For a quantized weight of any
+ // other type (K-quants, IQ) it declines the op and returns nullptr below, so
+ // KleidiAI does not accelerate it. Another CPU backend may still take the op,
+ // and this can run during graph planning, so the message says what KleidiAI
+ // did rather than what ends up executing. Warn once per process.
+ if (ggml_is_quantized(op->src[0]->type) &&
+ op->src[0]->type != GGML_TYPE_Q4_0 && op->src[0]->type != GGML_TYPE_Q8_0) {
+ static std::atomic<bool> warned(false);
+ if (!warned.exchange(true)) {
+ GGML_LOG_WARN("kleidiai: no kernel for tensor type %s, not accelerated by KleidiAI "
+ "(kernels available for Q4_0 and Q8_0)\n",
+ ggml_type_name(op->src[0]->type));
+ }
+ }
if (op->src[0]->type != GGML_TYPE_F16) {
return nullptr;
}
//PPC Implementation
#if defined(__MMA__)
-#define SAVE_ACC(ACC, ii, jj) \
- __builtin_mma_disassemble_acc(vec_C, ACC); \
- for (int I = 0; I < 4; I++) { \
- for (int J = 0; J < 4; J++) { \
- *((float*)(C+ii+((jj+J)*ldc)+I)) = *((float*)&vec_C[I]+J); \
- } \
- } \
-
template<typename T>
struct mma_instr;
}
void matmul(int64_t m, int64_t n) {
- mnpack(0, m, 0, n);
+ int64_t mc = 256;
+ int64_t nc = 256;
+ int64_t kc = 256;
+ #if defined(_AIX) || defined(__BIG_ENDIAN__)
+ mc = 128;
+ nc = 128;
+ kc = 128;
+ #endif
+ if (k < kc) {
+ kc = k;
+ }
+ bool can_use_tiled = (m % mc == 0) && (n % nc == 0) && (k % kc == 0);
+ if (can_use_tiled) {
+ matmul_tiled(m, n, mc, nc, kc);
+ } else {
+ mnpack(0, m, 0, n);
+ }
}
private:
+ __attribute__((always_inline))
+ inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) {
+ vec_t vec_C[4];
+ __builtin_mma_disassemble_acc(vec_C, ACC);
+ for (int I = 0; I < 4; I++) {
+ for (int J = 0; J < 4; J++) {
+ *((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J);
+ }
+ }
+ }
+
+ __attribute__((always_inline))
+ inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) {
+ vec_t vec_C[4];
+ __builtin_mma_disassemble_acc(vec_C, ACC);
+ for (int I = 0; I < 4; I++) {
+ for (int J = 0; J < 4; J++) {
+ float * c_ptr = (float *)(C+ii+((jj+J)*ldc)+I);
+ *c_ptr += *((float *)&vec_C[I]+J);
+ }
+ }
+ }
+
void vector_permute_store(vec_t *c, int numVec, unsigned char *vecOffset) {
vec_t t[8], s[8];
vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23};
j = (rows >> 3);
if (j > 0) {
do {
+ aoffsets[0] = aoffset;
if (cols == 4) {
aoffsets[0] = aoffset;
for (int it = 1; it < 4; ++it)
}
i = (cols >> 3);
if (i > 0) {
- aoffsets[0] = aoffset;
for (int it = 1; it < 8; ++it) {
aoffsets[it] = aoffsets[it-1] + lda;
}
aoffset += 8 * lda;
+
do {
for (int it = 0; it < 8; ++it)
c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]);
vector_permute_store(c_arr, 8, vecOffset);
for (int it = 0; it < 8; ++it)
- aoffsets[it] = aoffsets[it] + 8*lda;
+ aoffsets[it] = aoffsets[it] + 8;
vecOffset += 128;
i--;
} while(i > 0);
mma_instr<TA>::outer_product(&acc_1, vec_A[x], vec_B[x+4]);
}
}
- SAVE_ACC(&acc_0, ii, jj);
- SAVE_ACC(&acc_1, ii, jj+4);
+ save_acc(&acc_0, ii, jj);
+ save_acc(&acc_1, ii, jj+4);
}
void KERNEL_8x4(int64_t ii, int64_t jj) {
mma_instr<TA>::outer_product(&acc_1, vec_A[x+4], vec_B[x]);
}
}
- SAVE_ACC(&acc_0, ii, jj);
- SAVE_ACC(&acc_1, ii+4, jj);
+ save_acc(&acc_0, ii, jj);
+ save_acc(&acc_1, ii+4, jj);
}
mma_instr<TA>::outer_product(&acc_3, vec_A[x+4], vec_B[x+4]);
}
}
-
- SAVE_ACC(&acc_0, ii, jj);
- SAVE_ACC(&acc_1, ii, jj+4);
- SAVE_ACC(&acc_2, ii+4, jj);
- SAVE_ACC(&acc_3, ii+4, jj+4);
+ save_acc(&acc_0, ii, jj);
+ save_acc(&acc_1, ii, jj+4);
+ save_acc(&acc_2, ii+4, jj);
+ save_acc(&acc_3, ii+4, jj+4);
}
+ inline void MMA_16x8(vec_t * vec_A0, vec_t * vec_A1, vec_t * vec_B, acc_t * acc) {
+ for (int x = 0; x < 4; x ++) {
+ mma_instr<TA>::outer_product(&acc[0], vec_A0[x], vec_B[x]);
+ mma_instr<TA>::outer_product(&acc[1], vec_A0[x], vec_B[x+4]);
+ mma_instr<TA>::outer_product(&acc[2], vec_A0[x+4], vec_B[x]);
+ mma_instr<TA>::outer_product(&acc[3], vec_A0[x+4], vec_B[x+4]);
+ mma_instr<TA>::outer_product(&acc[4], vec_A1[x], vec_B[x]);
+ mma_instr<TA>::outer_product(&acc[5], vec_A1[x], vec_B[x+4]);
+ mma_instr<TA>::outer_product(&acc[6], vec_A1[x+4], vec_B[x]);
+ mma_instr<TA>::outer_product(&acc[7], vec_A1[x+4], vec_B[x+4]);
+ }
+ }
+ void KERNEL(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, vec_t * vec_A, vec_t * vec_B, int64_t kk) {
+ for (int64_t i = 0; i < mc; i += 16) {
+ int A_base_addr = (mc / 8) * (i / 8) * 8;
+ for (int64_t j = 0; j < nc; j += 8) {
+ int B_base_addr = (nc / 8) * (j / 8) * 8;
+ acc_t acc[8];
+ vec_t A0_block[8]; vec_t A1_block[8];
+ for (int x = 0; x < 8; x++)
+ __builtin_mma_xxsetaccz(&acc[x]);
+ for (int64_t l = 0; l < kc; l += 8) {
+ int A0_block_idx = A_base_addr + (l / 8) * 8;
+ int A1_block_idx = A0_block_idx + (mc / 8) * 8;
+ int B_block_idx = B_base_addr + (l / 8) * 8;
+ vec_t* A0_block = &vec_A[A0_block_idx];
+ vec_t* A1_block = &vec_A[A1_block_idx];
+ vec_t* B_block = &vec_B[B_block_idx];
+ MMA_16x8(A0_block, A1_block, B_block, acc);
+ }
+ if (kk == 0) {
+ save_acc(&acc[0], ii + i, jj + j);
+ save_acc(&acc[1], ii + i, jj + j + 4);
+ save_acc(&acc[2], ii + i + 4, jj + j);
+ save_acc(&acc[3], ii + i + 4, jj + j + 4);
+ save_acc(&acc[4], ii + i + 8, jj + j);
+ save_acc(&acc[5], ii + i + 8, jj + j + 4);
+ save_acc(&acc[6], ii + i + 12, jj + j);
+ save_acc(&acc[7], ii + i + 12, jj + j + 4);
+ } else {
+ add_save_acc(&acc[0], ii + i, jj + j);
+ add_save_acc(&acc[1], ii + i, jj + j + 4);
+ add_save_acc(&acc[2], ii + i + 4, jj + j);
+ add_save_acc(&acc[3], ii + i + 4, jj + j + 4);
+ add_save_acc(&acc[4], ii + i + 8, jj + j);
+ add_save_acc(&acc[5], ii + i + 8, jj + j + 4);
+ add_save_acc(&acc[6], ii + i + 12, jj + j);
+ add_save_acc(&acc[7], ii + i + 12, jj + j + 4);
+ }
+ }
+ }
+ }
template<int RM, int RN>
void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n) {
int64_t ytiles = (m - m0) / RM;
}
}
+ void matmul_tiled(int64_t m, int64_t n, int64_t mc, int64_t nc, int64_t kc) {
+ int64_t ytiles = m / mc;
+ int64_t xtiles = n / nc;
+ int64_t tiles = xtiles * ytiles;
+ int64_t duty = (tiles + nth - 1) / nth;
+ int64_t start = duty * ith;
+ int64_t end = start + duty;
+ if (end > tiles) {
+ end = tiles;
+ }
+ for (int64_t job = start; job < end; ++job) {
+ int64_t ii = (job / xtiles) * mc;
+ int64_t jj = (job % xtiles) * nc;
+ for (int64_t kk = 0; kk < k; kk += kc) {
+ vec_t A_pack[kc * mc / 8];
+ vec_t B_pack[kc * nc / 8];
+ packNormal(A + (ii * lda) + kk, lda, kc, mc, (uint8_t *)A_pack);
+ packNormal(B + (jj * ldb) + kk, ldb, kc, nc, (uint8_t *)B_pack);
+ KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack, kk);
+ }
+ }
+ }
+
template <int RM, int RN>
NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) {
int64_t ytiles = (m - m0) / RM;
mc = 32;
nc = 32;
kc = 32;
- n_chunk = 32
+ n_chunk = 32;
#endif
int64_t n_aligned = 0;
if (n % n_chunk == 0) {
const ggml_tensor * src1 = dst->src[1];
if (ggml_is_quantized(src0->type)) {
- GGML_ASSERT(ggml_is_contiguous(src0));
- GGML_ASSERT(ggml_is_contiguous(src1));
+ GGML_ASSERT(ggml_is_contiguous_rows(src0));
+ GGML_ASSERT(ggml_is_contiguous_rows(src1));
GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
}
}
}
+static void ggml_compute_forward_out_prod_f16_f32(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+
+ const ggml_tensor * src0 = dst->src[0];
+ const ggml_tensor * src1 = dst->src[1];
+
+ GGML_TENSOR_BINARY_OP_LOCALS;
+
+ const int ith = params->ith;
+ const int nth = params->nth;
+
+ GGML_ASSERT(src0->type == GGML_TYPE_F16);
+ GGML_ASSERT(src1->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ GGML_ASSERT(ne02 == ne12);
+ GGML_ASSERT(ne03 == ne13);
+ GGML_ASSERT(ne2 == ne12);
+ GGML_ASSERT(ne3 == ne13);
+
+ GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
+ GGML_ASSERT(nb0 == sizeof(float));
+
+ GGML_ASSERT(ne0 == ne00);
+ GGML_ASSERT(ne1 == ne10);
+ GGML_ASSERT(ne2 == ne02);
+ GGML_ASSERT(ne3 == ne03);
+
+ if (ith == 0) {
+ ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0);
+ }
+ ggml_barrier(params->threadpool);
+
+ const int64_t nr = ne1*ne2*ne3;
+ const int64_t dr = (nr + nth - 1)/nth;
+ const int64_t ir0 = dr*ith;
+ const int64_t ir1 = MIN(ir0 + dr, nr);
+
+ float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith;
+
+ for (int64_t ir = ir0; ir < ir1; ++ir) {
+ const int64_t i3 = ir/(ne2*ne1);
+ const int64_t i2 = (ir - i3*ne2*ne1)/ne1;
+ const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1);
+
+ const int64_t i02 = i2;
+ const int64_t i03 = i3;
+
+ const int64_t i12 = i2;
+ const int64_t i13 = i3;
+
+ float * d = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb3));
+
+ for (int64_t i01 = 0; i01 < ne01; ++i01) {
+ const int64_t i11 = i01;
+ ggml_fp16_t * s0 = (ggml_fp16_t *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03));
+ float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13));
+ ggml_fp16_to_fp32_row(s0, wdata, ne0);
+ ggml_vec_mad_f32(ne0, d, wdata, *s1);
+ }
+ }
+}
+
void ggml_compute_forward_out_prod(
const ggml_compute_params * params,
ggml_tensor * dst) {
} break;
case GGML_TYPE_F16:
{
- GGML_ABORT("fatal error"); // todo
- // ggml_compute_forward_out_prod_f16_f32(params, dst);
- }
+ ggml_compute_forward_out_prod_f16_f32(params, dst);
+ } break;
case GGML_TYPE_F32:
{
ggml_compute_forward_out_prod_f32(params, dst);
assert(ne0 == nc);
assert(ne2 == ne02);
assert(ne3 == ne03);
- GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
+ GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
assert(ne02 % ne11 == 0);
assert(ne03 % ne12 == 0);
(const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc);
} else if constexpr (std::is_same_v<src_t, ggml_fp16_t>) {
- memcpy(
+ if (dst->type == GGML_TYPE_F16) {
+ memcpy(
((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3),
((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
rs);
+ } else {
+ float * wdata = (float *) params->wdata + (nc + CACHE_LINE_SIZE_F32) * ith;
+ ggml_fp16_to_fp32_row(
+ (const ggml_fp16_t *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
+ wdata, nc);
+ from_float(wdata,
+ ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc);
+ }
} else {
GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type));
}
} break;
case GGML_TYPE_F16:
{
- if (dst->type == GGML_TYPE_F16) {
- if (src1->type == GGML_TYPE_I64) {
- ggml_compute_forward_set_rows_impl<ggml_fp16_t, int64_t>(params, dst);
- } else if (src1->type == GGML_TYPE_I32) {
- ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(params, dst);
- } else {
- GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
- }
+ if (src1->type == GGML_TYPE_I64) {
+ ggml_compute_forward_set_rows_impl<ggml_fp16_t, int64_t>(params, dst);
+ } else if (src1->type == GGML_TYPE_I32) {
+ ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(params, dst);
} else {
- GGML_ABORT("dst->type = %d (%s) not supported with src0->type = %d (%s)", dst->type, ggml_type_name(dst->type), src0->type, ggml_type_name(src0->type));
+ GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
}
} break;
default:
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
- GGML_ASSERT(src0->type == GGML_TYPE_F16);
GGML_ASSERT(src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32);
GGML_ASSERT( dst->type == GGML_TYPE_F16);
int ofs0 = is_2D ? nb13 : nb12;
int ofs1 = is_2D ? nb12 : nb11;
- GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
GGML_ASSERT(nb10 == ggml_type_size(src1->type));
// im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW]
const ggml_tensor * src1 = dst->src[1]; // convolution kernel
GGML_ASSERT(src0->type == GGML_TYPE_F32);
- GGML_ASSERT(src1->type == GGML_TYPE_F32);
+ GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
GGML_ASSERT( dst->type == GGML_TYPE_F32);
GGML_TENSOR_BINARY_OP_LOCALS;
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
- GGML_ASSERT(src0->type == GGML_TYPE_F16);
GGML_ASSERT(src1->type == GGML_TYPE_F32);
GGML_ASSERT( dst->type == GGML_TYPE_F16);
}
}
+
+// ggml_compute_forward_dsv4_hc_comb
+
+static void ggml_dsv4_hc_comb_norm_cols(float * comb, float eps) {
+ constexpr int64_t hc = 4;
+
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ float sum = eps;
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ sum += comb[idst + hc*isrc];
+ }
+
+ const float inv_sum = 1.0f / sum;
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ comb[idst + hc*isrc] *= inv_sum;
+ }
+ }
+}
+
+static void ggml_dsv4_hc_comb_norm_rows(float * comb, float eps) {
+ constexpr int64_t hc = 4;
+
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ float sum = eps;
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ sum += comb[idst + hc*isrc];
+ }
+
+ const float inv_sum = 1.0f / sum;
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ comb[idst + hc*isrc] *= inv_sum;
+ }
+ }
+}
+
+static void ggml_compute_forward_dsv4_hc_comb_f32(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+ const ggml_tensor * mixes = dst->src[0];
+ const ggml_tensor * scale = dst->src[1];
+ const ggml_tensor * base = dst->src[2];
+
+ GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+ GGML_ASSERT(scale->type == GGML_TYPE_F32);
+ GGML_ASSERT(base->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ constexpr int64_t hc = 4;
+ constexpr int64_t comb_offset = 2*hc;
+ constexpr int64_t hc_mix_dim = (2 + hc)*hc;
+
+ const int64_t n_tokens = mixes->ne[1];
+
+ GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
+ GGML_ASSERT(dst->ne[0] == hc);
+ GGML_ASSERT(dst->ne[1] == hc);
+ GGML_ASSERT(dst->ne[2] == n_tokens);
+ GGML_ASSERT(scale->ne[0] >= 3);
+ GGML_ASSERT(base->ne[0] == hc_mix_dim);
+
+ GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
+ GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
+ GGML_TENSOR_LOCALS(size_t, nbb, base, nb);
+ GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
+
+ const float eps = ggml_get_op_params_f32(dst, 0);
+ const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
+ GGML_ASSERT(n_iter > 0);
+
+ const int ith = params->ith;
+ const int nth = params->nth;
+
+ const int64_t dr = (n_tokens + nth - 1) / nth;
+ const int64_t it0 = dr * ith;
+ const int64_t it1 = MIN(it0 + dr, n_tokens);
+
+ const float scale_comb = *(const float *) ((const char *) scale->data + 2*nbs0);
+
+ for (int64_t it = it0; it < it1; ++it) {
+ float comb[hc*hc];
+
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ float max = -INFINITY;
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ const int64_t idx = idst + hc*isrc;
+ const float xv = *(const float *) ((const char *) mixes->data + (comb_offset + idx)*nbm0 + it*nbm1);
+ const float bv = *(const float *) ((const char *) base->data + (comb_offset + idx)*nbb0);
+ const float v = xv * scale_comb + bv;
+ comb[idx] = v;
+ max = MAX(max, v);
+ }
+
+ float sum = 0.0f;
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ const int64_t idx = idst + hc*isrc;
+ const float v = expf(comb[idx] - max);
+ comb[idx] = v;
+ sum += v;
+ }
+
+ const float inv_sum = 1.0f / sum;
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ const int64_t idx = idst + hc*isrc;
+ comb[idx] = comb[idx] * inv_sum + eps;
+ }
+ }
+
+ ggml_dsv4_hc_comb_norm_cols(comb, eps);
+ for (int32_t i = 1; i < n_iter; ++i) {
+ ggml_dsv4_hc_comb_norm_rows(comb, eps);
+ ggml_dsv4_hc_comb_norm_cols(comb, eps);
+ }
+
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ for (int64_t idst = 0; idst < hc; ++idst) {
+ const int64_t idx = idst + hc*isrc;
+ *(float *) ((char *) dst->data + idst*nbd0 + isrc*nbd1 + it*nbd2) = comb[idx];
+ }
+ }
+ }
+}
+
+void ggml_compute_forward_dsv4_hc_comb(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+ const ggml_tensor * src0 = dst->src[0];
+
+ switch (src0->type) {
+ case GGML_TYPE_F32:
+ {
+ ggml_compute_forward_dsv4_hc_comb_f32(params, dst);
+ } break;
+ default:
+ {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
+// ggml_compute_forward_dsv4_hc_pre
+
+static void ggml_compute_forward_dsv4_hc_pre_f32(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+ const ggml_tensor * x = dst->src[0];
+ const ggml_tensor * weights = dst->src[1];
+
+ GGML_ASSERT(x->type == GGML_TYPE_F32);
+ GGML_ASSERT(weights->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ const int64_t n_embd = x->ne[0];
+ const int64_t hc = x->ne[1];
+ const int64_t n_tokens = x->ne[2];
+
+ GGML_ASSERT(dst->ne[0] == n_embd);
+ GGML_ASSERT(dst->ne[1] == n_tokens);
+ GGML_ASSERT(weights->ne[0] == hc);
+ GGML_ASSERT(weights->ne[1] == n_tokens);
+
+ GGML_TENSOR_LOCALS(size_t, nbx, x, nb);
+ GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
+ GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
+
+ const int ith = params->ith;
+ const int nth = params->nth;
+
+ const int64_t nr = n_embd * n_tokens;
+ const int64_t dr = (nr + nth - 1) / nth;
+ const int64_t ir0 = dr * ith;
+ const int64_t ir1 = MIN(ir0 + dr, nr);
+
+ for (int64_t ir = ir0; ir < ir1; ++ir) {
+ const int64_t i0 = ir % n_embd;
+ const int64_t it = ir / n_embd;
+
+ float sum = 0.0f;
+ for (int64_t ih = 0; ih < hc; ++ih) {
+ const float xv = *(const float *) ((const char *) x->data + i0*nbx0 + ih*nbx1 + it*nbx2);
+ const float wv = *(const float *) ((const char *) weights->data + ih*nbw0 + it*nbw1);
+ sum += xv * wv;
+ }
+
+ *(float *) ((char *) dst->data + i0*nbd0 + it*nbd1) = sum;
+ }
+}
+
+void ggml_compute_forward_dsv4_hc_pre(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+ const ggml_tensor * src0 = dst->src[0];
+
+ switch (src0->type) {
+ case GGML_TYPE_F32:
+ {
+ ggml_compute_forward_dsv4_hc_pre_f32(params, dst);
+ } break;
+ default:
+ {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
+// ggml_compute_forward_dsv4_hc_post
+
+static void ggml_compute_forward_dsv4_hc_post_f32(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+ const ggml_tensor * x = dst->src[0];
+ const ggml_tensor * residual = dst->src[1];
+ const ggml_tensor * post = dst->src[2];
+ const ggml_tensor * comb = dst->src[3];
+
+ GGML_ASSERT(x->type == GGML_TYPE_F32);
+ GGML_ASSERT(residual->type == GGML_TYPE_F32);
+ GGML_ASSERT(post->type == GGML_TYPE_F32);
+ GGML_ASSERT(comb->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ const int64_t n_embd = x->ne[0];
+ const int64_t n_tokens = x->ne[1];
+ const int64_t hc = residual->ne[1];
+
+ GGML_ASSERT(dst->ne[0] == n_embd);
+ GGML_ASSERT(dst->ne[1] == hc);
+ GGML_ASSERT(dst->ne[2] == n_tokens);
+ GGML_ASSERT(residual->ne[0] == n_embd);
+ GGML_ASSERT(residual->ne[2] == n_tokens);
+ GGML_ASSERT(post->ne[0] == hc);
+ GGML_ASSERT(post->ne[1] == n_tokens);
+ GGML_ASSERT(comb->ne[0] == hc);
+ GGML_ASSERT(comb->ne[1] == hc);
+ GGML_ASSERT(comb->ne[2] == n_tokens);
+
+ GGML_TENSOR_LOCALS(size_t, nbx, x, nb);
+ GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
+ GGML_TENSOR_LOCALS(size_t, nbp, post, nb);
+ GGML_TENSOR_LOCALS(size_t, nbc, comb, nb);
+ GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
+
+ const int ith = params->ith;
+ const int nth = params->nth;
+
+ const int64_t nr = n_embd * hc * n_tokens;
+ const int64_t dr = (nr + nth - 1) / nth;
+ const int64_t ir0 = dr * ith;
+ const int64_t ir1 = MIN(ir0 + dr, nr);
+
+ for (int64_t ir = ir0; ir < ir1; ++ir) {
+ const int64_t i0 = ir % n_embd;
+ const int64_t idst = (ir / n_embd) % hc;
+ const int64_t it = ir / (n_embd * hc);
+
+ const float xv = *(const float *) ((const char *) x->data + i0*nbx0 + it*nbx1);
+ const float pv = *(const float *) ((const char *) post->data + idst*nbp0 + it*nbp1);
+
+ float sum = xv * pv;
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ const float rv = *(const float *) ((const char *) residual->data + i0*nbr0 + isrc*nbr1 + it*nbr2);
+ const float cv = *(const float *) ((const char *) comb->data + idst*nbc0 + isrc*nbc1 + it*nbc2);
+ sum += rv * cv;
+ }
+
+ *(float *) ((char *) dst->data + i0*nbd0 + idst*nbd1 + it*nbd2) = sum;
+ }
+}
+
+void ggml_compute_forward_dsv4_hc_post(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+ const ggml_tensor * src0 = dst->src[0];
+
+ switch (src0->type) {
+ case GGML_TYPE_F32:
+ {
+ ggml_compute_forward_dsv4_hc_post_f32(params, dst);
+ } break;
+ default:
+ {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
// ggml_compute_forward_rwkv_wkv7
static void ggml_compute_forward_rwkv_wkv7_f32(
}
}
}
+
+// 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]);
+ }
+ }
+ }
+}
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_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_hc_post(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);
return out;
}
-static block_q4_0x4 make_block_q4_0x4(block_q4_0 * in, unsigned int blck_size_interleave) {
+static block_q4_0x4 make_block_q4_0x4(block_q4_0 * in, int blck_size_interleave) {
block_q4_0x4 out;
for (int i = 0; i < 4; i++) {
#endif // STRIDED_ITERATOR_AVAILABLE
#ifdef GGML_CUDA_USE_CUB
+
+// returns the suggested maximum number of rows to process during one argsort_f32_i32_cuda_cub() call
+int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows) {
+ // perform argsort in chunks up to approximately this size (currently 64MB)
+ // to avoid excessive temporary buffers memory usage
+ const int chunk_bytes = 1 << 26;
+
+ // calculate how many rows will fit in one chunk (must be at least one)
+ const int chunk_nrows = std::max((int) (chunk_bytes / nb01), 1);
+
+ // limit the resulting amount to total nrows
+ return std::min((int64_t) chunk_nrows, nrows);
+}
+
void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
const float * x,
int * dst,
const size_t shared_mem = ncols_pad * sizeof(int);
const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;
- if (shared_mem > max_shared_mem || ncols > 1024) {
- ggml_cuda_pool & pool = ctx.pool();
- argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream);
- } else {
+ // early return if we can use bitonic argsort
+ if (shared_mem <= max_shared_mem && ncols <= 1024) {
argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
+ return;
+ }
+
+ const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows);
+
+ ggml_cuda_pool & pool = ctx.pool();
+
+ for (int64_t i = 0; i < nrows; i += chunk_nrows) {
+ int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i);
+
+ argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, iter_nrows, order, stream);
+
+ src0_d += ncols * iter_nrows;
+ dst_d += ncols * iter_nrows;
}
#else
argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
#ifdef GGML_CUDA_USE_CUB
+int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows);
void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
const float * x,
int * dst,
ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN;
}
+// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) {
+ GGML_ASSERT(tensor != nullptr);
+ return (reinterpret_cast<uintptr_t>(tensor->data) % alignment) == 0 &&
+ tensor->nb[1] % alignment == 0 &&
+ tensor->nb[2] % alignment == 0 &&
+ tensor->nb[3] % alignment == 0;
+}
+
static constexpr __device__ int ggml_cuda_get_physical_warp_size() {
#if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
return 64;
typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v);
+template<typename dst_t>
+using dequantize_kq_t = void (*)(const void * vx, const int64_t ib, dst_t * y, const int tid);
+
static __device__ __forceinline__ float get_alibi_slope(
const float max_bias, const uint32_t h, const uint32_t n_head_log2, const float m0, const float m1
) {
static constexpr int qi = QI1_0;
};
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_Q2_0> {
+ static constexpr int qk = QK2_0;
+ static constexpr int qr = QR2_0;
+ static constexpr int qi = QI2_0;
+};
+
template<>
struct ggml_cuda_type_traits<GGML_TYPE_Q4_0> {
static constexpr int qk = QK4_0;
//////////////////////
struct ggml_cuda_device_info {
- int device_count;
+ int device_count; // number of (possibly virtual) devices exposed to the rest of ggml
+ int physical_device_count; // number of physical CUDA devices actually present
struct cuda_device_info {
int cc; // compute capability
size_t total_vram;
int warp_size; // Number of threads in a dispatch
bool supports_cooperative_launch; // whether cooperative launch is supported
+ int physical_device; // backing physical CUDA device for this (virtual) device
+ int physical_share_count; // number of (virtual) devices sharing this device's physical GPU
+ int virtual_index; // index of this (virtual) device among those sharing its physical GPU
};
cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {};
template <typename T>
static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, int dim, cudaStream_t stream) {
- if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
+ if (dim != 3 && ggml_is_contiguous_to_3(src0) && ggml_is_contiguous_to_3(src1)) {
const T * src0_d = (const T *) src0->data;
const T * src1_d = (const T *) src1->data;
T * dst_d = (T *) dst->data;
- if (dim != 3) {
- for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) {
- concat_cont_cuda(
- src0_d + i3*(src0->nb[3] / sizeof(T)),
- src1_d + i3*(src1->nb[3] / sizeof(T)),
- dst_d + i3*( dst->nb[3] / sizeof(T)),
- ggml_row_size(src0->type, src0->ne[0])/sizeof(T), src0->ne[1], src0->ne[2],
- ggml_row_size(dst->type, dst->ne[0])/sizeof(T), dst->ne[1], dst->ne[2], dim, stream);
- }
- } else {
- const size_t size0 = ggml_nbytes(src0);
- const size_t size1 = ggml_nbytes(src1);
-
- CUDA_CHECK(cudaMemcpyAsync((char *) dst->data, src0->data, size0, cudaMemcpyDeviceToDevice, stream));
- CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream));
+ for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) {
+ concat_cont_cuda(
+ src0_d + i3*(src0->nb[3] / sizeof(T)),
+ src1_d + i3*(src1->nb[3] / sizeof(T)),
+ dst_d + i3*( dst->nb[3] / sizeof(T)),
+ ggml_row_size(src0->type, src0->ne[0])/sizeof(T), src0->ne[1], src0->ne[2],
+ ggml_row_size(dst->type, dst->ne[0])/sizeof(T), dst->ne[1], dst->ne[2], dim, stream);
}
+ } else if (dim == 3 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
+ const size_t size0 = ggml_nbytes(src0);
+ const size_t size1 = ggml_nbytes(src1);
+
+ CUDA_CHECK(cudaMemcpyAsync((char *) dst->data, src0->data, size0, cudaMemcpyDeviceToDevice, stream));
+ CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream));
} else {
GGML_ASSERT(!ggml_is_quantized(src0->type));
GGML_ASSERT(dst->type == src0->type);
if (ggml_is_quantized(src0->type)) {
- GGML_ASSERT(ggml_is_contiguous(src0));
- GGML_ASSERT(ggml_is_contiguous(src1));
+ if (dim == 3) {
+ GGML_ASSERT(ggml_is_contiguous(src0));
+ GGML_ASSERT(ggml_is_contiguous(src1));
+ } else {
+ GGML_ASSERT(ggml_is_contiguous_to_3(src0));
+ GGML_ASSERT(ggml_is_contiguous_to_3(src1));
+ }
GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
- // if tensors are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors
+ // if first 3 dimensions are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors
concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
} else {
GGML_ASSERT(ggml_blck_size(src0->type) == 1);
template<typename dst_t>
static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_q2_K * x = (const block_q2_K *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t n = tid/32;
- const int64_t l = tid - 32*n;
- const int64_t is = 8*n + l/16;
-
- const uint8_t q = x[i].qs[32*n + l];
- dst_t * y = yy + i*QK_K + 128*n;
-
- float dall = __low2half(x[i].dm);
- float dmin = __high2half(x[i].dm);
- y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4));
- y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4));
- y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4));
- y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4));
+ dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
const int64_t i = blockIdx.x;
- const block_q3_K * x = (const block_q3_K *) vx;
-
- const int64_t r = threadIdx.x/4;
- const int64_t tid = r/2;
- const int64_t is0 = r%2;
- const int64_t l0 = 16*is0 + 4*(threadIdx.x%4);
- const int64_t n = tid / 4;
- const int64_t j = tid - 4*n;
-
- uint8_t m = 1 << (4*n + j);
- int64_t is = 8*n + 2*j + is0;
- int shift = 2*j;
-
- int8_t us = is < 4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) :
- is < 8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) :
- is < 12 ? (x[i].scales[is-8] >> 4) | (((x[i].scales[is+0] >> 4) & 3) << 4) :
- (x[i].scales[is-8] >> 4) | (((x[i].scales[is-4] >> 6) & 3) << 4);
- float d_all = x[i].d;
- float dl = d_all * (us - 32);
-
- dst_t * y = yy + i*QK_K + 128*n + 32*j;
- const uint8_t * q = x[i].qs + 32*n;
- const uint8_t * hm = x[i].hmask;
-
- for (int l = l0; l < l0+4; ++l) {
- y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
- }
-}
-static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
- if (j < 4) {
- d = q[j] & 63; m = q[j + 4] & 63;
- } else {
- d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
- m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4);
- }
+ dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
- const block_q4_K * x = (const block_q4_K *) vx;
-
const int64_t i = blockIdx.x;
- // assume 32 threads
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8;
- const int64_t ir = tid%8;
- const int64_t is = 2*il;
- const int64_t n = 4;
-
- dst_t * y = yy + i*QK_K + 64*il + n*ir;
-
- const float dall = __low2half(x[i].dm);
- const float dmin = __high2half(x[i].dm);
-
- const uint8_t * q = x[i].qs + 32*il + n*ir;
-
- uint8_t sc, m;
- get_scale_min_k4(is + 0, x[i].scales, sc, m);
- const float d1 = dall * sc; const float m1 = dmin * m;
- get_scale_min_k4(is + 1, x[i].scales, sc, m);
- const float d2 = dall * sc; const float m2 = dmin * m;
- for (int l = 0; l < n; ++l) {
- y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
- y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >> 4) - m2);
- }
+ dequantize_q4_K(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
- const block_q5_K * x = (const block_q5_K *) vx;
-
const int64_t i = blockIdx.x;
- // assume 64 threads - this is very slightly better than the one below
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/16; // il is in 0...3
- const int64_t ir = tid%16; // ir is in 0...15
- const int64_t is = 2*il; // is is in 0...6
-
- dst_t * y = yy + i*QK_K + 64*il + 2*ir;
-
- const float dall = __low2half(x[i].dm);
- const float dmin = __high2half(x[i].dm);
-
- const uint8_t * ql = x[i].qs + 32*il + 2*ir;
- const uint8_t * qh = x[i].qh + 2*ir;
-
- uint8_t sc, m;
- get_scale_min_k4(is + 0, x[i].scales, sc, m);
- const float d1 = dall * sc; const float m1 = dmin * m;
- get_scale_min_k4(is + 1, x[i].scales, sc, m);
- const float d2 = dall * sc; const float m2 = dmin * m;
-
- uint8_t hm = 1 << (2*il);
- y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
- y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
- hm <<= 1;
- y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2);
- y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2);
+ dequantize_q5_K(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
- const block_q6_K * x = (const block_q6_K *) vx;
-
const int64_t i = blockIdx.x;
- // assume 64 threads - this is very slightly better than the one below
- const int64_t tid = threadIdx.x;
- const int64_t ip = tid/32; // ip is 0 or 1
- const int64_t il = tid - 32*ip; // 0...32
- const int64_t is = 8*ip + il/16;
-
- dst_t * y = yy + i*QK_K + 128*ip + il;
-
- const float d = x[i].d;
-
- const uint8_t * ql = x[i].ql + 64*ip + il;
- const uint8_t qh = x[i].qh[32*ip + il];
- const int8_t * sc = x[i].scales + is;
-
- y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
- y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
- y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32));
- y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32));
+ dequantize_q6_K(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq2_xxs * x = (const block_iq2_xxs *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const uint16_t * q2 = x[i].qs + 4*ib;
- const uint8_t * aux8 = (const uint8_t *)q2;
- const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
- const uint32_t aux32 = q2[2] | (q2[3] << 16);
- const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.25f;
- const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
- for (int j = 0; j < 8; ++j) {
- y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
- }
+ dequantize_iq2_xxs(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq2_xs * x = (const block_iq2_xs *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const uint16_t * q2 = x[i].qs + 4*ib;
- const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
- const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
- const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
- for (int j = 0; j < 8; ++j) {
- y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
- }
+ dequantize_iq2_xs(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq2_s * x = (const block_iq2_s *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
- const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
- const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
- for (int j = 0; j < 8; ++j) {
- y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
- }
+ dequantize_iq2_s(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq3_xxs * x = (const block_iq3_xxs *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const uint8_t * q3 = x[i].qs + 8*ib;
- const uint16_t * gas = (const uint16_t *)(x[i].qs + QK_K/4) + 2*ib;
- const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
- const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
- const uint32_t aux32 = gas[0] | (gas[1] << 16);
- const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.5f;
- const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
- for (int j = 0; j < 4; ++j) {
- y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
- y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
- }
+ dequantize_iq3_xxs(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq3_s * x = (const block_iq3_s *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const uint8_t * qs = x[i].qs + 8*ib;
- const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[i].qh[ib] << (8-2*il)) & 256)));
- const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[i].qh[ib] << (7-2*il)) & 256)));
- const float d = (float)x[i].d * (1 + 2*((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf));
- const uint8_t signs = x[i].signs[4*ib + il];
- for (int j = 0; j < 4; ++j) {
- y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
- y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
- }
+ dequantize_iq3_s(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq1_s * x = (const block_iq1_s *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
- const float d = (float)x[i].d * (2*((x[i].qh[ib] >> 12) & 7) + 1);
- uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
- grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[ib] >> 3*il) & 7) << 8)];
- grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
- grid32[0] &= 0x0f0f0f0f;
- for (int j = 0; j < 8; ++j) {
- y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
- }
+ dequantize_iq1_s(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq1_m * x = (const block_iq1_m *) vx;
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 8*il;
- const uint16_t * sc = (const uint16_t *)x[i].scales;
- iq1m_scale_t scale;
- scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
- const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
- const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
- const float delta = x[i].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
- uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
- grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
- grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
- grid32[0] &= 0x0f0f0f0f;
- for (int j = 0; j < 8; ++j) {
- y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
- }
+ dequantize_iq1_m(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_iq4_nl * x = (const block_iq4_nl *) vx + i*(QK_K/QK4_NL);
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 4*il;
- const uint8_t * q4 = x[ib].qs + 4*il;
- const float d = (float)x[ib].d;
- for (int j = 0; j < 4; ++j) {
- y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
- y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >> 4]);
- }
+ dequantize_iq4_nl(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
- const int64_t i = blockIdx.x;
- const block_iq4_xs * x = (const block_iq4_xs *)vx;
+ const int64_t i = blockIdx.x;
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 4*il;
- const uint8_t * q4 = x[i].qs + 16*ib + 4*il;
- const float d = (float)x[i].d * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
- for (int j = 0; j < 4; ++j) {
- y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
- y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >> 4]);
- }
+ dequantize_iq4_xs(vx, i, yy + i*QK_K, threadIdx.x);
}
template<typename dst_t>
static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+ const int64_t i = blockIdx.x;
- const int64_t i = blockIdx.x;
- const block_mxfp4 * x = (const block_mxfp4 *) vx + i*(QK_K/QK_MXFP4);
-
- const int64_t tid = threadIdx.x;
- const int64_t il = tid/8; // 0...3
- const int64_t ib = tid%8; // 0...7
- dst_t * y = yy + i*QK_K + 32*ib + 4*il;
- const uint8_t * q4 = x[ib].qs + 4*il;
- const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
- for (int j = 0; j < 4; ++j) {
- y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
- y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >> 4]*0.5f);
- }
+ dequantize_mxfp4(vx, i, yy + i*QK_K, threadIdx.x);
}
template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
switch (type) {
case GGML_TYPE_Q1_0:
return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+ case GGML_TYPE_Q2_0:
+ return dequantize_block_cont_cuda<QK2_0, QR2_0, dequantize_q2_0>;
case GGML_TYPE_Q4_0:
return dequantize_row_q4_0_cuda;
case GGML_TYPE_Q4_1:
switch (type) {
case GGML_TYPE_Q1_0:
return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+ case GGML_TYPE_Q2_0:
+ return dequantize_block_cont_cuda<QK2_0, QR2_0, dequantize_q2_0>;
case GGML_TYPE_Q4_0:
return dequantize_row_q4_0_cuda;
case GGML_TYPE_Q4_1:
switch (type) {
case GGML_TYPE_Q1_0:
return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+ case GGML_TYPE_Q2_0:
+ return dequantize_block_cont_cuda<QK2_0, QR2_0, dequantize_q2_0>;
case GGML_TYPE_Q4_0:
return dequantize_row_q4_0_cuda;
case GGML_TYPE_Q4_1:
return convert_unary_cuda<float>;
case GGML_TYPE_Q1_0:
return dequantize_block_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+ case GGML_TYPE_Q2_0:
+ return dequantize_block_cuda<QK2_0, QR2_0, dequantize_q2_0>;
case GGML_TYPE_Q4_0:
return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
case GGML_TYPE_Q4_1:
return convert_unary_cuda<float, nv_bfloat16>;
case GGML_TYPE_Q1_0:
return dequantize_block_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+ case GGML_TYPE_Q2_0:
+ return dequantize_block_cuda<QK2_0, QR2_0, dequantize_q2_0>;
case GGML_TYPE_Q4_0:
return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
case GGML_TYPE_Q4_1:
return convert_unary_cuda<half, float>;
case GGML_TYPE_Q1_0:
return dequantize_block_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+ case GGML_TYPE_Q2_0:
+ return dequantize_block_cuda<QK2_0, QR2_0, dequantize_q2_0>;
case GGML_TYPE_Q4_0:
return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
case GGML_TYPE_Q4_1:
#include "common.cuh"
+#include "convert.cuh"
static __device__ __forceinline__ void dequantize_q1_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
const block_q1_0 * x = (const block_q1_0 *) vx;
v.y = (2*bit_1 - 1) * d;
}
+static __device__ __forceinline__ void dequantize_q2_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
+ const block_q2_0 * x = (const block_q2_0 *) vx;
+
+ const float d = x[ib].d;
+
+ // Q2_0: 2 bits per element, 4 elements per byte.
+ // Stored code c in {0,1,2,3} maps to symbol s = c - 1 in {-1, 0, +1, +2}.
+ const int byte_index_0 = iqs / 4;
+ const int bit_offset_0 = (iqs % 4) * 2;
+
+ const int byte_index_1 = (iqs + 1) / 4;
+ const int bit_offset_1 = ((iqs + 1) % 4) * 2;
+
+ const int c0 = (x[ib].qs[byte_index_0] >> bit_offset_0) & 0x3;
+ const int c1 = (x[ib].qs[byte_index_1] >> bit_offset_1) & 0x3;
+
+ v.x = (c0 - 1) * d;
+ v.y = (c1 - 1) * d;
+}
+
static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
const block_q4_0 * x = (const block_q4_0 *) vx;
v.x *= d;
v.y *= d;
}
+
+//================================== k-quants
+
+// Each call dequantizes one super-block of QK_K values into y using the
+// thread layout of the caller: 32 threads for q4_K, 64 threads otherwise.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+ const block_q2_K * x = (const block_q2_K *) vx;
+
+ const int64_t n = tid/32;
+ const int64_t l = tid - 32*n;
+ const int64_t is = 8*n + l/16;
+
+ const uint8_t q = x[ib].qs[32*n + l];
+ dst_t * y = yy + 128*n;
+
+ float dall = __low2half(x[ib].dm);
+ float dmin = __high2half(x[ib].dm);
+ y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[ib].scales[is+0] >> 4));
+ y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[ib].scales[is+2] >> 4));
+ y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[ib].scales[is+4] >> 4));
+ y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[ib].scales[is+6] >> 4));
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+ const block_q3_K * x = (const block_q3_K *) vx;
+
+ const int64_t r = tid/4;
+ const int64_t t = r/2;
+ const int64_t is0 = r%2;
+ const int64_t l0 = 16*is0 + 4*(tid%4);
+ const int64_t n = t / 4;
+ const int64_t j = t - 4*n;
+
+ uint8_t m = 1 << (4*n + j);
+ int64_t is = 8*n + 2*j + is0;
+ int shift = 2*j;
+
+ int8_t us = is < 4 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+8] >> 0) & 3) << 4) :
+ is < 8 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+4] >> 2) & 3) << 4) :
+ is < 12 ? (x[ib].scales[is-8] >> 4) | (((x[ib].scales[is+0] >> 4) & 3) << 4) :
+ (x[ib].scales[is-8] >> 4) | (((x[ib].scales[is-4] >> 6) & 3) << 4);
+ float d_all = x[ib].d;
+ float dl = d_all * (us - 32);
+
+ dst_t * y = yy + 128*n + 32*j;
+ const uint8_t * q = x[ib].qs + 32*n;
+ const uint8_t * hm = x[ib].hmask;
+
+ for (int l = l0; l < l0+4; ++l) {
+ y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
+ }
+}
+
+static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
+ if (j < 4) {
+ d = q[j] & 63; m = q[j + 4] & 63;
+ } else {
+ d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
+ m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4);
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+ const block_q4_K * x = (const block_q4_K *) vx;
+
+ // assume 32 threads
+ const int64_t il = tid/8;
+ const int64_t ir = tid%8;
+ const int64_t is = 2*il;
+ const int64_t n = 4;
+
+ dst_t * y = yy + 64*il + n*ir;
+
+ const float dall = __low2half(x[ib].dm);
+ const float dmin = __high2half(x[ib].dm);
+
+ const uint8_t * q = x[ib].qs + 32*il + n*ir;
+
+ uint8_t sc, m;
+ get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+ const float d1 = dall * sc; const float m1 = dmin * m;
+ get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+ const float d2 = dall * sc; const float m2 = dmin * m;
+ for (int l = 0; l < n; ++l) {
+ y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
+ y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >> 4) - m2);
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+ const block_q5_K * x = (const block_q5_K *) vx;
+
+ // assume 64 threads - this is very slightly better than the one below
+ const int64_t il = tid/16; // il is in 0...3
+ const int64_t ir = tid%16; // ir is in 0...15
+ const int64_t is = 2*il; // is is in 0...6
+
+ dst_t * y = yy + 64*il + 2*ir;
+
+ const float dall = __low2half(x[ib].dm);
+ const float dmin = __high2half(x[ib].dm);
+
+ const uint8_t * ql = x[ib].qs + 32*il + 2*ir;
+ const uint8_t * qh = x[ib].qh + 2*ir;
+
+ uint8_t sc, m;
+ get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+ const float d1 = dall * sc; const float m1 = dmin * m;
+ get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+ const float d2 = dall * sc; const float m2 = dmin * m;
+
+ uint8_t hm = 1 << (2*il);
+ y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
+ y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
+ hm <<= 1;
+ y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2);
+ y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2);
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+ const block_q6_K * x = (const block_q6_K *) vx;
+
+ // assume 64 threads - this is very slightly better than the one below
+ const int64_t ip = tid/32; // ip is 0 or 1
+ const int64_t il = tid - 32*ip; // 0...32
+ const int64_t is = 8*ip + il/16;
+
+ dst_t * y = yy + 128*ip + il;
+
+ const float d = x[ib].d;
+
+ const uint8_t * ql = x[ib].ql + 64*ip + il;
+ const uint8_t qh = x[ib].qh[32*ip + il];
+ const int8_t * sc = x[ib].scales + is;
+
+ y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
+ y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
+ y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32));
+ y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32));
+}
+
+//================================== i-quants
+
+// Each call dequantizes one super-block of QK_K values into y with 32
+// threads; iq4_nl packs QK_K/QK4_NL sub-blocks per super-block.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq2_xxs * x = (const block_iq2_xxs *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const uint16_t * q2 = x[ibs].qs + 4*ib;
+ const uint8_t * aux8 = (const uint8_t *)q2;
+ const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
+ const uint32_t aux32 = q2[2] | (q2[3] << 16);
+ const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.25f;
+ const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+ for (int j = 0; j < 8; ++j) {
+ y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq2_xs * x = (const block_iq2_xs *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const uint16_t * q2 = x[ibs].qs + 4*ib;
+ const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
+ const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+ const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
+ for (int j = 0; j < 8; ++j) {
+ y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq2_s * x = (const block_iq2_s *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[ibs].qs[4*ib+il] | ((x[ibs].qh[ib] << (8-2*il)) & 0x300)));
+ const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+ const uint8_t signs = x[ibs].qs[QK_K/8+4*ib+il];
+ for (int j = 0; j < 8; ++j) {
+ y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq3_xxs * x = (const block_iq3_xxs *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const uint8_t * q3 = x[ibs].qs + 8*ib;
+ const uint16_t * gas = (const uint16_t *)(x[ibs].qs + QK_K/4) + 2*ib;
+ const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
+ const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
+ const uint32_t aux32 = gas[0] | (gas[1] << 16);
+ const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.5f;
+ const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+ for (int j = 0; j < 4; ++j) {
+ y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+ y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq3_s * x = (const block_iq3_s *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const uint8_t * qs = x[ibs].qs + 8*ib;
+ const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[ibs].qh[ib] << (8-2*il)) & 256)));
+ const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[ibs].qh[ib] << (7-2*il)) & 256)));
+ const float d = (float)x[ibs].d * (1 + 2*((x[ibs].scales[ib/2] >> 4*(ib%2)) & 0xf));
+ const uint8_t signs = x[ibs].signs[4*ib + il];
+ for (int j = 0; j < 4; ++j) {
+ y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+ y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq1_s * x = (const block_iq1_s *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
+ const float d = (float)x[ibs].d * (2*((x[ibs].qh[ib] >> 12) & 7) + 1);
+ uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+ grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[ib] >> 3*il) & 7) << 8)];
+ grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+ grid32[0] &= 0x0f0f0f0f;
+ for (int j = 0; j < 8; ++j) {
+ y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq1_m * x = (const block_iq1_m *) vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 8*il;
+ const uint16_t * sc = (const uint16_t *)x[ibs].scales;
+ iq1m_scale_t scale;
+ scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+ const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
+ const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
+ const float delta = x[ibs].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
+ uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+ grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
+ grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+ grid32[0] &= 0x0f0f0f0f;
+ for (int j = 0; j < 8; ++j) {
+ y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_iq4_nl * x = (const block_iq4_nl *) vx + ibs*(QK_K/QK4_NL);
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 4*il;
+ const uint8_t * q4 = x[ib].qs + 4*il;
+ const float d = (float)x[ib].d;
+ for (int j = 0; j < 4; ++j) {
+ y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+ y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >> 4]);
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+ const block_iq4_xs * x = (const block_iq4_xs *)vx;
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 4*il;
+ const uint8_t * q4 = x[ibs].qs + 16*ib + 4*il;
+ const float d = (float)x[ibs].d * ((((x[ibs].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[ibs].scales_h >> 2*ib) & 3) << 4)) - 32);
+ for (int j = 0; j < 4; ++j) {
+ y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+ y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >> 4]);
+ }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+ const block_mxfp4 * x = (const block_mxfp4 *) vx + ibs*(QK_K/QK_MXFP4);
+
+ const int64_t il = tid/8; // 0...3
+ const int64_t ib = tid%8; // 0...7
+ dst_t * y = yy + 32*ib + 4*il;
+ const uint8_t * q4 = x[ib].qs + 4*il;
+ const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
+ for (int j = 0; j < 4; ++j) {
+ y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
+ y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >> 4]*0.5f);
+ }
+}
--- /dev/null
+#include "common.cuh"
+#include "dsv4-hc.cuh"
+
+
+static constexpr int DSV4_HC = 4;
+
+
+static __device__ void dsv4_hc_comb_norm_cols(float * comb, float eps) {
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ float sum = eps;
+ for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+ sum += comb[idst + DSV4_HC*isrc];
+ }
+
+ const float inv_sum = 1.0f / sum;
+ for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+ comb[idst + DSV4_HC*isrc] *= inv_sum;
+ }
+ }
+}
+
+static __device__ void dsv4_hc_comb_norm_rows(float * comb, float eps) {
+ for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+ float sum = eps;
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ sum += comb[idst + DSV4_HC*isrc];
+ }
+
+ const float inv_sum = 1.0f / sum;
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ comb[idst + DSV4_HC*isrc] *= inv_sum;
+ }
+ }
+}
+
+static __global__ void dsv4_hc_comb_f32(
+ const float * mixes,
+ const float * scale,
+ const float * base,
+ float * dst,
+ int64_t n_tokens,
+ int64_t sm0,
+ int64_t sm1,
+ int64_t ss0,
+ int64_t sb0,
+ int64_t sd0,
+ int64_t sd1,
+ int64_t sd2,
+ float eps,
+ int32_t n_iter) {
+ constexpr int comb_offset = 2*DSV4_HC;
+
+ ggml_cuda_pdl_lc();
+ const int64_t it = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+
+ if (it >= n_tokens) {
+ return;
+ }
+
+ ggml_cuda_pdl_sync();
+
+ const float scale_comb = scale[2*ss0];
+ float comb[DSV4_HC*DSV4_HC];
+
+ for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+ float max = -INFINITY;
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ const int idx = idst + DSV4_HC*isrc;
+ const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0];
+ comb[idx] = v;
+ max = fmaxf(max, v);
+ }
+
+ float sum = 0.0f;
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ const int idx = idst + DSV4_HC*isrc;
+ const float v = expf(comb[idx] - max);
+ comb[idx] = v;
+ sum += v;
+ }
+
+ const float inv_sum = 1.0f / sum;
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ const int idx = idst + DSV4_HC*isrc;
+ comb[idx] = comb[idx] * inv_sum + eps;
+ }
+ }
+
+ dsv4_hc_comb_norm_cols(comb, eps);
+ for (int32_t i = 1; i < n_iter; ++i) {
+ dsv4_hc_comb_norm_rows(comb, eps);
+ dsv4_hc_comb_norm_cols(comb, eps);
+ }
+
+ for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+ for (int idst = 0; idst < DSV4_HC; ++idst) {
+ const int idx = idst + DSV4_HC*isrc;
+ dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx];
+ }
+ }
+}
+
+static __global__ void dsv4_hc_pre_f32(
+ const float * x,
+ const float * weights,
+ float * dst,
+ int64_t n_embd,
+ int64_t hc,
+ int64_t n_tokens,
+ int64_t sx0,
+ int64_t sx1,
+ int64_t sx2,
+ int64_t sw0,
+ int64_t sw1,
+ int64_t sd0,
+ int64_t sd1) {
+ ggml_cuda_pdl_lc();
+ const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+ const int64_t nr = n_embd * n_tokens;
+
+ if (ir >= nr) {
+ return;
+ }
+
+ ggml_cuda_pdl_sync();
+
+ const int64_t i0 = ir % n_embd;
+ const int64_t it = ir / n_embd;
+
+ float sum = x[i0*sx0 + it*sx2] * weights[it*sw1];
+ for (int64_t ih = 1; ih < hc; ++ih) {
+ const float xv = x[i0*sx0 + ih*sx1 + it*sx2];
+ const float wv = weights[ih*sw0 + it*sw1];
+ sum += xv * wv;
+ }
+
+ dst[i0*sd0 + it*sd1] = sum;
+}
+
+static __global__ void dsv4_hc_post_f32(
+ const float * x,
+ const float * residual,
+ const float * post,
+ const float * comb,
+ float * dst,
+ int64_t n_embd,
+ int64_t hc,
+ int64_t n_tokens,
+ int64_t sx0,
+ int64_t sx1,
+ int64_t sr0,
+ int64_t sr1,
+ int64_t sr2,
+ int64_t sp0,
+ int64_t sp1,
+ int64_t sc0,
+ int64_t sc1,
+ int64_t sc2,
+ int64_t sd0,
+ int64_t sd1,
+ int64_t sd2) {
+ ggml_cuda_pdl_lc();
+ const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+ const int64_t nr = n_embd * hc * n_tokens;
+
+ if (ir >= nr) {
+ return;
+ }
+
+ ggml_cuda_pdl_sync();
+
+ const int64_t i0 = ir % n_embd;
+ const int64_t idst = (ir / n_embd) % hc;
+ const int64_t it = ir / (n_embd * hc);
+
+ float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1];
+ for (int64_t isrc = 0; isrc < hc; ++isrc) {
+ sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2];
+ }
+
+ dst[i0*sd0 + idst*sd1 + it*sd2] = sum;
+}
+
+void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+ const ggml_tensor * mixes = dst->src[0];
+ const ggml_tensor * scale = dst->src[1];
+ const ggml_tensor * base = dst->src[2];
+
+ GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+ GGML_ASSERT(scale->type == GGML_TYPE_F32);
+ GGML_ASSERT(base->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ constexpr int64_t hc_mix_dim = (2 + DSV4_HC)*DSV4_HC;
+
+ GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
+ GGML_ASSERT(dst->ne[0] == DSV4_HC);
+ GGML_ASSERT(dst->ne[1] == DSV4_HC);
+ GGML_ASSERT(dst->ne[2] == mixes->ne[1]);
+ GGML_ASSERT(scale->ne[0] >= 3);
+ GGML_ASSERT(base->ne[0] == hc_mix_dim);
+
+ GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
+ GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
+ GGML_TENSOR_LOCALS(size_t, nbb, base, nb);
+ GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
+
+ const int64_t n_tokens = mixes->ne[1];
+ const float eps = ggml_get_op_params_f32(dst, 0);
+ const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
+
+ const int block_size = 256;
+ const dim3 block_dims(block_size, 1, 1);
+ const dim3 grid_dims((n_tokens + block_size - 1) / block_size, 1, 1);
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+ ggml_cuda_kernel_launch(dsv4_hc_comb_f32, launch_params,
+ (const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data,
+ n_tokens,
+ nbm0 / sizeof(float), nbm1 / sizeof(float),
+ nbs0 / sizeof(float),
+ nbb0 / sizeof(float),
+ nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float),
+ eps, n_iter);
+}
+
+void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+ const ggml_tensor * x = dst->src[0];
+ const ggml_tensor * weights = dst->src[1];
+
+ GGML_ASSERT(x->type == GGML_TYPE_F32);
+ GGML_ASSERT(weights->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ GGML_TENSOR_LOCALS(size_t, nbx, x, nb);
+ GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
+ GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
+
+ const int64_t n_embd = x->ne[0];
+ const int64_t hc = x->ne[1];
+ const int64_t n_tokens = x->ne[2];
+
+ const int block_size = 256;
+ const int64_t nr = n_embd * n_tokens;
+ const dim3 block_dims(block_size, 1, 1);
+ const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+ ggml_cuda_kernel_launch(dsv4_hc_pre_f32, launch_params,
+ (const float *) x->data, (const float *) weights->data, (float *) dst->data,
+ n_embd, hc, n_tokens,
+ nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float),
+ nbw0 / sizeof(float), nbw1 / sizeof(float),
+ nbd0 / sizeof(float), nbd1 / sizeof(float));
+}
+
+void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+ const ggml_tensor * x = dst->src[0];
+ const ggml_tensor * residual = dst->src[1];
+ const ggml_tensor * post = dst->src[2];
+ const ggml_tensor * comb = dst->src[3];
+
+ GGML_ASSERT(x->type == GGML_TYPE_F32);
+ GGML_ASSERT(residual->type == GGML_TYPE_F32);
+ GGML_ASSERT(post->type == GGML_TYPE_F32);
+ GGML_ASSERT(comb->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ GGML_TENSOR_LOCALS(size_t, nbx, x, nb);
+ GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
+ GGML_TENSOR_LOCALS(size_t, nbp, post, nb);
+ GGML_TENSOR_LOCALS(size_t, nbc, comb, nb);
+ GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
+
+ const int64_t n_embd = x->ne[0];
+ const int64_t n_tokens = x->ne[1];
+ const int64_t hc = residual->ne[1];
+
+ const int block_size = 256;
+ const int64_t nr = n_embd * hc * n_tokens;
+ const dim3 block_dims(block_size, 1, 1);
+ const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+ ggml_cuda_kernel_launch(dsv4_hc_post_f32, launch_params,
+ (const float *) x->data, (const float *) residual->data,
+ (const float *) post->data, (const float *) comb->data, (float *) dst->data,
+ n_embd, hc, n_tokens,
+ nbx0 / sizeof(float), nbx1 / sizeof(float),
+ nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float),
+ nbp0 / sizeof(float), nbp1 / sizeof(float),
+ nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float),
+ nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float));
+}
--- /dev/null
+#include "common.cuh"
+#include "ggml.h"
+
+void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
#include "common.cuh"
#include "fattn-tile.cuh"
-#include "fattn-wmma-f16.cuh"
void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * K = dst->src[1];
#include "common.cuh"
#include "fattn-common.cuh"
-#include "fattn-wmma-f16.cuh"
// nbatch_fa == number of KQ rows to process per iteration
// nbatch_K == number of K columns to load in parallel for KQ calculation
// Skip unused kernel variants for faster compilation:
- if (
-#ifdef GGML_USE_WMMA_FATTN
- (ncols2 != 1 && DV != 40 && DV != 72 && DV != 512) ||
-#endif // GGML_USE_WMMA_FATTN
- (use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512))
- ) {
+ if ((use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512))) {
GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
+++ /dev/null
-// Old and deprecated WMMA FlashAttention implementation.
-// It is still needed for Volta since the memory layout of NVIDIA tensor cores changed with Turing.
-// Long-term the WMMA code should be replaced with a dedicated Volta implementation.
-
-#include "common.cuh"
-#include "fattn-common.cuh"
-#include "fattn-wmma-f16.cuh"
-
-#ifdef GGML_USE_WMMA_FATTN
-#if !defined(GGML_USE_HIP)
-#include <mma.h>
-#if defined(GGML_USE_MUSA)
-namespace wmma = mtmusa::wmma;
-#else // GGML_USE_MUSA
-namespace wmma = nvcuda::wmma;
-#endif // GGML_USE_MUSA
-#elif defined(GGML_USE_HIP)
-#include <rocwmma/rocwmma.hpp>
-namespace wmma = rocwmma;
-#endif // !defined(GGML_USE_HIP)
-#endif // GGML_USE_WMMA_FATTN
-
-// D == head size, VKQ_stride == num VKQ rows calculated in parallel:
-template<int D, int ncols, int nwarps, int VKQ_stride, typename KQ_acc_t, bool use_logit_softcap>
-__launch_bounds__(nwarps*ggml_cuda_get_physical_warp_size(), 1)
-static __global__ void flash_attn_ext_f16(
- const char * Q_ptr,
- const char * K_ptr,
- const char * V_ptr,
- const char * mask_ptr,
- const char * sinks_ptr,
- const int * KV_max_ptr,
- float * dst_ptr,
- float2 * dst_meta_ptr,
- const float scale,
- const float max_bias,
- const float m0,
- const float m1,
- const uint32_t n_head_log2,
- const float logit_softcap,
- const int32_t ne00, const uint3 ne01, const int32_t ne02, const int32_t ne03,
- const int32_t nb01, const int32_t nb02, const int32_t nb03,
- const int32_t ne10, const int32_t ne11, const int32_t ne12, const int32_t ne13,
- const int32_t nb11, const int32_t nb12, const int64_t nb13,
- const int32_t nb21, const int32_t nb22, const int64_t nb23,
- const int32_t ne31, const int32_t ne32, const int32_t ne33,
- const int32_t nb31, const int32_t nb32, const int64_t nb33) {
-#if defined(FLASH_ATTN_AVAILABLE) && (defined(GGML_HIP_ROCWMMA_FATTN) && defined(GGML_USE_WMMA_FATTN))
- const char * GGML_CUDA_RESTRICT Q = Q_ptr;
- const char * GGML_CUDA_RESTRICT K = K_ptr;
- const char * GGML_CUDA_RESTRICT V = V_ptr;
- const char * GGML_CUDA_RESTRICT mask = mask_ptr;
- const char * GGML_CUDA_RESTRICT sinks = sinks_ptr;
- const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr;
- float * GGML_CUDA_RESTRICT dst = dst_ptr;
- float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
- // Skip unused kernel variants for faster compilation:
- if (use_logit_softcap && !(D == 128 || D == 256)) {
- NO_DEVICE_CODE;
- return;
- }
-
- //In this kernel Q, K, V are matrices while i, j, k are matrix indices.
-
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- const int ic0 = ncols*blockIdx.x; // Index of the first Q/QKV column to work on.
-
- static_assert(D <= FATTN_KQ_STRIDE, "D must be <= FATTN_KQ_STRIDE.");
- static_assert(ncols == 8 || ncols % 16 == 0, "ncols must be 8 or a multiple of 16.");
- constexpr int frag_m = ncols == 8 ? 32 : 16;
- constexpr int frag_n = ncols == 8 ? 8 : 16;
- static_assert(D % frag_m == 0, "If ncols == 8 then D % frag_m must be 0.");
-#if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000
- typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, _Float16, wmma::row_major> frag_a_K;
- typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, _Float16, wmma::col_major> frag_a_V;
- typedef wmma::fragment<wmma::matrix_b, frag_m, frag_n, 16, _Float16, wmma::col_major> frag_b;
- typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, KQ_acc_t> frag_c_KQ;
- typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, _Float16> frag_c_VKQ;
-#else
- typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, half, wmma::row_major> frag_a_K;
- typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, half, wmma::col_major> frag_a_V;
- typedef wmma::fragment<wmma::matrix_b, frag_m, frag_n, 16, half, wmma::col_major> frag_b;
- typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, KQ_acc_t> frag_c_KQ;
- typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, half> frag_c_VKQ;
-#endif
-
- constexpr int KQ_stride_tc = nwarps*frag_m; // Number of KQ rows calculated in parallel.
- constexpr int VKQ_ratio = KQ_stride_tc/VKQ_stride; // Number of parallel VKQ accumulators needed to keep all warps busy.
- static_assert(VKQ_ratio <= nwarps, "VKQ_ratio must be <= nwarps.");
-
- // Pad internal representation of KQ, KQV to reduce shared memory bank conflicts:
- constexpr int D_padded = D + 8;
- constexpr int kqs_padded = FATTN_KQ_STRIDE + 8;
- constexpr int kqar = sizeof(KQ_acc_t)/sizeof(half);
-
- ggml_cuda_pdl_sync();
- const int sequence = blockIdx.z / ne02;
- const int head = blockIdx.z - sequence*ne02;
- const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix.
- const float * Q_f = (const float *) (Q + nb03* sequence + nb02* head + nb01*ic0);
- const half * K_h = (const half *) (K + nb13* sequence + nb12*(head / gqa_ratio));
- const half * V_h = (const half *) (V + nb13* sequence + nb12*(head / gqa_ratio)); // K and V have same shape
- const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
- const half2 * mask2 = (const half2 *) maskh;
- const float * sinksf = (const float *) sinks;
-
- const int stride_Q = nb01 / sizeof(float);
- const int stride_KV = nb11 / sizeof(half);
-
- const float slopef = get_alibi_slope(max_bias, head, n_head_log2, m0, m1);
- const half slopeh = __float2half(slopef);
- const half2 slope2 = make_half2(slopef, slopef);
-
- const half2 logit_softcap_2 = make_half2(logit_softcap, logit_softcap);
-
- frag_b Q_b[D/16][ncols/frag_n];
-
- // A single buffer for temporarily holding tiles of KQ and VKQ parts:
- constexpr int mem_KQ = ncols*kqs_padded*kqar;
- constexpr int mem_VKQ_parts = VKQ_ratio*ncols*D_padded;
- __shared__ half KQ[mem_KQ >= mem_VKQ_parts ? mem_KQ : mem_VKQ_parts];
- float * KQ_f = (float *) KQ;
- half2 * KQ2 = (half2 *) KQ;
-
- float KQ_rowsum_f[ncols/nwarps] = {0.0f};
- float KQ_max_f[ncols/nwarps];
- float KQ_max_scale_f[ncols/nwarps] = {0.0f};
-
-#pragma unroll
- for (int j = 0; j < ncols/nwarps; ++j) {
- KQ_max_f[j] = -FLT_MAX/2.0f;
- }
-
- half2 KQ_rowsum_h2[ncols/nwarps] = {{0.0f, 0.0f}};
- half2 KQ_max_h2[ncols/nwarps];
- half2 KQ_max_scale_h2[ncols/nwarps] = {{0.0f, 0.0f}};
-
-#pragma unroll
- for (int j = 0; j < ncols/nwarps; ++j) {
- KQ_max_h2[j] = make_half2(-HALF_MAX_HALF, -HALF_MAX_HALF);
- }
-
- __shared__ half VKQ[ncols*D_padded]; // Accumulator for final VKQ slice.
- half2 * VKQ2 = (half2 *) VKQ;
-
-#if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000
- const _Float16 * K_h_f16 = reinterpret_cast<const _Float16 *>(K_h);
- const _Float16 * V_h_f16 = reinterpret_cast<const _Float16 *>(V_h);
- _Float16 * KQ_f16 = reinterpret_cast<_Float16 *>(KQ);
- _Float16 * VKQ_f16 = reinterpret_cast<_Float16 *>(VKQ);
-#else
- const half * K_h_f16 = K_h;
- const half * V_h_f16 = V_h;
- half * KQ_f16 = KQ;
- half * VKQ_f16 = VKQ;
-#endif
-
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-#pragma unroll
- for (int i0 = 0; i0 < D/2; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- if (i0 + warp_size > D/2 && i >= D/2) {
- break;
- }
- VKQ2[j*(D_padded/2) + i] = make_half2(0.0f, 0.0f);
- }
- }
-
- // Convert Q to half and apply scale, temporarily store in KQ:
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-#pragma unroll
- for (int i0 = 0; i0 < D; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- if (i0 + warp_size > D && i >= D) {
- break;
- }
- KQ[j*D_padded + i] = ic0 + j < int(ne01.z) ? Q_f[j*stride_Q + i] * scale : 0.0f;
- }
- }
-
- __syncthreads();
-
- // Load Q into tensor core fragments/registers since it will be used frequently:
-#pragma unroll
- for (int i0 = 0; i0 < D; i0 += 16) {
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += frag_n) {
- wmma::load_matrix_sync(Q_b[i0/16][j0/frag_n], KQ_f16 + j0*D_padded + i0, D_padded);
- }
- }
-
- __syncthreads();
-
- // Iterate over ne11 == previous tokens:
- const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11;
- for (int k_VKQ_0 = blockIdx.y*FATTN_KQ_STRIDE; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*FATTN_KQ_STRIDE) {
- // Calculate tile of KQ:
-#pragma unroll
- for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE; i_KQ_0 += KQ_stride_tc) {
- frag_c_KQ KQ_c[ncols/frag_n];
-#pragma unroll
- for (int j = 0; j < ncols/frag_n; ++j) {
- wmma::fill_fragment(KQ_c[j], static_cast<KQ_acc_t>(0.0f));
- }
-#pragma unroll
- for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += 16) {
- frag_a_K K_a;
- wmma::load_matrix_sync(K_a, K_h_f16 + int64_t(k_VKQ_0 + i_KQ_0 + frag_m*threadIdx.y)*stride_KV + k_KQ_0, stride_KV);
-#pragma unroll
- for (int j = 0; j < ncols/frag_n; ++j) {
- wmma::mma_sync(KQ_c[j], K_a, Q_b[k_KQ_0/16][j], KQ_c[j]);
- }
- }
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += frag_n) {
- wmma::store_matrix_sync((KQ_acc_t *) KQ + j0*kqs_padded + i_KQ_0 + frag_m*threadIdx.y, KQ_c[j0/frag_n], kqs_padded, wmma::mem_col_major);
- }
- }
-
- __syncthreads();
-
- // Calculate softmax for each KQ column using the current max. value.
- // The divisor is stored in KQ_rowsum and will be applied at the end.
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
- if (std::is_same<KQ_acc_t, float>::value) {
- float KQ_f_tmp[FATTN_KQ_STRIDE / warp_size];
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
- const int k = k0 + threadIdx.x;
-
- KQ_f_tmp[k0/warp_size] = KQ_f[j*kqs_padded + k];
-
- if (use_logit_softcap) {
- KQ_f_tmp[k0/warp_size] = logit_softcap*tanhf(KQ_f_tmp[k0/warp_size]);
- }
- }
-
- float KQ_max_new = KQ_max_f[j0/nwarps];
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
- const int k = k0 + threadIdx.x;
-
- KQ_f_tmp[k0/warp_size] += mask && ic0 + j < int(ne01.z) ?
- __half2float(slopeh*maskh[j*(nb31/sizeof(half)) + k_VKQ_0 + k]) : 0.0f;
- KQ_max_new = max(KQ_max_new, KQ_f_tmp[k0/warp_size] + FATTN_KQ_MAX_OFFSET);
- }
- KQ_max_new = warp_reduce_max<warp_size>(KQ_max_new);
-
- const float diff = KQ_max_f[j0/nwarps] - KQ_max_new;
- KQ_max_scale_f[j0/nwarps] = expf(diff);
- if (diff <= SOFTMAX_FTZ_THRESHOLD) {
- KQ_max_scale_f[j0/nwarps] = 0.0f;
- }
- KQ_max_f[j0/nwarps] = KQ_max_new;
-
- float KQ_rowsum_add = 0.0f;
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
- const int k = k0 + threadIdx.x;
-
- const float diff = KQ_f_tmp[k0/warp_size] - KQ_max_f[j0/nwarps];
- KQ_f_tmp[k0/warp_size] = expf(diff);
- if (diff <= SOFTMAX_FTZ_THRESHOLD) {
- KQ_f_tmp[k0/warp_size] = 0.0f;
- }
- KQ_rowsum_add += KQ_f_tmp[k0/warp_size];
- KQ[j*(kqar*kqs_padded) + k] = KQ_f_tmp[k0/warp_size];
- }
- KQ_rowsum_add = warp_reduce_sum<warp_size>(KQ_rowsum_add);
-
- // Scale previous KQ_rowsum to account for a potential increase in KQ_max:
- KQ_rowsum_f[j0/nwarps] = KQ_max_scale_f[j0/nwarps]*KQ_rowsum_f[j0/nwarps] + KQ_rowsum_add;
- } else {
- half2 KQ2_tmp[FATTN_KQ_STRIDE/(2*warp_size)];
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
- const int k = k0 + threadIdx.x;
-
- KQ2_tmp[k0/warp_size] = KQ2[j*(kqs_padded/2) + k];
-
- if (use_logit_softcap) {
- // There is no dedicated tangens hyperbolicus function for half2.
- KQ2_tmp[k0/warp_size] = h2exp(KQ2_tmp[k0/warp_size]*make_half2(2.0f, 2.0f));
- KQ2_tmp[k0/warp_size] = (KQ2_tmp[k0/warp_size] - make_half2(1.0f, 1.0f))
- /(KQ2_tmp[k0/warp_size] + make_half2(1.0f, 1.0f));
-
- KQ2_tmp[k0/warp_size] *= logit_softcap_2;
- }
- }
-
- half2 KQ_max_new = KQ_max_h2[j0/nwarps];
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
- const int k = k0 + threadIdx.x;
-
- KQ2_tmp[k0/warp_size] += mask && ic0 + j < int(ne01.z) ? slope2*mask2[(j*ne11 + k_VKQ_0)/2 + k] : make_half2(0.0f, 0.0f);
- KQ_max_new = ggml_cuda_hmax2(KQ_max_new, KQ2_tmp[k0/warp_size]);
- }
- KQ_max_new = __half2half2(warp_reduce_max<warp_size>(ggml_cuda_hmax(__low2half(KQ_max_new), __high2half(KQ_max_new))));
- const half2 diff = KQ_max_h2[j0/nwarps] - KQ_max_new;
- KQ_max_scale_h2[j0/nwarps] = h2exp(diff);
- const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
- *((uint32_t *) &KQ_max_scale_h2[j0/nwarps]) &= ftz_mask;
- KQ_max_h2[j0/nwarps] = KQ_max_new;
-
- half2 KQ_rowsum_add = make_half2(0.0f, 0.0f);
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
- const int k = k0 + threadIdx.x;
-
- const half2 diff = KQ2_tmp[k0/warp_size] - KQ_max_h2[j0/nwarps];
- KQ2_tmp[k0/warp_size] = h2exp(diff);
- const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
- *((uint32_t *) &KQ2_tmp[k0/warp_size]) &= ftz_mask;
- KQ_rowsum_add += KQ2_tmp[k0/warp_size];
- KQ2[j*(kqs_padded/2) + k] = KQ2_tmp[k0/warp_size];
- }
- KQ_rowsum_add = warp_reduce_sum<warp_size>(KQ_rowsum_add);
-
- // Scale previous KQ_rowsum to account for a potential increase in KQ_max:
- KQ_rowsum_h2[j0/nwarps] = KQ_max_scale_h2[j0/nwarps]*KQ_rowsum_h2[j0/nwarps] + KQ_rowsum_add;
- }
- }
-
- __syncthreads();
-
- frag_b KQ_b[FATTN_KQ_STRIDE/(VKQ_ratio*16)][ncols/frag_n];
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += frag_n) {
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += VKQ_ratio*16) {
- const int k = k0 + (threadIdx.y % VKQ_ratio)*16;
- wmma::load_matrix_sync(
- KQ_b[k0/(VKQ_ratio*16)][j0/frag_n],
- KQ_f16 + j0*(kqar*kqs_padded) + k,
- kqar*kqs_padded);
- }
- }
-
- frag_c_VKQ VKQ_c[D/VKQ_stride][ncols/frag_n];
-#pragma unroll
- for (int i_VKQ_0 = 0; i_VKQ_0 < D; i_VKQ_0 += VKQ_stride) {
-#pragma unroll
- for (int j = 0; j < ncols/frag_n; ++j) {
- wmma::fill_fragment(VKQ_c[i_VKQ_0/VKQ_stride][j], static_cast<half>(0.0f));
- }
-
-#pragma unroll
- for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += VKQ_ratio*16) {
- const int k = k0 + (threadIdx.y % VKQ_ratio)*16;
-
- frag_a_V v_a;
- wmma::load_matrix_sync(v_a, V_h_f16 + int64_t(k_VKQ_0 + k)*stride_KV + i_VKQ_0 + frag_m*(threadIdx.y/VKQ_ratio), stride_KV);
-#pragma unroll
- for (int j = 0; j < ncols/frag_n; ++j) {
- wmma::mma_sync(VKQ_c[i_VKQ_0/VKQ_stride][j], v_a, KQ_b[k0/(VKQ_ratio*16)][j], VKQ_c[i_VKQ_0/VKQ_stride][j]);
- }
- }
- }
-
- __syncthreads();
-
- const int offset_k = (threadIdx.y % VKQ_ratio) * (ncols*D_padded);
-#pragma unroll
- for (int i_KQ_0 = 0; i_KQ_0 < D; i_KQ_0 += VKQ_stride) {
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += frag_n) {
- wmma::store_matrix_sync(
- KQ_f16 + offset_k + j0*D_padded + i_KQ_0 + frag_m*(threadIdx.y/VKQ_ratio),
- VKQ_c[i_KQ_0/VKQ_stride][j0/frag_n],
- D_padded, wmma::mem_col_major);
- }
- }
-
- __syncthreads();
-
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
- half2 VKQ_scale;
- if (std::is_same<KQ_acc_t, float>::value) {
- VKQ_scale = make_half2(KQ_max_scale_f[j0/nwarps], KQ_max_scale_f[j0/nwarps]);
- } else {
- VKQ_scale = KQ_max_scale_h2[j0/nwarps];
- }
-
-#pragma unroll
- for (int i0 = 0; i0 < D/2; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- if (i0 + warp_size > D/2 && i >= D/2) {
- break;
- }
-
- half2 VKQ_add = make_half2(0.0f, 0.0f);
-#pragma unroll
- for (int l = 0; l < VKQ_ratio; ++l) {
- VKQ_add += KQ2[l*(ncols*D_padded/2) + j*(D_padded/2) + i];
- }
- VKQ2[j*(D_padded/2) + i] = VKQ_scale*VKQ2[j*(D_padded/2) + i] + VKQ_add;
- }
- }
-
- __syncthreads();
- }
-
- // Apply attention sinks
- if (sinksf && blockIdx.y == 0) {
- const float sinkf = sinksf[head];
- const half sinkh = __float2half(sinkf);
-
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
- if (std::is_same<KQ_acc_t, float>::value) {
- float kqmax_new = fmaxf(KQ_max_f[j0/nwarps], sinkf);
-
- const float KQ_max_scale = expf(KQ_max_f[j0/nwarps] - kqmax_new);
- KQ_max_f[j0/nwarps] = kqmax_new;
-
- KQ_rowsum_f[j0/nwarps] = KQ_rowsum_f[j0/nwarps] * KQ_max_scale + expf(sinkf - KQ_max_f[j0/nwarps]);
-
- const half2 scale_h2 = make_half2(KQ_max_scale, KQ_max_scale);
-#pragma unroll
- for (int i0 = 0; i0 < D/2; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- if (i0 + warp_size > D/2 && i >= D/2) break;
- VKQ2[j*(D_padded/2) + i] *= scale_h2;
- }
- } else {
- half kqmax_old = __low2half(KQ_max_h2[j0/nwarps]);
- half kqmax_new = fmaxf(kqmax_old, sinkh);
- KQ_max_h2[j0/nwarps] = __half2half2(kqmax_new);
-
- const half KQ_max_scale_h = hexp(kqmax_old - kqmax_new);
- const half2 KQ_max_scale = __half2half2(KQ_max_scale_h);
-
- KQ_rowsum_h2[j0/nwarps] = KQ_rowsum_h2[j0/nwarps] * KQ_max_scale;
- const half val = hexp(sinkh - kqmax_new);
- KQ_rowsum_h2[j0/nwarps].x = __hadd(KQ_rowsum_h2[j0/nwarps].x, val);
-
-#pragma unroll
- for (int i0 = 0; i0 < D/2; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- if (i0 + warp_size > D/2 && i >= D/2) break;
- VKQ2[j*(D_padded/2) + i] *= KQ_max_scale;
- }
- }
- }
-
- __syncthreads();
- }
-#pragma unroll
- for (int j0 = 0; j0 < ncols; j0 += nwarps) {
- const int j_VKQ = j0 + threadIdx.y;
- if (ic0 + j_VKQ >= int(ne01.z)) {
- return;
- }
-
- float KQ_rowsum_j;
- if (std::is_same<KQ_acc_t, float>::value) {
- KQ_rowsum_j = KQ_rowsum_f[j0/nwarps];
- } else {
- KQ_rowsum_j = __low2float(KQ_rowsum_h2[j0/nwarps]) + __high2float(KQ_rowsum_h2[j0/nwarps]);
- }
-
- const int j_dst_unrolled = ((sequence*int(ne01.z) + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < D; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- if (i0 + warp_size > D && i >= D) {
- break;
- }
- float dst_val = VKQ[j_VKQ*D_padded + i];
- if (gridDim.y == 1) {
- dst_val /= KQ_rowsum_j;
- }
- dst[j_dst_unrolled*D + i] = dst_val;
- }
-
- if (gridDim.y == 1 || threadIdx.x != 0) {
- continue;
- }
-
- float2 dst_meta_val;
- if (std::is_same<KQ_acc_t, float>::value) {
- dst_meta_val.x = KQ_max_f[j0/nwarps];
- } else {
- dst_meta_val.x = __low2float(KQ_max_h2[j0/nwarps]);
- }
- dst_meta_val.y = KQ_rowsum_j;
- dst_meta[j_dst_unrolled] = dst_meta_val;
- }
-#else
- GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
- max_bias, m0, m1, n_head_log2, logit_softcap,
- ne00, ne01, ne02, ne03,
- nb01, nb02, nb03,
- ne10, ne11, ne12, ne13,
- nb11, nb12, nb13,
- nb21, nb22, nb23,
- ne31, ne32, ne33,
- nb31, nb32, nb33);
- NO_DEVICE_CODE;
-#endif // defined(FLASH_ATTN_AVAILABLE) && (defined(GGML_HIP_ROCWMMA_FATTN) && defined(GGML_USE_WMMA_FATTN))
-}
-
-constexpr int get_max_power_of_2(int x) {
- return x % 2 == 0 ? 2*get_max_power_of_2(x/2) : 1;
-}
-
-static_assert(get_max_power_of_2(1) == 1, "Test failed.");
-static_assert(get_max_power_of_2(2) == 2, "Test failed.");
-static_assert(get_max_power_of_2(4) == 4, "Test failed.");
-static_assert(get_max_power_of_2(6) == 2, "Test failed.");
-
-// Number of VKQ rows calculated in parallel:
-constexpr int get_VKQ_stride(int D, int nwarps, int frag_m) {
- return (get_max_power_of_2(D/frag_m) < nwarps ? get_max_power_of_2(D/frag_m) : nwarps)*frag_m;
-}
-
-static_assert(get_VKQ_stride(128, 1, 32) == 32, "Test failed.");
-static_assert(get_VKQ_stride(128, 2, 32) == 64, "Test failed.");
-static_assert(get_VKQ_stride(128, 4, 32) == 128, "Test failed.");
-static_assert(get_VKQ_stride( 64, 1, 32) == 32, "Test failed.");
-static_assert(get_VKQ_stride( 64, 2, 32) == 64, "Test failed.");
-static_assert(get_VKQ_stride( 64, 4, 32) == 64, "Test failed.");
-static_assert(get_VKQ_stride( 80, 1, 16) == 16, "Test failed.");
-static_assert(get_VKQ_stride( 80, 2, 16) == 16, "Test failed.");
-static_assert(get_VKQ_stride( 80, 4, 16) == 16, "Test failed.");
-
-template <int D, int cols_per_block, typename KQ_acc_t>
-void ggml_cuda_flash_attn_ext_wmma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
- const ggml_tensor * KQV = dst;
-
- constexpr int nwarps = 4;
-
- constexpr int frag_m = cols_per_block == 8 && D % 32 == 0 ? 32 : 16;
- const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
-
- float logit_softcap;
- memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
-
- fattn_kernel_t fattn_kernel;
- if (logit_softcap == 0.0f) {
- constexpr bool use_logit_softcap = false;
- fattn_kernel = flash_attn_ext_f16<
- D, cols_per_block, nwarps, get_VKQ_stride(D, nwarps, frag_m), KQ_acc_t, use_logit_softcap>;
- } else {
- constexpr bool use_logit_softcap = true;
- fattn_kernel = flash_attn_ext_f16<
- D, cols_per_block, nwarps, get_VKQ_stride(D, nwarps, frag_m), KQ_acc_t, use_logit_softcap>;
- }
- launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, 0, FATTN_KQ_STRIDE, true, true, false, warp_size);
-}
-
-void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
- const ggml_tensor * KQV = dst;
- const ggml_tensor * Q = dst->src[0];
-
- const enum ggml_prec prec = ggml_flash_attn_ext_get_prec(KQV);
- const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
-
- if (prec != GGML_PREC_DEFAULT) {
- if (Q->ne[1] <= 32 || Q->ne[0] > 128) {
- constexpr int cols_per_block = 16;
- switch (Q->ne[0]) {
- case 64:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, float>(ctx, dst);
- break;
- case 80:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, float>(ctx, dst);
- break;
- case 96:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, float>(ctx, dst);
- break;
- case 112:
- ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, float>(ctx, dst);
- break;
- case 128:
- ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, float>(ctx, dst);
- break;
- case 256:
- ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, float>(ctx, dst);
- break;
- default:
- GGML_ABORT("fatal error");
- break;
- }
- } else {
- constexpr int cols_per_block = 32;
- switch (Q->ne[0]) {
- case 64:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, float>(ctx, dst);
- break;
- case 80:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, float>(ctx, dst);
- break;
- case 96:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, float>(ctx, dst);
- break;
- case 112:
- ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, float>(ctx, dst);
- break;
- case 128:
- ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, float>(ctx, dst);
- break;
- // case 256:
- // ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, float>(ctx, dst);
- // break;
- default:
- GGML_ABORT("fatal error");
- break;
- }
- }
- return;
- }
-
-#if !defined(GGML_USE_HIP)
- if (Q->ne[1] <= 8 && Q->ne[0] % warp_size == 0) {
- constexpr int cols_per_block = 8;
- switch (Q->ne[0]) {
- case 64:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
- break;
- case 96:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
- break;
- case 128:
- ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
- break;
- case 256:
- ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
- break;
- default:
- GGML_ABORT("fatal error");
- break;
- }
- return;
- }
-#endif // !defined(GGML_USE_HIP)
-
- if (Q->ne[1] <= 32) {
- constexpr int cols_per_block = 16;
- switch (Q->ne[0]) {
- case 64:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
- break;
- case 80:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, half>(ctx, dst);
- break;
- case 96:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
- break;
- case 112:
- ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, half>(ctx, dst);
- break;
- case 128:
- ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
- break;
- case 256:
- ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
- break;
- default:
- GGML_ABORT("fatal error");
- break;
- }
- return;
- }
-
- constexpr int cols_per_block = 32;
- switch (Q->ne[0]) {
- case 64:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
- break;
- case 80:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, half>(ctx, dst);
- break;
- case 96:
- ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
- break;
- case 112:
- ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, half>(ctx, dst);
- break;
- case 128:
- ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
- break;
- case 256:
- ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
- break;
- default:
- GGML_ABORT("fatal error");
- break;
- }
-}
+++ /dev/null
-#pragma once
-
-#include "common.cuh"
-
-#if defined(GGML_USE_MUSA)
-#define GGML_USE_WMMA_FATTN
-#endif // defined(GGML_USE_MUSA)
-
-#if defined(GGML_HIP_ROCWMMA_FATTN)
-#if defined(CDNA) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
-#define GGML_USE_WMMA_FATTN
-#elif defined(CDNA)
-#warning "rocwmma fattn on CDNA is broken on rocwmma v2.0.0, expect degraded performance"
-#endif // defined(CDNA) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
-#if defined(RDNA3)
-#define GGML_USE_WMMA_FATTN
-#endif // defined(RDNA3)
-#if defined(RDNA4) && ROCWMMA_VERSION_MAJOR > 1
-#define GGML_USE_WMMA_FATTN
-#elif defined(RDNA4)
-#warning "rocwmma fattn is not supported on RDNA4 on rocwmma < v2.0.0, expect degraded performance"
-#endif // defined(RDNA4) && ROCWMMA_VERSION_MAJOR > 1
-#endif // defined(GGML_HIP_ROCWMMA_FATTN)
-
-// WMMA flash attention requires FP16 matrix instructions to be available for ggml code.
-static bool ggml_cuda_should_use_wmma_fattn(const int cc) {
-#if defined(GGML_USE_HIP) && !defined(GGML_HIP_ROCWMMA_FATTN)
- return false;
-#else
- if ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_VOLTA) ||
- GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_MTHREADS(cc)) {
- return true;
- } else if (GGML_CUDA_CC_IS_CDNA(cc)){
-#if defined(GGML_HIP_ROCWMMA_FATTN) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
- return true;
-#else
- return false;
-#endif // defined(GGML_HIP_ROCWMMA_FATTN) (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
- } else if (GGML_CUDA_CC_IS_RDNA4(cc)) {
-#if defined(GGML_HIP_ROCWMMA_FATTN) && ROCWMMA_VERSION_MAJOR > 1
- return true;
-#else
- return false;
-#endif // defined(GGML_HIP_ROCWMMA_FATTN) && ROCWMMA_VERSION_MAJOR > 1
- } else {
- return false;
- }
-#endif // defined(GGML_USE_HIP) && !defined(GGML_HIP_ROCWMMA_FATTN)
-}
-
-void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
#include "fattn-mma-f16.cuh"
#include "fattn-tile.cuh"
#include "fattn-vec.cuh"
-#include "fattn-wmma-f16.cuh"
#include "fattn.cuh"
template <int DKQ, int DV, int ncols2>
// Best FlashAttention kernel for a specific GPU:
enum best_fattn_kernel {
- BEST_FATTN_KERNEL_NONE = 0,
- BEST_FATTN_KERNEL_TILE = 200,
- BEST_FATTN_KERNEL_VEC = 100,
- BEST_FATTN_KERNEL_WMMA_F16 = 300,
- BEST_FATTN_KERNEL_MMA_F16 = 400,
+ BEST_FATTN_KERNEL_NONE = 0,
+ BEST_FATTN_KERNEL_TILE = 200,
+ BEST_FATTN_KERNEL_VEC = 100,
+ BEST_FATTN_KERNEL_MMA_F16 = 400,
};
static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
return BEST_FATTN_KERNEL_MMA_F16;
}
- // Use the WMMA kernel if possible:
- if (ggml_cuda_should_use_wmma_fattn(cc) && K->ne[1] % FATTN_KQ_STRIDE == 0 && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[0] != 192 && Q->ne[0] != 512 && Q->ne[0] != 576) {
- if (can_use_vector_kernel && Q->ne[1] <= 2) {
- return BEST_FATTN_KERNEL_VEC;
- }
- return BEST_FATTN_KERNEL_WMMA_F16;
- }
-
// AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
if ((amd_mfma_available(cc) && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72) {
if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {
switch (kernel) {
case BEST_FATTN_KERNEL_TILE:
- case BEST_FATTN_KERNEL_WMMA_F16:
case BEST_FATTN_KERNEL_MMA_F16:
need_f16_K = true;
need_f16_V = true;
case BEST_FATTN_KERNEL_VEC:
ggml_cuda_flash_attn_ext_vec(ctx, dst);
break;
- case BEST_FATTN_KERNEL_WMMA_F16:
- ggml_cuda_flash_attn_ext_wmma_f16(ctx, dst);
- break;
case BEST_FATTN_KERNEL_MMA_F16:
ggml_cuda_flash_attn_ext_mma_f16(ctx, dst);
break;
}
}
+template<typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static __global__ void k_get_rows_kq(
+ const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,
+ const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/
+ /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/
+ /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,
+ /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
+ const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
+
+ ggml_cuda_pdl_sync();
+ const int64_t nsb = ne00/QK_K; // super-blocks per row
+ for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+ // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+ const int i10 = blockIdx.x;
+ const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+ const int i11 = dm.x;
+ const int i12 = dm.y;
+
+ const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+ dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+ const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+ for (int64_t ib = blockIdx.y; ib < nsb; ib += gridDim.y) {
+ dequantize_kq(src0_row, ib, dst_row + ib*QK_K, threadIdx.x);
+ }
+ }
+}
+
template<typename src0_t, typename dst_t>
static __global__ void k_get_rows_float(
const src0_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
dst_t * GGML_CUDA_RESTRICT dst = dst_ptr;
ggml_cuda_pdl_sync();
for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+ // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+ const int i10 = blockIdx.x;
+ const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+ const int i11 = dm.x;
+ const int i12 = dm.y;
+
+ const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+ dst_t * GGML_CUDA_RESTRICT dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+ const src0_t * GGML_CUDA_RESTRICT src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
+
for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) {
- // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
- const int i10 = blockIdx.x;
- const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
- const int i11 = dm.x;
- const int i12 = dm.y;
+ dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
+ }
+ }
+}
- if (i00 >= ne00) {
- return;
- }
+template<typename dst_t>
+static __global__ void k_get_rows_float_vec(
+ const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
+ const int64_t ne00v,
+ const int64_t ne11, const uint3 ne12_fdv,
+ const size_t s1, const size_t s2, const size_t s3,
+ const size_t nb01, const size_t nb02, const size_t nb03,
+ const size_t s10, const size_t s11, const size_t s12) {
- const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+ ggml_cuda_pdl_lc();
+ ggml_cuda_pdl_sync();
+ for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+ const int i10 = blockIdx.x;
+ const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+ const int i11 = dm.x;
+ const int i12 = dm.y;
- dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
- const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
+ const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12];
- dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
+ int4 * GGML_CUDA_RESTRICT dst_row = (int4 *) (dst_ptr + i10*s1 + i11*s2 + i12*s3);
+ const int4 * GGML_CUDA_RESTRICT src0_row = (const int4 *)((const char *) src0_ptr + i01*nb01 + i11*nb02 + i12*nb03);
+
+ for (int64_t i = blockIdx.y*blockDim.x + threadIdx.x; i < ne00v; i += gridDim.y*blockDim.x) {
+ dst_row[i] = src0_row[i];
}
}
}
s10, s11, s12/*, s13*/);
}
+template<int block_dim, typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static void get_rows_cuda_kq(
+ const void * src0_d, const int32_t * src1_d, dst_t * dst_d,
+ const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+ const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+ const size_t nb1, const size_t nb2, const size_t nb3,
+ cudaStream_t stream) {
+ GGML_ASSERT(ne00 % QK_K == 0);
+ const int64_t nsb = ne00/QK_K;
+
+ const dim3 block_dims(block_dim, 1, 1);
+ const dim3 block_nums(ne10, MIN(nsb, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
+ // strides in elements
+ // const size_t s0 = nb0 / sizeof(dst_t);
+ const size_t s1 = nb1 / sizeof(dst_t);
+ const size_t s2 = nb2 / sizeof(dst_t);
+ const size_t s3 = nb3 / sizeof(dst_t);
+
+ const size_t s10 = nb10 / sizeof(int32_t);
+ const size_t s11 = nb11 / sizeof(int32_t);
+ const size_t s12 = nb12 / sizeof(int32_t);
+ // const size_t s13 = nb13 / sizeof(int32_t);
+
+ GGML_ASSERT(ne12 > 0);
+ GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
+ const uint3 ne12_fdv = init_fastdiv_values(ne12);
+
+ k_get_rows_kq<dst_t, dequantize_kq><<<block_nums, block_dims, 0, stream>>>(
+ src0_d, src1_d, dst_d,
+ ne00, /*ne01, ne02, ne03,*/
+ /*ne10,*/ ne11, ne12_fdv, /*ne13,*/
+ /* s0,*/ s1, s2, s3,
+ /* nb00,*/ nb01, nb02, nb03,
+ s10, s11, s12/*, s13*/);
+}
+
template<typename src0_t, typename dst_t>
static void get_rows_cuda_float(
const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d,
const size_t nb1, const size_t nb2, const size_t nb3,
cudaStream_t stream) {
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
- const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
- const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
// strides in elements
// const size_t s0 = nb0 / sizeof(dst_t);
GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
const uint3 ne12_fdv = init_fastdiv_values(ne12);
+ if constexpr (std::is_same<src0_t, dst_t>::value) {
+ constexpr int VEC = 16 / sizeof(dst_t);
+ const int64_t ne00v = ne00 / VEC;
+ const int64_t vec_block_num_y = (ne00v + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+ const bool enough_blocks = vec_block_num_y * ne10 * ne11 * ne12 >= 128;
+ const bool can_vec = VEC > 1 && enough_blocks &&
+ (ne00 % VEC == 0) &&
+ (nb01 % 16 == 0) && (nb02 % 16 == 0) && (nb03 % 16 == 0) &&
+ (nb1 % 16 == 0) && (nb2 % 16 == 0) && (nb3 % 16 == 0) &&
+ (((uintptr_t) src0_d) % 16 == 0) && (((uintptr_t) dst_d) % 16 == 0);
+
+ if (can_vec) {
+ const int block_num_y = vec_block_num_y;
+ const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
+ ggml_cuda_kernel_launch(k_get_rows_float_vec<dst_t>, launch_params,
+ (const dst_t *) src0_d, src1_d, dst_d,
+ ne00v, ne11, ne12_fdv,
+ s1, s2, s3,
+ nb01, nb02, nb03,
+ s10, s11, s12);
+ return;
+ }
+ }
+
+ const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+ const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
ggml_cuda_kernel_launch(k_get_rows_float<src0_t, dst_t>, launch_params,
src0_d, src1_d, dst_d,
get_rows_cuda_q<QK1_0, QR1_0, dequantize_q1_0>(src0_d, src1_d, dst_d,
ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
break;
+ case GGML_TYPE_Q2_0:
+ get_rows_cuda_q<QK2_0, QR2_0, dequantize_q2_0>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
case GGML_TYPE_Q4_0:
get_rows_cuda_q<QK4_0, QR4_0, dequantize_q4_0>(src0_d, src1_d, dst_d,
ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
get_rows_cuda_q<QK8_0, QR8_0, dequantize_q8_0>(src0_d, src1_d, dst_d,
ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
break;
+ case GGML_TYPE_Q2_K:
+ get_rows_cuda_kq<64, dst_t, dequantize_q2_K<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_Q3_K:
+ get_rows_cuda_kq<64, dst_t, dequantize_q3_K<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_Q4_K:
+ get_rows_cuda_kq<32, dst_t, dequantize_q4_K<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_Q5_K:
+ get_rows_cuda_kq<64, dst_t, dequantize_q5_K<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_Q6_K:
+ get_rows_cuda_kq<64, dst_t, dequantize_q6_K<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ2_XXS:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq2_xxs<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ2_XS:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq2_xs<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ2_S:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq2_s<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ3_XXS:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq3_xxs<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ3_S:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq3_s<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ1_S:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq1_s<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ1_M:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq1_m<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ4_NL:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq4_nl<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_IQ4_XS:
+ get_rows_cuda_kq<32, dst_t, dequantize_iq4_xs<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
+ case GGML_TYPE_MXFP4:
+ get_rows_cuda_kq<32, dst_t, dequantize_mxfp4<dst_t>>(src0_d, src1_d, dst_d,
+ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+ break;
default:
- // TODO: k-quants
GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type));
break;
}
#include "ggml-cuda/wkv.cuh"
#include "ggml-cuda/gla.cuh"
#include "ggml-cuda/gated_delta_net.cuh"
+#include "ggml-cuda/dsv4-hc.cuh"
#include "ggml-cuda/set.cuh"
#include "ggml-cuda/set-rows.cuh"
#include "ggml-cuda/pad_reflect_1d.cuh"
#include "ggml-cuda/tri.cuh"
#include "ggml-cuda/cumsum.cuh"
#include "ggml-cuda/fill.cuh"
+#include "ggml-cuda/lightning-indexer.cuh"
#include "ggml.h"
#include <algorithm>
GGML_ABORT(GGML_CUDA_NAME " error");
}
+// map a (possibly virtual) device id to the physical CUDA device that backs it
+static int ggml_cuda_get_physical_device(int device) {
+ const ggml_cuda_device_info & info = ggml_cuda_info();
+ GGML_ASSERT(device >= 0 && device < info.device_count);
+ return info.devices[device].physical_device;
+}
+
// this is faster on Windows
// probably because the Windows CUDA libraries forget to make this check before invoking the drivers
void ggml_cuda_set_device(int device) {
+ // translate the (possibly virtual) device id to the physical CUDA device that backs it
+ const int physical_device = ggml_cuda_get_physical_device(device);
+
int current_device;
CUDA_CHECK(cudaGetDevice(¤t_device));
- if (device == current_device) {
+ if (physical_device == current_device) {
return;
}
- CUDA_CHECK(cudaSetDevice(device));
+ CUDA_CHECK(cudaSetDevice(physical_device));
}
int ggml_cuda_get_device() {
static ggml_cuda_device_info ggml_cuda_init() {
ggml_cuda_device_info info = {};
- cudaError_t err = cudaGetDeviceCount(&info.device_count);
+ cudaError_t err = cudaGetDeviceCount(&info.physical_device_count);
if (err != cudaSuccess) {
GGML_LOG_ERROR("%s: failed to initialize " GGML_CUDA_NAME ": %s\n", __func__, cudaGetErrorString(err));
return info;
}
- GGML_ASSERT(info.device_count <= GGML_CUDA_MAX_DEVICES);
+ GGML_ASSERT(info.physical_device_count <= GGML_CUDA_MAX_DEVICES);
- int64_t total_vram = 0;
+ // by default expose exactly the physical devices; GGML_CUDA_DEVICES can request a different
+ // number of (virtual) devices to emulate multi-GPU systems on a machine with fewer GPUs
+ info.device_count = info.physical_device_count;
+
+ const char * devices_env = getenv("GGML_CUDA_DEVICES");
+ if (devices_env != nullptr && info.physical_device_count > 0) {
+ const int requested = atoi(devices_env);
+ if (requested > 0) {
+ info.device_count = requested;
+ } else {
+ GGML_LOG_WARN("%s: ignoring invalid GGML_CUDA_DEVICES=\"%s\"\n", __func__, devices_env);
+ }
+ }
+
+ if (info.device_count > GGML_CUDA_MAX_DEVICES) {
+ GGML_LOG_WARN("%s: requested %d devices, clamping to GGML_CUDA_MAX_DEVICES=%d\n",
+ __func__, info.device_count, GGML_CUDA_MAX_DEVICES);
+ info.device_count = GGML_CUDA_MAX_DEVICES;
+ }
+
+ // map each (virtual) device to a backing physical device (round-robin), assign each its index
+ // among the (virtual) devices sharing that physical GPU, and store the per-physical share count
+ int physical_share_count[GGML_CUDA_MAX_DEVICES] = {};
+ GGML_ASSERT(info.device_count == 0 || info.physical_device_count > 0);
for (int id = 0; id < info.device_count; ++id) {
+ info.devices[id].physical_device = id % info.physical_device_count;
+ info.devices[id].virtual_index = physical_share_count[info.devices[id].physical_device]++;
+ }
+
+ int64_t total_vram = 0;
+ for (int id = 0; id < info.physical_device_count; ++id) {
cudaDeviceProp prop;
CUDA_CHECK(cudaGetDeviceProperties(&prop, id));
total_vram += prop.totalGlobalMem;
}
GGML_LOG_INFO("%s: found %d " GGML_CUDA_NAME " devices (Total VRAM: %zu MiB):\n",
- __func__, info.device_count, (size_t)(total_vram / (1024 * 1024)));
+ __func__, info.physical_device_count, (size_t)(total_vram / (1024 * 1024)));
+ if (info.device_count != info.physical_device_count) {
+ GGML_LOG_INFO("%s: emulating %d virtual device(s) on %d physical device(s) (GGML_CUDA_DEVICES)\n",
+ __func__, info.device_count, info.physical_device_count);
+ }
total_vram = 0;
std::vector<std::pair<int, std::string>> turing_devices_without_mma;
for (int id = 0; id < info.device_count; ++id) {
+ const int physical_id = info.devices[id].physical_device;
+
int device_vmm = 0;
#if defined(GGML_USE_VMM)
CUdevice device;
- CU_CHECK(cuDeviceGet(&device, id));
+ CU_CHECK(cuDeviceGet(&device, physical_id));
CU_CHECK(cuDeviceGetAttribute(&device_vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, device));
if (device_vmm) {
CUmemAllocationProp alloc_prop = {};
alloc_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
alloc_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
- alloc_prop.location.id = id;
+ alloc_prop.location.id = physical_id;
CU_CHECK(cuMemGetAllocationGranularity(&info.devices[id].vmm_granularity, &alloc_prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
}
#endif // defined(GGML_USE_VMM)
info.devices[id].vmm = !!device_vmm;
cudaDeviceProp prop;
- CUDA_CHECK(cudaGetDeviceProperties(&prop, id));
+ CUDA_CHECK(cudaGetDeviceProperties(&prop, physical_id));
+
+ // a virtual device owns only a share of its physical GPU's memory; report that share so the
+ // logged per-device VRAM sums to the physical total above.
+ GGML_ASSERT(physical_share_count[physical_id] > 0);
+ info.devices[id].physical_share_count = physical_share_count[physical_id];
+ const size_t device_vram = prop.totalGlobalMem / info.devices[id].physical_share_count;
+ const size_t device_vram_mib = device_vram / (1024 * 1024);
info.default_tensor_split[id] = total_vram;
- total_vram += prop.totalGlobalMem;
+ total_vram += device_vram;
+#if defined(GGML_USE_HIP)
+ info.devices[id].integrated = prop.integrated;
+#else
info.devices[id].integrated = false; // Temporarily disabled due to issues with corrupted output (e.g. #15034)
+#endif
info.devices[id].nsm = prop.multiProcessorCount;
info.devices[id].smpb = prop.sharedMemPerBlock;
info.devices[id].warp_size = prop.warpSize;
#ifndef GGML_USE_MUSA
int supports_coop_launch = 0;
- CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, id));
+ CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, physical_id));
info.devices[id].supports_cooperative_launch = !!supports_coop_launch;
#else
info.devices[id].supports_cooperative_launch = false;
GGML_LOG_INFO(" Device %d: %s, %s (0x%x), VMM: %s, Wave Size: %d, VRAM: %zu MiB\n",
id, prop.name, prop.gcnArchName, info.devices[id].cc & 0xffff,
device_vmm ? "yes" : "no", prop.warpSize,
- (size_t)(prop.totalGlobalMem / (1024 * 1024)));
+ device_vram_mib);
#elif defined(GGML_USE_MUSA)
// FIXME: Ensure compatibility with varying warp sizes across different MUSA archs.
info.devices[id].warp_size = 32;
info.devices[id].cc += prop.minor * 0x10;
GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n",
id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no",
- (size_t)(prop.totalGlobalMem / (1024 * 1024)));
+ device_vram_mib);
#else
info.devices[id].smpbo = prop.sharedMemPerBlockOptin;
info.devices[id].cc = 100*prop.major + 10*prop.minor;
GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n",
id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no",
- (size_t)(prop.totalGlobalMem / (1024 * 1024)));
+ device_vram_mib);
std::string device_name(prop.name);
if (device_name == "NVIDIA GeForce MX450") {
turing_devices_without_mma.push_back({ id, device_name });
// TODO: Check for future drivers the default scheduling strategy and
// remove this call again when cudaDeviceScheduleSpin is default.
if (prop.major == 12 && prop.minor == 1) {
- CUDA_CHECK(cudaSetDevice(id));
+ CUDA_CHECK(cudaSetDevice(physical_id));
CUDA_CHECK(cudaSetDeviceFlags(cudaDeviceScheduleSpin));
}
// CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr));
if (getenv("GGML_CUDA_P2P") != nullptr) {
- for (int id = 0; id < info.device_count; ++id) {
- ggml_cuda_set_device(id);
- for (int id_other = 0; id_other < info.device_count; ++id_other) {
+ for (int id = 0; id < info.physical_device_count; ++id) {
+ CUDA_CHECK(cudaSetDevice(id));
+ for (int id_other = 0; id_other < info.physical_device_count; ++id_other) {
if (id == id_other) {
continue;
}
static const size_t CUDA_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB
int device;
+ int physical_device;
CUdeviceptr pool_addr = 0;
size_t pool_used = 0;
size_t pool_size = 0;
explicit ggml_cuda_pool_vmm(int device) :
device(device),
+ physical_device(ggml_cuda_get_physical_device(device)),
granularity(ggml_cuda_info().devices[device].vmm_granularity) {
}
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
- prop.location.id = device;
+ prop.location.id = physical_device;
CUmemGenericAllocationHandle handle;
CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0));
// NCCL implicitly enables peer access (cudaDeviceEnablePeerAccess), and
// GGML_CUDA_P2P enables it explicitly. Unlike cudaMalloc buffers, VMM
// allocations do not become peer-accessible from that alone, so access
- // must be granted explicitly here.
+ // must be granted explicitly here. With virtual devices, grant access
+ // on the backing *physical* devices (deduplicated, since several
+ // virtual devices can map to the same physical GPU).
std::vector<CUmemAccessDesc> access_descs;
+ bool physical_seen[GGML_CUDA_MAX_DEVICES] = {};
const int device_count = ggml_cuda_info().device_count;
for (int id = 0; id < device_count; ++id) {
- if (id != device) {
+ const int id_physical = ggml_cuda_get_physical_device(id);
+ if (id_physical != physical_device) {
int can_access_peer = 0;
- CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, device));
+ CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id_physical, physical_device));
if (!can_access_peer) {
continue;
}
}
+ if (physical_seen[id_physical]) {
+ continue;
+ }
+ physical_seen[id_physical] = true;
CUmemAccessDesc access = {};
access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
- access.location.id = id;
+ access.location.id = id_physical;
access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
access_descs.push_back(access);
}
// set access for non P2P
CUmemAccessDesc access = {};
access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
- access.location.id = device;
+ access.location.id = physical_device;
access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
CU_CHECK(cuMemSetAccess(start_ptr, reserve_size, &access, 1));
}
if (ggml_backend_buffer_is_cuda(src->buffer)) {
ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context;
ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context;
- if (src_ctx->device == dst_ctx->device) {
+ // compare the backing physical devices: distinct virtual devices may share one physical GPU,
+ // in which case a same-device copy (not a peer copy) is required
+ const int src_physical = ggml_cuda_get_physical_device(src_ctx->device);
+ const int dst_physical = ggml_cuda_get_physical_device(dst_ctx->device);
+ if (src_physical == dst_physical) {
CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread));
} else {
#ifdef GGML_CUDA_NO_PEER_COPY
return false;
#else
- CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_ctx->device, src->data, src_ctx->device, ggml_nbytes(src), cudaStreamPerThread));
+ CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_physical, src->data, src_physical, ggml_nbytes(src), cudaStreamPerThread));
#endif
}
CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
static void ggml_backend_cuda_comm_init_nccl(ggml_backend_cuda_comm_context * ret) {
#ifdef GGML_USE_NCCL
+ // Disabling NCCL path when CUDA virtual devices are in use since NCCL requires one distinct physical GPU per rank.
+ const ggml_cuda_device_info & info = ggml_cuda_info();
+ if (info.device_count > info.physical_device_count) {
+ GGML_LOG_WARN("NCCL disabled: virtual devices in use; "
+ "falling back to internal AllReduce\n");
+ ggml_backend_cuda_comm_init_internal(ret);
+ return;
+ }
+
const size_t n = ret->dev_ids.size();
ret->comms.resize(n);
ncclResult_t rc = ncclCommInitAll(ret->comms.data(), (int) n, ret->dev_ids.data());
case GGML_OP_GATED_DELTA_NET:
ggml_cuda_op_gated_delta_net(ctx, dst);
break;
+ case GGML_OP_DSV4_HC_COMB:
+ ggml_cuda_op_dsv4_hc_comb(ctx, dst);
+ break;
+ case GGML_OP_DSV4_HC_PRE:
+ ggml_cuda_op_dsv4_hc_pre(ctx, dst);
+ break;
+ case GGML_OP_DSV4_HC_POST:
+ ggml_cuda_op_dsv4_hc_post(ctx, dst);
+ break;
case GGML_OP_RWKV_WKV7:
ggml_cuda_op_rwkv_wkv7(ctx, dst);
break;
case GGML_OP_FILL:
ggml_cuda_op_fill(ctx, dst);
break;
+ case GGML_OP_LIGHTNING_INDEXER:
+ ggml_cuda_lightning_indexer(ctx, dst);
+ break;
default:
return false;
}
if (backend_src != backend_dst) {
// copy on src stream
- if (cuda_ctx_src->device == cuda_ctx_dst->device) {
+ // compare the backing physical devices: distinct virtual devices may share one physical GPU,
+ // in which case a same-device copy (not a peer copy) is required
+ const int src_physical = ggml_cuda_get_physical_device(cuda_ctx_src->device);
+ const int dst_physical = ggml_cuda_get_physical_device(cuda_ctx_dst->device);
+ if (src_physical == dst_physical) {
CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream()));
} else {
#ifdef GGML_CUDA_NO_PEER_COPY
return false;
#else
- CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, cuda_ctx_dst->device, src->data, cuda_ctx_src->device, ggml_nbytes(dst), cuda_ctx_src->stream()));
+ CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_physical, src->data, src_physical, ggml_nbytes(dst), cuda_ctx_src->stream()));
#endif // GGML_CUDA_NO_PEER_COPY
}
static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
args.sigmoid = false;
+ args.sqrt_softplus = false;
args.softmax = false;
args.delayed_softmax = false;
args.prob_bias = false;
}
if (nodes[node_idx]->op == GGML_OP_UNARY) {
- if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) {
+ const ggml_unary_op unary_op = ggml_get_unary_op(nodes[node_idx]);
+ if (unary_op == GGML_UNARY_OP_SIGMOID) {
+ args.sigmoid = true;
+ } else if (unary_op == GGML_UNARY_OP_SOFTPLUS && node_idx + 1 < n_nodes &&
+ nodes[node_idx + 1]->op == GGML_OP_SQRT && nodes[node_idx + 1]->src[0] == nodes[node_idx]) {
+ // sqrt(softplus(x)) scoring (DeepSeek-V4)
+ args.sqrt_softplus = true;
+ node_idx++;
+ } else {
return false;
}
- args.sigmoid = true;
}
if (nodes[node_idx]->op == GGML_OP_ARGSORT) {
node_idx++;
- if (args.sigmoid || args.softmax) {
+ if (args.sigmoid || args.sqrt_softplus || args.softmax) {
// SOFTMAX -> RESHAPE
if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
const ggml_tensor * scale = nullptr;
if (!args.delayed_softmax) {
- ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX;
- int out_nodes[2]; // nodes which can't be elided
+ int out_nodes[2]; // nodes which can't be elided
+
+ if (args.sigmoid) {
+ ops.insert(ops.end(), { GGML_OP_UNARY });
+ } else if (args.sqrt_softplus) {
+ ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT });
+ } else {
+ ops.insert(ops.end(), { GGML_OP_SOFT_MAX });
+ }
+ const int i_probs = i + (int) ops.size() - 1; // last node of the gating activation
if (args.prob_bias) {
- bias = cgraph->nodes[i + 2]->src[1];
- ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
+ bias = cgraph->nodes[i_probs + 2]->src[1];
+ ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
GGML_OP_GET_ROWS });
- out_nodes[0] = i + 4;
- ids = cgraph->nodes[i + 4];
+ out_nodes[0] = i_probs + 4;
} else {
- ops.insert(ops.end(),
- { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
- out_nodes[0] = i + 3;
- ids = cgraph->nodes[i + 3];
+ ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
+ out_nodes[0] = i_probs + 3;
}
+ ids = cgraph->nodes[out_nodes[0]];
if (args.norm) {
ops.insert(ops.end(),
ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
if (graph->graph == nullptr) {
- if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) {
+ if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) {
if (!graph->disable_due_to_gpu_arch) {
GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__);
}
return ggml_cuda_info().device_count;
}
-void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) {
+static std::string ggml_cuda_device_description(int device) {
cudaDeviceProp prop;
- CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
- snprintf(description, description_size, "%s", prop.name);
+ CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(device)));
+
+ const ggml_cuda_device_info & info = ggml_cuda_info();
+ std::string description = prop.name;
+ if (info.device_count > info.physical_device_count) {
+ description += " (physical device " + std::to_string(info.devices[device].physical_device) +
+ ", virtual device " + std::to_string(info.devices[device].virtual_index) + ")";
+ }
+ return description;
+}
+
+void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) {
+ snprintf(description, description_size, "%s", ggml_cuda_device_description(device).c_str());
+}
+
+static int ggml_cuda_physical_device_share_count(int device) {
+ const ggml_cuda_device_info & info = ggml_cuda_info();
+ GGML_ASSERT(device >= 0 && device < info.device_count);
+ return info.devices[device].physical_share_count;
}
void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total) {
ggml_cuda_set_device(device);
CUDA_CHECK(cudaMemGetInfo(free, total));
+
+ // virtual devices sharing one physical GPU share its memory pool; split it between them
+ const int share_count = ggml_cuda_physical_device_share_count(device);
+ *free /= share_count;
+ *total /= share_count;
}
bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) {
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__)
// Check if this is a UMA (Unified Memory Architecture) system
cudaDeviceProp prop;
- CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device));
+ CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
// Check if UMA is explicitly enabled via environment variable
bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr;
}
#endif // defined(__linux__)
+ // virtual devices sharing one physical GPU share its memory pool; split it between them
+ const int share_count = ggml_cuda_physical_device_share_count(ctx->device);
+ *free /= share_count;
+ *total /= share_count;
}
static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) {
ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *) dev->context;
cudaDeviceProp prop;
- CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device));
+ CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
return prop.integrated
? GGML_BACKEND_DEVICE_TYPE_IGPU
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_BF16:
case GGML_TYPE_I32:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q8_0:
+ case GGML_TYPE_Q2_K:
+ case GGML_TYPE_Q3_K:
+ case GGML_TYPE_Q4_K:
+ case GGML_TYPE_Q5_K:
+ case GGML_TYPE_Q6_K:
+ case GGML_TYPE_IQ2_XXS:
+ case GGML_TYPE_IQ2_XS:
+ case GGML_TYPE_IQ2_S:
+ case GGML_TYPE_IQ3_XXS:
+ case GGML_TYPE_IQ3_S:
+ case GGML_TYPE_IQ1_S:
+ case GGML_TYPE_IQ1_M:
+ case GGML_TYPE_IQ4_XS:
return true;
+ case GGML_TYPE_IQ4_NL:
+ case GGML_TYPE_MXFP4:
+ // 32-value sub-blocks, the row size does not guarantee
+ // the QK_K super-blocks the get_rows kernel iterates on
+ return op->src[0]->ne[0] % QK_K == 0;
default:
return false;
}
{
ggml_type src0_type = op->src[0]->type;
ggml_type src1_type = op->src[1]->type;
+ const int32_t dim = op->op_params[0];
return src0_type == src1_type &&
src0_type == op->type &&
(
(
ggml_is_quantized(src0_type) &&
- ggml_is_contiguous(op->src[0]) &&
- ggml_is_contiguous(op->src[1]) &&
+ (
+ (
+ dim == 3 &&
+ ggml_is_contiguous(op->src[0]) &&
+ ggml_is_contiguous(op->src[1])
+ ) || (
+ dim != 3 &&
+ ggml_is_contiguous_to_3(op->src[0]) &&
+ ggml_is_contiguous_to_3(op->src[1])
+ )
+ ) &&
op->src[0]->ne[0] % ggml_blck_size(src0_type) == 0 &&
op->src[1]->ne[0] % ggml_blck_size(src0_type) == 0
) || (
#else
return true;
#endif // GGML_USE_MUSA
+ case GGML_OP_DSV4_HC_COMB:
+ return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+ op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
+ case GGML_OP_DSV4_HC_PRE:
+ return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+ op->type == GGML_TYPE_F32;
+ case GGML_OP_DSV4_HC_POST:
+ return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+ op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 &&
+ op->type == GGML_TYPE_F32;
case GGML_OP_FLASH_ATTN_EXT:
return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op);
case GGML_OP_CROSS_ENTROPY_LOSS:
case GGML_OP_DIAG:
case GGML_OP_SOLVE_TRI:
return true;
+ case GGML_OP_LIGHTNING_INDEXER:
+ return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
default:
return false;
ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context;
const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;
- for (int i = 0; i < ggml_cuda_info().device_count; i++) {
+ const ggml_cuda_device_info & info = ggml_cuda_info();
+ const bool virtual_devices = info.device_count > info.physical_device_count;
+
+ for (int i = 0; i < info.device_count; i++) {
+ const int physical_id = info.devices[i].physical_device;
+
ggml_backend_cuda_device_context * dev_ctx = new ggml_backend_cuda_device_context;
dev_ctx->device = i;
dev_ctx->name = GGML_CUDA_NAME + std::to_string(i);
-
- cudaDeviceProp prop;
- CUDA_CHECK(cudaGetDeviceProperties(&prop, i));
- dev_ctx->description = prop.name;
+ dev_ctx->description = ggml_cuda_device_description(i);
char pci_bus_id[32] = {};
- CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), i));
+ CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), physical_id));
dev_ctx->pci_bus_id = pci_bus_id;
+ if (virtual_devices) {
+ // make the pci bus id unique for virtual devices
+ dev_ctx->pci_bus_id += "-v" + std::to_string(i);
+ }
for (char & c : dev_ctx->pci_bus_id) {
c = std::tolower(c);
}
--- /dev/null
+#include "common.cuh"
+#include "lightning-indexer.cuh"
+#include "fattn-common.cuh"
+#include "convert.cuh"
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if defined(TURING_MMA_AVAILABLE)
+
+typedef union {
+ int2 i2;
+ half2 h2[2];
+} half4;
+
+// TODO add support for AMD cards via rocWMMA
+#include <mma.h>
+namespace wmma = nvcuda::wmma;
+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_wmma(
+ const float * Q, const char * K, const float * W, const half * M, float * dst,
+ int64_t n_stream, int64_t n_batch, int64_t n_kv,
+ size_t nb1, size_t nb2, size_t nb3,
+ size_t nbq1, size_t nbq2, size_t nbq3,
+ size_t nbk1, size_t nbk2, size_t nbk3,
+ size_t nbw1, size_t nbw2, size_t nbw3,
+ size_t nbm1, size_t nbm2, size_t nbm3,
+ int64_t nem3
+ ) {
+
+ constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE;
+ constexpr int HEADS_PER_INNER_LOOP = 8;
+ constexpr int K_EMBD_PER_INNER_LOOP = 16;
+ constexpr int N_EMBD_PADDED = N_EMBD + 8;
+
+ const int i_batch = blockIdx.y;
+ const int i_stream = blockIdx.z;
+ const int i_warp = threadIdx.y;
+ const int i_lane = threadIdx.x;
+ const int tid = i_warp * WARP_SIZE + i_lane;
+
+ // each block processes K_VECS_PER_BLOCK K vectors
+ const int start_kv = blockIdx.x * K_VECS_PER_BLOCK;
+
+ const char * q_base = (const char *) Q + i_batch*nbq2 + i_stream*nbq3;
+ const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3);
+
+ // phase 1 - load weights and first Q tile to shared memory
+
+ __shared__ float w_shared[N_HEAD];
+ __shared__ int2 q_shared_h[HEADS_PER_INNER_LOOP][N_EMBD_PADDED / 4];
+
+ if (tid < N_HEAD) {
+ w_shared[tid] = w_base[tid];
+ }
+
+ // total number of half4 elements in HEADS_PER_INNER_LOOP x N_EMBD Q tile
+ constexpr int N_Q_TILE = HEADS_PER_INNER_LOOP * (N_EMBD / 4);
+ // number of registers needed in each thread to store Q tile in thread block
+ constexpr int N_Q_NEXT = (N_Q_TILE + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+
+#pragma unroll
+ for (int i_q = tid; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
+ const int i_head = i_q / (N_EMBD / 4);
+ const int i_embd = i_q % (N_EMBD / 4);
+ const float4 q = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
+ half4 q_packed;
+ q_packed.h2[0] = __float22half2_rn(make_float2(q.x, q.y));
+ q_packed.h2[1] = __float22half2_rn(make_float2(q.z, q.w));
+ q_shared_h[i_head][i_embd] = q_packed.i2;
+ }
+
+ // phase 2 - load (and dequantize if needed) K to shared mem
+
+ __shared__ half2 k_shared_h[K_VECS_PER_BLOCK][N_EMBD_PADDED / 4][2];
+
+ constexpr int n_k = K_VECS_PER_BLOCK * (N_EMBD / 4);
+
+ if constexpr (TYPE_K == GGML_TYPE_F16) {
+#pragma unroll
+ for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) {
+ const int i_k_vec = i_k / (N_EMBD / 4);
+ const int i_embd = i_k % (N_EMBD / 4);
+ const int i_kv = start_kv + i_k_vec;
+ if (i_kv < n_kv) {
+ const int2 * k_base = (const int2 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+ *(int2*) &k_shared_h[i_k_vec][i_embd] = k_base[i_embd];
+ } else {
+ *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0);
+ }
+ }
+ } else {
+ constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, half, 4>();
+#pragma unroll
+ for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) {
+ const int i_k_vec = i_k / (N_EMBD / 4);
+ const int i_embd = i_k % (N_EMBD / 4);
+ const int i_kv = start_kv + i_k_vec;
+ if (i_kv < n_kv) {
+ const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+ dequantize_k(k_base, &k_shared_h[i_k_vec][i_embd][0], i_embd * 4);
+ } else {
+ *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0);
+ }
+ }
+ }
+
+ __syncthreads();
+
+ // phase 3 - calculate lightning indexer scores
+
+ __shared__ float qk_shared[WARPS_PER_BLOCK][HEADS_PER_INNER_LOOP][K_VECS_PER_BLOCK];
+
+ // load K fragment
+ wmma::fragment<wmma::matrix_b, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::col_major> frag_k;
+ wmma::load_matrix_sync(frag_k, (half*) &k_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED);
+
+ float score_k = 0.0f;
+
+ for (int i_head_0 = 0; i_head_0 < N_HEAD; i_head_0 += HEADS_PER_INNER_LOOP) {
+ const int i_head_next = i_head_0 + HEADS_PER_INNER_LOOP;
+
+ // we don't use accumulator for anything, fill it with zeros
+ wmma::fragment<wmma::accumulator, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, float> frag_acc;
+ wmma::fill_fragment(frag_acc, 0.0f);
+
+ // load Q fragment
+ wmma::fragment<wmma::matrix_a, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::row_major> frag_q;
+ wmma::load_matrix_sync(frag_q, (half*) &q_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED);
+
+ // preload next Q tile to registers during matrix multiplication
+ float4 q_next[N_Q_NEXT];
+
+ if (i_head_next < N_HEAD) {
+#pragma unroll
+ for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
+ const int i_head = i_head_next + i_q / (N_EMBD / 4);
+ const int i_embd = i_q % (N_EMBD / 4);
+ q_next[i_q_next++] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
+ }
+ }
+
+ // perform matrix multiplication
+ wmma::mma_sync(frag_acc, frag_q, frag_k, frag_acc);
+ wmma::store_matrix_sync((float*) &qk_shared[i_warp][0][0], frag_acc, K_VECS_PER_BLOCK, wmma::mem_row_major);
+
+ // make sure all threads finished using q_shared_h so we can store next tile
+ __syncthreads();
+
+ // write preloaded Q tile to shared memory
+ if (i_head_next < N_HEAD) {
+#pragma unroll
+ for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
+ const int i_head = i_q / (N_EMBD / 4);
+ const int i_embd = i_q % (N_EMBD / 4);
+ half4 q_packed;
+ q_packed.h2[0] = __float22half2_rn(make_float2(q_next[i_q_next].x, q_next[i_q_next].y));
+ q_packed.h2[1] = __float22half2_rn(make_float2(q_next[i_q_next].z, q_next[i_q_next].w));
+ q_shared_h[i_head][i_embd] = q_packed.i2;
+ ++i_q_next;
+ }
+ }
+
+ // accumulate QK multiplication results from all block warps
+ // (there are 256 threads in block and 256 matmul outputs)
+ // TODO it will break if WARP_SIZE is not 32
+ const int h = tid / K_VECS_PER_BLOCK;
+ const int k = tid % K_VECS_PER_BLOCK;
+ const float w_val = w_shared[i_head_0 + h];
+
+ float sum = 0.0f;
+#pragma unroll
+ for (int w = 0; w < WARPS_PER_BLOCK; ++w) {
+ sum += qk_shared[w][h][k];
+ }
+
+ // ReLU, weight
+ sum = sum > 0.0f ? sum : 0.0f;
+ sum *= w_val;
+
+ // wait until qk_shared[0] is no longer used
+ __syncthreads();
+
+ // reuse qk_shared[0] for storing partial results
+ qk_shared[0][h][k] = sum;
+
+ // wait until all threads write their results
+ __syncthreads();
+
+ // accumulate result over heads
+ if (tid < K_VECS_PER_BLOCK) {
+#pragma unroll
+ for (int i_head = 0; i_head < HEADS_PER_INNER_LOOP; ++i_head) {
+ score_k += qk_shared[0][i_head][tid];
+ }
+ }
+
+ // make sure all threads finished using qk_shared
+ __syncthreads();
+ }
+
+ // phase 4 - store output to VRAM
+
+ if (tid < K_VECS_PER_BLOCK) {
+ const int i_kv = start_kv + tid;
+ if (i_kv < n_kv) {
+ const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
+ float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
+ dst_base[i_kv] = score_k + __half2float(m_base[i_kv]);
+ }
+ }
+}
+
+#else // defined(TURING_MMA_AVAILABLE)
+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_wmma(
+ const float * Q, const char * K, const float * W, const half * M, float * dst,
+ int64_t n_stream, int64_t n_batch, int64_t n_kv,
+ size_t nb1, size_t nb2, size_t nb3,
+ size_t nbq1, size_t nbq2, size_t nbq3,
+ size_t nbk1, size_t nbk2, size_t nbk3,
+ size_t nbw1, size_t nbw2, size_t nbw3,
+ size_t nbm1, size_t nbm2, size_t nbm3,
+ int64_t nem3
+ ) {
+ GGML_UNUSED_VARS(Q, K, W, M, dst,
+ n_stream, n_batch, n_kv,
+ nb1, nb2, nb3,
+ nbq1, nbq2, nbq3,
+ nbk1, nbk2, nbk3,
+ nbw1, nbw2, nbw3,
+ nem3);
+ NO_DEVICE_CODE;
+}
+
+#endif // defined(TURING_MMA_AVAILABLE)
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
+// TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32
+// thanks to that one warp operating on float4 processes whole indexer K/Q vectors
+// 32 * 4 = 128 (N_EMBD)
+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_vec(
+ const float * Q, const char * K, const float * W, const half * M, float * dst,
+ int64_t n_stream, int64_t n_batch, int64_t n_kv,
+ size_t nb1, size_t nb2, size_t nb3,
+ size_t nbq1, size_t nbq2, size_t nbq3,
+ size_t nbk1, size_t nbk2, size_t nbk3,
+ size_t nbw1, size_t nbw2, size_t nbw3,
+ size_t nbm1, size_t nbm2, size_t nbm3,
+ int64_t nem3
+ ) {
+
+ constexpr int K_VECS_PER_WARP = K_VECS_PER_BLOCK / WARPS_PER_BLOCK;
+ constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE;
+
+ const int i_batch = blockIdx.y;
+ const int i_stream = blockIdx.z;
+ const int i_warp = threadIdx.y;
+ const int i_lane = threadIdx.x;
+ const int tid = i_warp * WARP_SIZE + i_lane;
+
+ // each warp processes K_VECS_PER_WARP K vectors
+ const int start_kv_block = blockIdx.x * K_VECS_PER_BLOCK;
+ const int start_kv = start_kv_block + i_warp * K_VECS_PER_WARP;
+
+ const char * q_base = (const char *) Q + i_batch*nbq2 + i_stream*nbq3;
+ const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3);
+
+ // phase 1 - load (and dequantize if needed) K to registers
+
+ float4 k_reg_f[K_VECS_PER_WARP];
+
+ if constexpr (TYPE_K == GGML_TYPE_F32) {
+ // direct copy of float4
+#pragma unroll
+ for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+ int i_kv = start_kv + k;
+ if (i_kv < n_kv) {
+ const float4 * k_base = (const float4 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+ k_reg_f[k] = k_base[i_lane];
+ } else {
+ k_reg_f[k] = make_float4(0, 0, 0, 0);
+ }
+ }
+ } else {
+ // dequantize remaining types to float
+ constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, float, 4>();
+#pragma unroll
+ for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+ int i_kv = start_kv + k;
+ if (i_kv < n_kv) {
+ const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+ dequantize_k(k_base, &k_reg_f[k], i_lane * 4);
+ } else {
+ k_reg_f[k] = make_float4(0, 0, 0, 0);
+ }
+ }
+ }
+
+ float score_k[K_VECS_PER_WARP] = { 0.0f };
+
+ // load weights and Q only for N_HEAD_INNER heads at once to reduce shared memory usage
+ constexpr int N_HEAD_INNER = N_HEAD / 4;
+
+ for (int i_head_0 = 0; i_head_0 < N_HEAD; i_head_0 += N_HEAD_INNER) {
+ // phase 2 - load weights and Q to shared memory
+
+ __shared__ float w_shared[N_HEAD_INNER];
+ __shared__ float4 q_shared_f[N_HEAD_INNER][N_EMBD / 4];
+
+ if (tid < N_HEAD_INNER) {
+ w_shared[tid] = w_base[i_head_0 + tid];
+ }
+
+ constexpr int n_q = N_HEAD_INNER * (N_EMBD / 4);
+#pragma unroll
+ for (int i_q = tid; i_q < n_q; i_q += THREADS_PER_BLOCK) {
+ const int i_head_inner = i_q / (N_EMBD / 4);
+ const int i_head = i_head_0 + i_head_inner;
+ const int i_embd = i_q % (N_EMBD / 4);
+ q_shared_f[i_head_inner][i_embd] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
+ }
+
+ __syncthreads();
+
+ // phase 3 - calculate lightning indexer scores
+
+ for (int i_head_inner = 0; i_head_inner < N_HEAD_INNER; ++i_head_inner) {
+ const float w_val = w_shared[i_head_inner];
+ float qk[K_VECS_PER_WARP] = { 0.0f };
+
+ // dot product of floats
+ const float4 q_vec = q_shared_f[i_head_inner][i_lane];
+
+#pragma unroll
+ for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+ ggml_cuda_mad(qk[k], q_vec.x, k_reg_f[k].x);
+ ggml_cuda_mad(qk[k], q_vec.y, k_reg_f[k].y);
+ ggml_cuda_mad(qk[k], q_vec.z, k_reg_f[k].z);
+ ggml_cuda_mad(qk[k], q_vec.w, k_reg_f[k].w);
+ }
+
+#pragma unroll
+ for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+ float sum = warp_reduce_sum(qk[k]);
+
+ // ReLU, weight
+ if (i_lane == 0) {
+ sum = (sum > 0.0f) ? sum : 0.0f;
+ score_k[k] += sum * w_val;
+ }
+ }
+ }
+
+ __syncthreads();
+ }
+
+ // phase 4 - store outputs to shared memory
+
+ __shared__ float dst_shared[K_VECS_PER_BLOCK];
+
+ if (i_lane == 0) {
+#pragma unroll
+ for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+ dst_shared[i_warp * K_VECS_PER_WARP + k] = score_k[k];
+ }
+ }
+
+ __syncthreads();
+
+ // phase 5 - write from shared memory to VRAM in coalesced manner
+
+ if (tid < K_VECS_PER_BLOCK) {
+ int i_kv = start_kv_block + tid;
+ if (i_kv < n_kv) {
+ const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
+ float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
+ dst_base[i_kv] = dst_shared[tid] + __half2float(m_base[i_kv]);
+ }
+ }
+}
+
+#define LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, K, type_K) \
+ if (K->type == (type_K)) { \
+ lightning_indexer_kernel<WARPS_PER_BLOCK, K_VECS_PER_BLOCK, n_embd, n_head, type_K> \
+ <<<grid, block, 0, ctx.stream()>>>( \
+ q_d, k_d, w_d, m_d, dst_d, \
+ n_stream, n_batch, n_kv, \
+ nb1, nb2, nb3, \
+ nbq1, nbq2, nbq3, \
+ nbk1, nbk2, nbk3, \
+ nbw1, nbw2, nbw3, \
+ nbm1, nbm2, nbm3, \
+ nem3 \
+ ); \
+ } else
+
+void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, 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)
+
+ // input tensor rows must be contiguous
+ 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));
+
+ // dst cannot be transposed or permuted
+ GGML_ASSERT(nb0 == sizeof(float));
+ GGML_ASSERT(nb0 <= nb1);
+ GGML_ASSERT(nb1 <= nb2);
+ GGML_ASSERT(nb2 <= nb3);
+
+ const int n_embd = q->ne[0];
+ const int n_head = q->ne[1];
+ const int n_batch = q->ne[2];
+ const int n_stream = q->ne[3];
+ const int n_kv = k->ne[2];
+
+ const float * q_d = (const float *) q->data;
+ const char * k_d = (const char *) k->data;
+ const float * w_d = (const float *) w->data;
+ const half * m_d = (const half *) m->data;
+ float * dst_d = ( float *) dst->data;
+
+ const int device = ggml_cuda_get_device();
+ const int cc = ggml_cuda_info().devices[device].cc;
+
+ if (n_embd == 128 && n_head == 64) {
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
+ // use wmma kernel
+ constexpr int K_VECS_PER_BLOCK = 32;
+ constexpr int WARPS_PER_BLOCK = 8;
+
+ dim3 block(32, WARPS_PER_BLOCK);
+ int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+ dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_F16)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q4_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q4_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q5_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q5_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q8_0)
+ GGML_ABORT("fatal error");
+ } else {
+#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ {
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ // use vector kernel
+ constexpr int K_VECS_PER_WARP = 8;
+ constexpr int WARPS_PER_BLOCK = 8;
+ constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;
+
+ dim3 block(32, WARPS_PER_BLOCK);
+ int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+ dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_F16)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q4_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q4_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q5_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q5_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q8_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_BF16)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_F32)
+ GGML_ABORT("fatal error");
+ }
+ } else if (n_embd == 128 && n_head == 32) {
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
+ // use wmma kernel
+ constexpr int K_VECS_PER_BLOCK = 32;
+ constexpr int WARPS_PER_BLOCK = 8;
+
+ dim3 block(32, WARPS_PER_BLOCK);
+ int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+ dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_F16)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q8_0)
+ GGML_ABORT("fatal error");
+ } else {
+#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ {
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ // use vector kernel
+ constexpr int K_VECS_PER_WARP = 8;
+ constexpr int WARPS_PER_BLOCK = 8;
+ constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;
+
+ dim3 block(32, WARPS_PER_BLOCK);
+ int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+ dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F16)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q4_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q4_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q5_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q5_1)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q8_0)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_BF16)
+ LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32)
+ GGML_ABORT("fatal error");
+ }
+ } else {
+ GGML_ABORT("fatal error");
+ }
+}
+
+bool ggml_cuda_lightning_indexer_supported(int device, const ggml_tensor * dst) {
+ GGML_UNUSED(device);
+
+ 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_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)
+
+ if (neq0 != 128) {
+ return false;
+ }
+
+ if (neq1 != 64 && neq1 != 32) {
+ return false;
+ }
+
+ // alignment checks
+ for (const ggml_tensor * t : {q, k}) {
+ if (ggml_is_quantized(t->type)) {
+ continue;
+ }
+ for (size_t i = 1; i < GGML_MAX_DIMS; ++i) {
+ if (t->nb[i] % 16 != 0) {
+ return false;
+ }
+ }
+ }
+
+ switch(k->type) {
+ case GGML_TYPE_F32:
+ case GGML_TYPE_BF16:
+ case GGML_TYPE_F16:
+ case GGML_TYPE_Q8_0:
+ case GGML_TYPE_Q5_1:
+ case GGML_TYPE_Q5_0:
+ case GGML_TYPE_Q4_1:
+ case GGML_TYPE_Q4_0:
+ return true;
+ default:
+ return false;
+ }
+}
--- /dev/null
+#include "common.cuh"
+
+void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+bool ggml_cuda_lightning_indexer_supported(int device, const ggml_tensor * dst);
GGML_ASSERT(sis1 > 0);
ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), expert_bounds_dev.get(),
- static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, ctx.stream());
+ static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, /*write_inverse =*/ false, ctx.stream());
CUDA_CHECK(cudaGetLastError());
ids_info.ids_src_compact = ids_src_compact_dev.get();
__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
- const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
+ const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
const mm_ids_helper_store store_it = store[itc];
const int it = store_it.it();
const int iex_used = store_it.iex_used();
- ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
- ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
+ ids_dst[nex_prev + itc] = it*n_expert_used + iex_used;
+ // ids_src1 holds the forward map, or the inverse map (token slot -> compact row) for quant dedup
+ if (write_inverse) {
+ ids_src1[it*n_expert_used + iex_used] = nex_prev + itc;
+ } else {
+ ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
+ }
}
if (threadIdx.x != 0) {
template <int n_expert_used_template>
static void launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
- const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
+ const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
GGML_ASSERT(n_tokens < (1 << 22) && "too few bits in mm_ids_helper_store");
GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
GGML_ASSERT(nbytes_shared <= smpbo);
mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
- (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
+ (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
}
void ggml_cuda_launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
- const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
+ const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
switch (n_expert_used) {
case 2:
- launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 4:
- launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 6:
- launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 8:
- launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 16:
- launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
case 32:
- launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
default:
- launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+ launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
break;
}
}
void ggml_cuda_launch_mm_ids_helper(
const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
- int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, cudaStream_t stream);
+ int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_ampere(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_blackwell(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+ CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, false);
+
+ return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_cdna(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q1_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_1, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q8_0, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q2_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q3_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q4_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q5_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_Q6_K, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ1_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ2_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_XXS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ3_S, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_XS, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_IQ4_NL, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_MXFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false);
+
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true);
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+ CASE(GGML_TYPE_NVFP4, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna2(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3_5(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna4(ggml_type type, int J, bool fallback) {
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+ CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+ return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
--- /dev/null
+#pragma once
+
+#include "vecdotq.cuh"
+
+#include "mmq.cuh"
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q1_0(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int blocks_per_iter = MMQ_ITER_K / QK1_0;
+ constexpr int threads_per_row = blocks_per_iter * QI1_0;
+ constexpr int nrows = warp_size / threads_per_row;
+ constexpr int scale_entries_per_block = QK1_0 / QK8_1;
+ constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block;
+
+ const int txi = threadIdx.x % threads_per_row;
+ const int kbx = txi / QI1_0;
+ const int kqsx = txi % QI1_0;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + kbx;
+ const int16_t * qxi = (const int16_t *) bxi->qs + kqsx * 2;
+
+ const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
+#pragma unroll
+ for (int j = 0; j < 2; ++j) {
+ const int q = qxi[j];
+
+ // unpack crumbs into nibble indices
+ const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8, 9, 12, 13]
+ const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15]
+ // unpack nibbles into byte values
+ const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >> 0);
+ const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >> 0);
+ const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16);
+ const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16);
+ // unshuffle values
+ const int v0 = __byte_perm(s0, s1, 0x5410);
+ const int v1 = __byte_perm(s0, s1, 0x7632);
+ const int v2 = __byte_perm(s2, s3, 0x5410);
+ const int v3 = __byte_perm(s2, s3, 0x7632);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + dst_offset + j*4+0] = v0;
+ x_qs[i*sram_stride + dst_offset + j*4+1] = v1;
+ x_qs[i*sram_stride + dst_offset + j*4+2] = v2;
+ x_qs[i*sram_stride + dst_offset + j*4+3] = v3;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+0] = v0;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+1] = v1;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+2] = v2;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+3] = v3;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+ }
+
+ const int ksx = threadIdx.x % scale_entries_per_row;
+ const int scale_block = ksx / scale_entries_per_block;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps) {
+ int i = i0 + threadIdx.y;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + scale_block;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + ksx] = bxi->d;
+#else
+ x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_0(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int blocks_per_iter = MMQ_ITER_K / QK2_0;
+ constexpr int threads_per_row = blocks_per_iter * QI2_0;
+ constexpr int nrows = warp_size / threads_per_row;
+ constexpr int scale_entries_per_block = QK2_0 / QK8_1;
+ constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block;
+
+ const int txi = threadIdx.x % threads_per_row;
+ const int kbx = txi / QI2_0;
+ const int kqsx = txi % QI2_0;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + kbx;
+ // Each 32-element chunk occupies 8 bytes of qs (32 elements * 2 bits = 64 bits)
+ const int qs_offset = 8*kqsx;
+ const int qs0 = bxi->qs[qs_offset + 0] | (bxi->qs[qs_offset + 1] << 8) |
+ (bxi->qs[qs_offset + 2] << 16) | (bxi->qs[qs_offset + 3] << 24);
+ const int qs1 = bxi->qs[qs_offset + 4] | (bxi->qs[qs_offset + 5] << 8) |
+ (bxi->qs[qs_offset + 6] << 16) | (bxi->qs[qs_offset + 7] << 24);
+
+ // Unpack 32 2-bit codes into 8 int32s, each holding 4 signed int8s in {-1,0,1,2}.
+ int unpacked_bytes[8];
+#pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ const int shift = j * 8;
+ const int codes = (qs0 >> shift) & 0xFF;
+ const int c0 = ((codes >> 0) & 0x3) - 1;
+ const int c1 = ((codes >> 2) & 0x3) - 1;
+ const int c2 = ((codes >> 4) & 0x3) - 1;
+ const int c3 = ((codes >> 6) & 0x3) - 1;
+ unpacked_bytes[j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
+ }
+#pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ const int shift = j * 8;
+ const int codes = (qs1 >> shift) & 0xFF;
+ const int c0 = ((codes >> 0) & 0x3) - 1;
+ const int c1 = ((codes >> 2) & 0x3) - 1;
+ const int c2 = ((codes >> 4) & 0x3) - 1;
+ const int c3 = ((codes >> 6) & 0x3) - 1;
+ unpacked_bytes[4 + j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
+ }
+
+ const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
+#pragma unroll
+ for (int j = 0; j < 8; ++j) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + dst_offset + j] = unpacked_bytes[j];
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j] = unpacked_bytes[j];
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+ }
+
+ const int ksx = threadIdx.x % scale_entries_per_row;
+ const int scale_block = ksx / scale_entries_per_block;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps) {
+ int i = i0 + threadIdx.y;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + scale_block;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + ksx] = bxi->d;
+#else
+ x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_0);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI4_0;
+ const int kqsx = txi % QI4_0;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbx;
+ const int qs0 = get_int_b2(bxi->qs, kqsx);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + 0] = __vsubss4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808);
+ x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsubss4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808);
+#else
+ x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kbxd] = bxi->d;
+#else
+ x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + kbxd] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I);
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_1);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI4_1;
+ const int kqsx = txi % QI4_1;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbx;
+ const int qs0 = get_int_b4(bxi->qs, kqsx);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kbx*(2*QI4_1) + kqsx + 0] = (qs0 >> 0) & 0x0F0F0F0F;
+ x_qs[i*sram_stride + kbx*(2*QI4_1) + kqsx + QI4_1] = (qs0 >> 4) & 0x0F0F0F0F;
+#else
+ x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_1;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_dm[i*sram_stride + kbxd] = bxi->dm;
+#else
+ x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + kbxd] = bxi->dm;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_0(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_0, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_0);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI5_0;
+ const int kqsx = txi % QI5_0;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbx;
+
+ const int ql = get_int_b2(bxi->qs, kqsx);
+ const int qh = get_int_b2(bxi->qh, 0) >> (4 * kqsx);
+
+ int qs0 = (ql >> 0) & 0x0F0F0F0F;
+ qs0 |= (qh << 4) & 0x00000010; // 0 -> 4
+ qs0 |= (qh << 11) & 0x00001000; // 1 -> 12
+ qs0 |= (qh << 18) & 0x00100000; // 2 -> 20
+ qs0 |= (qh << 25) & 0x10000000; // 3 -> 28
+ qs0 = __vsubss4(qs0, 0x10101010); // subtract 16
+
+ int qs1 = (ql >> 4) & 0x0F0F0F0F;
+ qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4
+ qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12
+ qs1 |= (qh << 2) & 0x00100000; // 18 -> 20
+ qs1 |= (qh << 9) & 0x10000000; // 19 -> 28
+ qs1 = __vsubss4(qs1, 0x10101010); // subtract 16
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + 0] = qs0;
+ x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + 0] = qs0;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_0;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kbxd] = bxi->d;
+#else
+ x_df[i*(MMQ_TILE_NE_K/QI5_0) + i/QI5_0 + kbxd] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_1(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I);
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_1);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI5_1;
+ const int kqsx = txi % QI5_1;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbx;
+
+ const int ql = get_int_b4(bxi->qs, kqsx);
+ const int qh = get_int_b4(bxi->qh, 0) >> (4 * kqsx);
+
+ int qs0 = (ql >> 0) & 0x0F0F0F0F;
+ qs0 |= (qh << 4) & 0x00000010; // 0 -> 4
+ qs0 |= (qh << 11) & 0x00001000; // 1 -> 12
+ qs0 |= (qh << 18) & 0x00100000; // 2 -> 20
+ qs0 |= (qh << 25) & 0x10000000; // 3 -> 28
+
+ int qs1 = (ql >> 4) & 0x0F0F0F0F;
+ qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4
+ qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12
+ qs1 |= (qh << 2) & 0x00100000; // 18 -> 20
+ qs1 |= (qh << 9) & 0x10000000; // 19 -> 28
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kbx*(2*QI5_1) + kqsx + 0] = qs0;
+ x_qs[i*sram_stride + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + 0] = qs0;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_1;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_dm[i*sram_stride + kbxd] = bxi->dm;
+#else
+ x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + kbxd] = bxi->dm;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ // MMQ_ITER_K / (4 * QR8_0) == 64 required. but NV has only 32 threads per warp
+ constexpr int threads_per_row = 32;
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI8_0;
+ const int kqsx = txi % QI8_0;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 0 + txi] = get_int_b2(bxi[0].qs, kqsx);
+ x_qs[i*sram_stride + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx);
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 0 + txi] = get_int_b2(bxi[0].qs, kqsx);
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kbxd] = bxi->d;
+#else
+ x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I);
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR2_K);
+ constexpr int nrows = ggml_cuda_get_physical_warp_size() / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q2_K * bxi = (const block_q2_K *) x + kbx0 + i*stride;
+
+ const int x_ql_0 = get_int_b2(bxi->qs, kqsx);
+
+#pragma unroll
+ for (int l = 0; l < QR2_K; ++l) {
+ const int k = (kqsx/8)*32 + l*8 + kqsx % 8;
+
+ const int x_qs_k = (x_ql_0 >> (2*l)) & 0x03030303;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + k] = x_qs_k;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const int sc_m = bxi->scales[kqsx];
+#ifdef FAST_FP16_AVAILABLE
+ const half2 x_dm_ik = __hmul2(bxi->dm, make_half2(sc_m & 0x0F, sc_m >> 4));
+#else
+ const float2 bxi_dmf = __half22float2(bxi->dm);
+ const half2 x_dm_ik = make_half2(bxi_dmf.x*(sc_m & 0x0F), bxi_dmf.y*(sc_m >> 4));
+#endif // FAST_FP16_AVAILABLE
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_dm[i*sram_stride + kqsx] = x_dm_ik;
+#else
+ x_dm[i*(MMQ_TILE_NE_K + 1) + kqsx] = x_dm_ik;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+ int * x_sc = (int *) (x_df + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR3_K);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
+
+ const int x_ql_0 = get_int_b2(bxi->qs, kqsx);
+ const int x_qh_0 = get_int_b2(bxi->hmask, kqsx % (QI3_K/2)) >> (4 * (kqsx / (QI3_K/2)));
+
+#pragma unroll
+ for (int l = 0; l < QR3_K; ++l) {
+ const int k = (kqsx/8)*32 + l*8 + kqsx % 8;
+
+ const int x_ql_k = (x_ql_0 >> (2*l)) & 0x03030303;
+ const int x_qh_k = ((x_qh_0 >> l) << 2) & 0x04040404;
+
+ const int x_qs_k = __vsubss4(x_ql_k | x_qh_k, 0x04040404);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + k] = x_qs_k;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+ }
+
+ constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+ int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/4;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
+
+ const int ksc = threadIdx.x % 4;
+
+ const int ksc_low = ksc % (QI3_K/8);
+ const int shift_low = 4 * (ksc / (QI3_K/8));
+ const int sc_low = (get_int_b2(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F;
+
+ const int ksc_high = QI3_K/8;
+ const int shift_high = 2 * ksc;
+ const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030;
+
+ const int sc = __vsubss4(sc_low | sc_high, 0x20202020);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ const int8_t * sc8 = (const int8_t *) ≻
+ const float d = bxi->d;
+
+#pragma unroll
+ for (int l = 0; l < int(sizeof(int)); ++l) {
+ x_df[i*sram_stride + sizeof(int)*ksc + l] = d*sc8[l];
+ }
+#else
+ x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = sc;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+#if !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE))
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+ int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
+
+ x_df[i] = bxi->d;
+ }
+#endif // !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)) || defined(AMD_WMMA_AVAILABLE)
+}
+
+static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, const int ksc) {
+ // scale arrangement after the following two lines:
+ // - ksc == 0: sc0, sc1, sc2, sc3
+ // - ksc == 1: sc4, sc5, sc6, sc7
+ // - ksc == 2: m0, m1, m2, m3
+ // - ksc == 3: m4, m5, m6, m7
+ return ((scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F) | // lower 4 bits
+ ((scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030); // upper 2 bits
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I);
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + txs.qs);
+ int * x_sc = (int *) (x_dm + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_K);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+ const int qs0 = get_int_b4(bxi->qs, txi);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 16*(txi/8) + txi % 8 + 0] = (qs0 >> 0) & 0x0F0F0F0F;
+ x_qs[i*sram_stride + 16*(txi/8) + txi % 8 + 8] = (qs0 >> 4) & 0x0F0F0F0F;
+#else
+ x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr int rows_per_warp = warp_size / 2;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ // Need if on AMD instead of % because warp_size == 64
+ // This causes double work and throughput loss (MI300X)
+ // H100 loses about 100 t/s with 'if' condition over '%'
+ int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2;
+ if (i < I) {
+#else
+ int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % I;
+ {
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+
+ const int * scales = (const int *) bxi->scales;
+ const int ksc = threadIdx.x % 2;
+
+ const int sc32 = unpack_scales_q45_K(scales, ksc + 0);
+ const int m32 = unpack_scales_q45_K(scales, ksc + 2);
+
+ const uint8_t * sc8 = (const uint8_t *) &sc32;
+ const uint8_t * m8 = (const uint8_t *) &m32;
+
+ const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
+
+ #pragma unroll
+ for (int l = 0; l < sizeof(int); ++l) {
+ x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
+ }
+ }
+ }
+#else
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+ int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+
+ x_dm[i] = bxi->dm;
+ }
+ constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+ int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
+
+ const int * scales = (const int *) bxi->scales;
+
+ const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8);
+ const int scales8 = unpack_scales_q45_K(scales, ksc);
+
+ x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8;
+ }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I);
+ int * x_qs = (int *) x_tile;
+ half2 * x_dm = (half2 *) (x_qs + txs.qs);
+ int * x_sc = (int *) (x_dm + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_K);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+ const int ky = QR5_K*txi;
+
+ const int ql = get_int_b4(bxi->qs, txi);
+ const int ql0 = (ql >> 0) & 0x0F0F0F0F;
+ const int ql1 = (ql >> 4) & 0x0F0F0F0F;
+
+ const int qh = get_int_b4(bxi->qh, txi % (QI5_K/4));
+ const int qh0 = ((qh >> (2 * (txi / (QI5_K/4)) + 0)) << 4) & 0x10101010;
+ const int qh1 = ((qh >> (2 * (txi / (QI5_K/4)) + 1)) << 4) & 0x10101010;
+
+ const int kq0 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + 0;
+ const int kq1 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + QI5_K/4;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kq0] = ql0 | qh0;
+ x_qs[i*sram_stride + kq1] = ql1 | qh1;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = ql0 | qh0;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = ql1 | qh1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr int rows_per_warp = warp_size / 2;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+#if defined(AMD_MFMA_AVAILABLE)
+ // Need if on AMD instead of % because warp_size == 64
+ // This causes double work and throughput loss (MI300X)
+ // H100 loses about 100 t/s with 'if' condition over '%'
+ int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2;
+ if (i < I) {
+#else
+ int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % I;
+ {
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+
+ const int * scales = (const int *) bxi->scales;
+ const int ksc = threadIdx.x % 2;
+
+ const int sc32 = unpack_scales_q45_K(scales, ksc + 0);
+ const int m32 = unpack_scales_q45_K(scales, ksc + 2);
+
+ const uint8_t * sc8 = (const uint8_t *) &sc32;
+ const uint8_t * m8 = (const uint8_t *) &m32;
+
+ const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
+
+#pragma unroll
+ for (int l = 0; l < int(sizeof(int)); ++l) {
+ x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
+ }
+ }
+ }
+#else
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+ int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+
+ x_dm[i] = bxi->dm;
+ }
+
+ constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+ int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+
+ const int * scales = (const int *) bxi->scales;
+
+ const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8);
+ const int scales8 = unpack_scales_q45_K(scales, ksc);
+
+ x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8;
+ }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+ int * x_sc = (int *) (x_df + MMQ_TILE_NE_K/QI6_K);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+ int * x_sc = (int *) (x_df + txs.dm);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR6_K);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+
+ const int ql = get_int_b2(bxi->ql, txi);
+ const int ql0 = (ql >> 0) & 0x0F0F0F0F;
+ const int ql1 = (ql >> 4) & 0x0F0F0F0F;
+
+ const int qh = get_int_b2(bxi->qh, (QI6_K/4) * (txi / (QI6_K/2)) + txi % (QI6_K/4));
+ const int qh0 = ((qh >> ((txi & 0x08) >> 2)) << 4) & 0x30303030;
+ const int qh1 = (qh >> ((txi & 0x08) >> 2)) & 0x30303030;
+
+ const int kq0 = 2*txi - txi % (QI6_K/2) + 0;
+ const int kq1 = 2*txi - txi % (QI6_K/2) + QI6_K/2;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
+ x_qs[i*sram_stride + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*warp_size) {
+ int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride] = bxi->d;
+#else
+ x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K] = bxi->d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int rows_per_warp = warp_size / 4;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps*rows_per_warp) {
+ int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % I;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_sc[i*sram_stride + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8));
+#else
+ x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + threadIdx.x%(MMQ_TILE_NE_K/8)] = get_int_b2(bxi->scales, threadIdx.x%(QI6_K/8));
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq1_s(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ half2 * x_ds = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, I);
+ int * x_qs = (int *) x_tile;
+ half2 * x_ds = (half2 *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR1_S);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq1_s * bxi = (const block_iq1_s *) x + kbx0 + i*stride;
+
+ const int qs_packed = get_int_b2(bxi->qs, kqsx);
+ const uint8_t * qs = (const uint8_t *) &qs_packed;
+
+ const int qh = bxi->qh[kqsx];
+
+ #pragma unroll
+ for (int l = 0; l < QR1_S/2; ++l) {
+ const int grid = iq1s_grid_gpu[qs[l] | (((qh >> (3*l)) & 0x07) << 8)];
+
+ const int grid0 = (grid >> 0) & 0x0F0F0F0F;
+ const int grid1 = (grid >> 4) & 0x0F0F0F0F;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 8*kqsx + (2*l+0)] = grid0;
+ x_qs[i*sram_stride + 8*kqsx + (2*l+1)] = grid1;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid0;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const float d1q = __half2float(bxi->d) * (((qh >> 11) & 0x0E) + 1);
+ const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_ds[i*sram_stride + kqsx] = make_half2(d1q, d1q*delta);
+#else
+ x_ds[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = make_half2(d1q, d1q*delta);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xxs(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XXS, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XXS)) / 2;
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq2_xxs * bxi = (const block_iq2_xxs *) x + kbx0 + i*stride;
+
+ const int q2 = get_int_b2(bxi->qs, 2*kqsx+0);
+ const uint8_t * aux8 = (const uint8_t *) &q2;
+ const uint32_t aux32 = get_int_b2(bxi->qs, 2*kqsx+1);
+
+#pragma unroll
+ for (int l = 0; l < QR2_XXS; ++l) {
+ const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[l]];
+ const uint32_t signs = unpack_ksigns(aux32 >> (7 * l));
+
+ const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+ const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0);
+
+ const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+ const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid0;
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid1;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid0;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid1;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const int ls = aux32 >> 27 | 1; // (scale * 2 + 1)
+ const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4
+#else
+ x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xs(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XS, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XS)) / 2;
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq2_xs * bxi = (const block_iq2_xs *) x + kbx0 + i*stride;
+
+ const int2 q2_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
+ const uint16_t * q2 = (const uint16_t *) &q2_packed;
+
+ #pragma unroll
+ for (int l = 0; l < QR2_XS; ++l) {
+ const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l] & 0x1FF];
+ const uint32_t signs = unpack_ksigns(q2[l] >> 9);
+
+ const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+ const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+
+ const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+ const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l;
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const int ls = bxi->scales[kqsx];
+ const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
+ x_df[i*sram_stride + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
+#else
+ x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
+ x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_s(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_S, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_S)) / 2;
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq2_s * bxi = (const block_iq2_s *) x + kbx0 + i*stride;
+
+ const int qs_packed = get_int_b2(bxi->qs, kqsx);
+ const uint8_t * qs = (const uint8_t *) &qs_packed;
+
+ const int qh = bxi->qh[kqsx];
+
+ const int signs_packed_32 = get_int_b2(bxi->qs, QK_K/32 + kqsx);
+ const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32;
+
+#pragma unroll
+ for (int l = 0; l < QR2_S; ++l) {
+ const int * grid_pos = (const int *)(iq2s_grid + (qs[l] | ((qh << (8-2*l)) & 0x300)));
+
+ const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000);
+ const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000);
+
+ const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0);
+ const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l;
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const int ls = bxi->scales[kqsx];
+ const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
+ x_df[i*sram_stride + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
+#else
+ x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
+ x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_xxs(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_XXS, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_XXS)) / 2;
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq3_xxs * bxi = (const block_iq3_xxs *) x + kbx0 + i*stride;
+
+ const int2 q3_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
+ const uint8_t * q3 = (const uint8_t *) &q3_packed;
+ const uint32_t aux32 = get_int_b2(bxi->qs, QK_K/16 + kqsx);
+
+#pragma unroll
+ for (int l = 0; l < QR3_XXS; ++l) {
+ const int2 grid_pos = make_int2(iq3xxs_grid[q3[2*l+0]], iq3xxs_grid[q3[2*l+1]]);
+ const uint32_t signs = unpack_ksigns(aux32 >> (7*l));
+
+ const int signs0 = __vcmpne4(signs & 0x08040201, 0);
+ const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+
+ const int signs1 = __vcmpne4(signs & 0x80402010, 0);
+ const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 0)] = grid_l;
+ x_qs[i*sram_stride + 8*kqsx + (2*l + 1)] = grid_h;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const int ls = aux32 >> 28;
+ const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kqsx] = (ls*d + d/2)/2;
+#else
+ x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = (ls*d + d/2)/2;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_s(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_S)) / 2;
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * nrows) {
+ int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq3_s * bxi = (const block_iq3_s *) x + kbx0 + i*stride;
+
+ const int2 qs_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
+ const uint8_t * qs = (const uint8_t *) &qs_packed;
+
+ const int qh = bxi->qh[kqsx];
+
+ const int signs_packed_32 = get_int_b2(bxi->signs, kqsx);
+ const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32;
+
+#pragma unroll
+ for (int l = 0; l < QR3_S; ++l) {
+ const int2 grid_pos = make_int2(
+ iq3s_grid[qs[2*l+0] | ((qh << (8 - 2*l)) & 0x100)],
+ iq3s_grid[qs[2*l+1] | ((qh << (7 - 2*l)) & 0x100)]);
+
+ const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000);
+ const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000);
+
+ const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
+ const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + 8*kqsx + (2*l+0)] = grid_l;
+ x_qs[i*sram_stride + 8*kqsx + (2*l+1)] = grid_h;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid_l;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid_h;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ const int ls = 1 + 2*((bxi->scales[kqsx/2] >> (((2*kqsx) << 1) & 0x04)) & 0x0F);
+ const float d = bxi->d;
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kqsx] = ls*d;
+#else
+ x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = ls*d;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_XS, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_XS);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int kqsx = threadIdx.x % threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+
+ const int aux_q4 = get_int_b4(bxi->qs, kqsx);
+ const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
+ const int k0 = 8 * (kqsx / 4) + kqsx % 4;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + k0 + 0] = v.x;
+ x_qs[i*sram_stride + k0 + 4] = v.y;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 4] = v.y;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int rows_per_warp = warp_size / 8;
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / (MMQ_TILE_NE_K/4);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+
+ const float d = __half2float(bxi->d);
+
+ const int ls = ((bxi->scales_l[(threadIdx.x % 8)/2] >> (4*(threadIdx.x % 2))) & 0x0F)
+ | (((bxi->scales_h >> (2*(threadIdx.x % 8))) & 0x03) << 4);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + threadIdx.x % 8] = d * (ls - 32);
+#else
+ x_df[i*(MMQ_TILE_NE_K/4) + i/4 + threadIdx.x % 8] = d * (ls - 32);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_nl(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_NL, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_NL);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI4_NL;
+ const int kqsx = txi % QI4_NL;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbx;
+
+ const int aux_q4 = get_int_b2(bxi->qs, kqsx);
+ const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
+ const int k0 = kbx * (2 * QI4_NL) + kqsx;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + k0 + 0] = v.x;
+ x_qs[i*sram_stride + k0 + QI4_NL] = v.y;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI4_NL] = v.y;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_NL;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kbxd] = __half2float(bxi->d);
+#else
+ x_df[i*(MMQ_TILE_NE_K/QI4_NL) + i/QI4_NL + kbxd] = __half2float(bxi->d);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+// ---------------------------------------------------------------------------------------------
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_MXFP4, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_MXFP4);
+ constexpr int nrows = warp_size / threads_per_row;
+ const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
+ const int kbx = txi / QI_MXFP4;
+ const int kqsx = txi % QI_MXFP4;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nrows*nwarps) {
+ int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbx;
+
+ const int aux_q4 = get_int_b1(bxi->qs, kqsx);
+ const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4);
+ const int k0 = kbx * (2 * QI_MXFP4) + kqsx;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + k0 + 0] = v.x;
+ x_qs[i*sram_stride + k0 + QI_MXFP4] = v.y;
+#else
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
+ x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_MXFP4] = v.y;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+
+ constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI_MXFP4;
+ constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
+ const int kbxd = threadIdx.x % blocks_per_tile_x_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
+ int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
+
+ if (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbxd;
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_df[i*sram_stride + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f;
+#else
+ x_df[i*(MMQ_TILE_NE_K/QI_MXFP4) + i/QI_MXFP4 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+ int * x_qs = (int *) x_tile;
+ uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
+
+ const int txi = threadIdx.x;
+
+ constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+
+ constexpr int threads_per_row = iter_k / QK_MXFP4; // each thread processes 1 block
+ constexpr int rows_per_warp = warp_size / threads_per_row;
+ const int kbx = txi % threads_per_row;
+ const int row_in_warp = txi / threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) {
+ int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
+
+ if constexpr (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx;
+
+ // quantize_mxfp4_mmq permutes nibbles to match the quantized format
+ const int k0 = kbx * 4;
+ memcpy(x_qs + i*sram_stride + k0, bxi->qs, 16);
+
+ // Load E8M0 scales: pack 2 consecutive scales into one uint32
+ if (kbx % 2 == 0) {
+ uint32_t e = bxi->e;
+ e |= ((bxi + 1)->e << 8);
+ x_sc[i*sram_stride + kbx / 2] = e;
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
+#else
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_NVFP4, I);
+ int * x_qs = (int *) x_tile;
+ float * x_df = (float *) (x_qs + txs.qs);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int threads_per_row = MMQ_ITER_K / QK_NVFP4;
+ constexpr int rows_per_warp = warp_size / threads_per_row;
+ const int kbx = threadIdx.x % threads_per_row;
+ const int row_in_warp = threadIdx.x / threads_per_row;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) {
+ int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
+
+ if constexpr (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_nvfp4 * bxi = (const block_nvfp4 *) x + kb0 + i * stride + kbx;
+ const uint32_t * __restrict__ src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
+ const int kqs = 16 * kbx;
+ const int ksc = 4 * kbx;
+
+#pragma unroll
+ for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) {
+ const int2 q0 = get_int_from_table_16(src_qs[2 * sub + 0], kvalues_mxfp4);
+ const int2 q1 = get_int_from_table_16(src_qs[2 * sub + 1], kvalues_mxfp4);
+
+#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ x_qs[i*sram_stride + kqs + 4 * sub + 0] = q0.x;
+ x_qs[i*sram_stride + kqs + 4 * sub + 1] = q1.x;
+ x_qs[i*sram_stride + kqs + 4 * sub + 2] = q0.y;
+ x_qs[i*sram_stride + kqs + 4 * sub + 3] = q1.y;
+ x_df[i*sram_stride + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]);
+#else
+ x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 0] = q0.x;
+ x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 1] = q1.x;
+ x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 2] = q0.y;
+ x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 3] = q1.y;
+ x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI_NVFP4) + i / (QK_NVFP4_SUB / QI_NVFP4) + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]);
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
+ const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+ constexpr int threads_per_row = iter_k / QK_NVFP4; // each thread processes 1 block
+ constexpr int rows_per_warp = warp_size / threads_per_row;
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+
+ uint32_t * x_u32 = (uint32_t *) x_tile;
+
+ const int txi = threadIdx.x;
+ const int kbx = txi % threads_per_row;
+ const int row_in_warp = txi / threads_per_row;
+
+ const block_nvfp4 * bxi_base = (const block_nvfp4 *) x + kbx0 + kbx;
+ uint32_t * x_u32_scale = x_u32 + 64 + kbx;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += rows_per_warp * nwarps) {
+ int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
+
+ if constexpr (fallback) {
+ i = min(i, i_max);
+ }
+
+ const block_nvfp4 * bxi = bxi_base + i * stride;
+
+ const uint32_t * src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
+
+#pragma unroll
+ for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) {
+ x_u32[i*sram_stride + 8*kbx + 2 * sub + 0] = src_qs[2 * sub + 0];
+ x_u32[i*sram_stride + 8*kbx + 2 * sub + 1] = src_qs[2 * sub + 1];
+ }
+
+ x_u32_scale[i*sram_stride] = get_int_b4(bxi->d, 0);
+ }
+}
--- /dev/null
+#pragma once
+
+#include "vecdotq.cuh"
+#include "mma.cuh"
+
+using namespace ggml_cuda_mma;
+
+#include "mmq.cuh"
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, I);
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + txs.qs;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_0*VDR_Q4_0_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+ const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2);
+
+ int u[2*VDR_Q4_0_Q8_1_MMQ];
+
+ constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes();
+ constexpr int mcpy_int = max_cpy / sizeof(int);
+ static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ");
+
+ int tmp0[4], tmp1[4];
+
+ #pragma unroll
+ for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) {
+ ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] );
+ ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_0 + l0 * mcpy_int]);
+ }
+
+ u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3];
+ u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3];
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMQ>
+ (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_0], u,
+ x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + k0/(QR4_0*QI4_0)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, I);
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_1*VDR_Q4_1_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+ const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2);
+
+ int u[2*VDR_Q4_1_Q8_1_MMQ];
+
+ constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes();
+ constexpr int mcpy_int = max_cpy / sizeof(int);
+ static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ");
+
+ int tmp0[4], tmp1[4];
+
+ #pragma unroll
+ for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) {
+ ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] );
+ ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_1 + l0 * mcpy_int]);
+ }
+
+ u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3];
+ u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3];
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMQ>
+ (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_1], u,
+ x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + k0/(QR4_1*QI4_1)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + txs.qs;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMQ>
+ (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K],
+ x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + k0/QI8_0], y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback, mmq_q8_1_ds_layout ds_layout>
+static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr data_layout input_layout = get_input_data_layout();
+ typedef tile<16, 8, int, input_layout> tile_A;
+ typedef tile<16, 8, int, input_layout> tile_B;
+ typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+ const half2 * y_ds = (const half2 *) y;
+
+ const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+ const int k0 = k00 + k01;
+
+ tile_A A[ntx];
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ tile_B B;
+ load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+ float dB;
+ const int j = j0 + tile_C::get_j(0);
+ if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
+ dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+ } else {
+ dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C;
+ mma(C, A[n], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ const int i = i0 + n*tile_A::I + tile_C::get_i(l);
+ const float dA = x_df[i*sram_stride + k0/QI8_0];
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA*dB;
+ }
+ }
+ }
+ }
+#else
+ typedef tile<16, 8, int> tile_A;
+ typedef tile< 8, 8, int> tile_B;
+ typedef tile<16, 8, int> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+ const half2 * y_ds = (const half2 *) y;
+
+ tile_A A[ntx][MMQ_TILE_NE_K/QI8_0];
+ float dA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_0];
+
+ const int i0 = (threadIdx.y/ntx)*rows_per_warp;
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+ const int k0 = k00 + k01;
+
+ load_ldmatrix(A[n][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+ const int k0 = k00 + k01;
+
+ dA[n][l][k01/QI8_0] = x_df[i*sram_stride + k0/QI8_0];
+ }
+ }
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+ tile_B B;
+ float dB[tile_C::ne/2];
+
+ load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int j = j0 + tile_C::get_j(l);
+
+ if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
+ dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+ } else {
+ dB[l] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C;
+ mma(C, A[n][k01/QI8_0], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l/2][k01/QI8_0]*dB[l%2];
+ }
+ }
+ }
+ }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, I);
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_1_q8_1_impl<QR5_1*VDR_Q5_1_Q8_1_MMQ>
+ (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
+ x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + k0/QI8_1], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr data_layout input_layout = get_input_data_layout();
+ typedef tile<16, 8, int, input_layout> tile_A;
+ typedef tile<16, 8, int, input_layout> tile_B;
+ typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_dm = (const half2 *) y;
+
+ const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+ const int k0 = k00 + k01;
+
+ tile_A A[ntx];
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ tile_B B;
+ load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+ const int j = j0 + tile_C::get_j(0);
+ const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C;
+ mma(C, A[n], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ const int i = i0 + n*tile_A::I + tile_C::get_i(l);
+ float2 dmA = __half22float2(x_dm[i*sram_stride + k0/QI8_1]);
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.x*dsB.x*C.x[l];
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.y*dsB.y;
+ }
+ }
+ }
+ }
+#else
+ typedef tile<16, 8, int> tile_A;
+ typedef tile< 8, 8, int> tile_B;
+ typedef tile<16, 8, int> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_dm = (const half2 *) y;
+
+ tile_A A[ntx][MMQ_TILE_NE_K/QI8_1];
+ float2 dmA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_1];
+
+ const int i0 = (threadIdx.y/ntx)*rows_per_warp;
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+ const int k0 = k00 + k01;
+
+ load_ldmatrix(A[n][k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+ const int k0 = k00 + k01;
+
+ dmA[n][l][k01/QI8_1] = __half22float2(x_dm[i*sram_stride + k0/QI8_1]);
+ }
+ }
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+ tile_B B;
+ float2 dsB[tile_C::ne/2];
+
+ load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int j = j0 + tile_C::get_j(l);
+
+ dsB[l] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C;
+ mma(C, A[n][k01/QI8_1], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].x*dsB[l%2].x*C.x[l];
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].y*dsB[l%2].y;
+ }
+ }
+ }
+ }
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+}
+
+// Used for NVFP4, Q3_K, IQ2_S, and IQ2_XS
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, I);
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + txs.qs;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q8_0_16_q8_1_impl<QI8_0>(
+ &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0],
+ &y_qs[j*MMQ_TILE_Y_K + k01],
+ &x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + k0/(QI8_0/2)],
+ y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+// Used for Q3_K, IQ2_S, and IQ2_XS:
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr data_layout input_layout = get_input_data_layout();
+ typedef tile<16, 4, int, input_layout> tile_A;
+ typedef tile<16, 4, int, input_layout> tile_B;
+ typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+ const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+ const int k0 = k00 + k01;
+
+ tile_A A[ntx];
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ tile_B B;
+ load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+ const int j = j0 + tile_C::get_j(0);
+ const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C;
+ mma(C, A[n], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*sram_stride + k0/4] * dB;
+ }
+ }
+ }
+ }
+#elif defined(TURING_MMA_AVAILABLE)
+
+ typedef tile<16, 4, int> tile_A;
+ typedef tile<16, 8, int> tile_A_8;
+ typedef tile< 8, 4, int> tile_B;
+ typedef tile<16, 8, int> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+ const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
+
+ tile_A A[ntx][8];
+ float dA[ntx][tile_C::ne/2][8];
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
+ const int k0 = k00 + k01;
+
+ load_ldmatrix(((tile_A_8 *) A[n])[k01/8], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+ const int k0 = k00 + k01;
+
+ dA[n][l][k01/4] = x_df[i*sram_stride + k0/4];
+ }
+ }
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) {
+ tile_B B[2];
+ float dB[tile_C::ne/2];
+
+ // Here load_generic is faster than load_ldmatrix.
+ load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K);
+ load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int j = j0 + tile_C::get_j(l);
+
+ dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C[2];
+ mma(C[0], A[n][k01/4 + 0], B[0]);
+ mma(C[1], A[n][k01/4 + 1], B[1]);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += dB[l%2]*(C[0].x[l]*dA[n][l/2][k01/4 + 0] + C[1].x[l]*dA[n][l/2][k01/4 + 1]);
+ }
+ }
+ }
+ }
+#else
+ GGML_UNUSED_VARS(x, y, sum, k00);
+ NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, I);
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+ float2 y_df[J/nwarps];
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+ y_df[j0/nwarps] = __half22float2(y_ds[j*MMQ_TILE_Y_K]);
+ }
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K/2; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ constexpr int ns = 2;
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>(
+ &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
+ &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y,
+ &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
+ }
+ }
+ }
+
+ // Some compilers fail to unroll the loop over k01 if there is a conditional statement for ns in the inner loop.
+ // As a workaround 2 separate loops are used instead.
+#pragma unroll
+ for (int k01 = MMQ_TILE_NE_K/2; k01 < MMQ_TILE_NE_K; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ constexpr int ns = 1;
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>(
+ &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
+ &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y,
+ &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr data_layout input_layout = get_input_data_layout();
+ typedef tile<16, 4, int, input_layout> tile_A;
+ typedef tile<16, 4, int, input_layout> tile_B;
+ typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+ const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+ const int k0 = k00 + k01;
+
+ tile_A A[ntx];
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ tile_B B;
+ load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+ const int j = j0 + tile_C::get_j(0);
+ const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y;
+ const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0
+ : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y
+ : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x);
+
+ tile_C Cm;
+ if (k01 >= MMQ_TILE_NE_K * 3/4) {
+ tile_A A1;
+#pragma unroll
+ for (int l = 0; l < tile_A::ne; ++l) {
+ A1.x[l] = 0x01010101;
+ }
+ mma(Cm, A1, B);
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C Cd;
+ mma(Cd, A[n], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+ const float2 dm = __half22float2(x_dm[i*sram_stride + k0/4]);
+ float tmp = Cd.x[l]*dm.x;
+ if (k01 >= MMQ_TILE_NE_K * 3/4) {
+ tmp -= Cm.x[l]*dm.y;
+ }
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB;
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB;
+ }
+ }
+ }
+ }
+#elif defined(TURING_MMA_AVAILABLE)
+
+ typedef tile<16, 4, int> tile_A;
+ typedef tile<16, 8, int> tile_A_8;
+ typedef tile< 8, 4, int> tile_B;
+ typedef tile<16, 8, int> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+ const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
+
+ tile_A A[ntx][8];
+ float dA[ntx][tile_C::ne/2][8];
+ float mA[ntx][tile_C::ne/2][8];
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+ const int k0 = k00 + k01;
+
+ load_ldmatrix(((tile_A_8 *) A[n])[k01/QI8_1], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1/2) {
+ const int k0 = k00 + k01;
+
+ const float2 dm = __half22float2(x_dm[i*sram_stride + k0/(QI8_1/2)]);
+
+ dA[n][l][k01/(QI8_1/2)] = dm.x;
+ mA[n][l][k01/(QI8_1/2)] = dm.y;
+ }
+ }
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ float2 dB[tile_C::ne/2];
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int j = j0 + tile_C::get_j(l);
+
+ dB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K]);
+ }
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
+ tile_B B[2];
+
+ // Here load_generic is faster than load_ldmatrix.
+ load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K);
+ load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K);
+
+ tile_C Cm[2];
+ if (k01 >= MMQ_TILE_NE_K * 3/4) {
+ tile_A A1;
+ A1.x[0] = 0x01010101;
+ A1.x[1] = 0x01010101;
+ mma(Cm[0], A1, B[0]);
+ mma(Cm[1], A1, B[1]);
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C Cd[2];
+
+ mma(Cd[0], A[n][k01/4 + 0], B[0]);
+ mma(Cd[1], A[n][k01/4 + 1], B[1]);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ float tmp = Cd[0].x[l]*dA[n][l/2][k01/4 + 0] + Cd[1].x[l]*dA[n][l/2][k01/4 + 1];
+ if (k01 >= MMQ_TILE_NE_K * 3/4) {
+ tmp -= Cm[0].x[l]*mA[n][l/2][k01/4 + 0] + Cm[1].x[l]*mA[n][l/2][k01/4 + 1];
+ }
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*(k01 < MMQ_TILE_NE_K/2 ? dB[l%2].x : dB[l%2].y);
+ }
+ }
+ }
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K * 3/4; k01 += QI8_1) {
+ float2 sB[tile_C::ne/2];
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int j = j0 + tile_C::get_j(l);
+
+ sB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 0]*sB[l%2].x;
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 1]*sB[l%2].y;
+ }
+ }
+ }
+ }
+#else
+ GGML_UNUSED_VARS(x, y, sum, k00);
+ NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, I);
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + txs.qs;
+ const int * x_sc = (const int *) x_df + txs.dm;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ const int8_t * scales = ((const int8_t *) (x_sc + i*(MMQ_TILE_NE_K/8) + i/8)) + k0/4;
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q3_K_q8_1_impl_mmq(
+ &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], scales,
+ x_df[i], y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, I);
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+ const int * x_sc = (const int *) x_dm + txs.dm;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_K*VDR_Q4_K_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ const uint8_t * sc = (const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/32] + 2*(k01/16);
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q4_K_q8_1_impl_mmq(
+ &x_qs[i*(MMQ_TILE_NE_K + 1) + k0/2], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8,
+ x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, I);
+ const int * x_qs = (const int *) x;
+ const half2 * x_dm = (const half2 *) x_qs + txs.qs;
+ const int * x_sc = (const int *) x_dm + txs.dm;
+ const int * y_qs = (const int *) y + 4;
+ const half2 * y_ds = (const half2 *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR5_K*VDR_Q5_K_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ const uint8_t * sc = ((const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k00/32]) + 2*(k01/16);
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q5_K_q8_1_impl_mmq(
+ &x_qs[i*(QR5_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8,
+ x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, I);
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + txs.qs;
+ const int * x_sc = (const int *) x_df + txs.dm;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+// #pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR6_K*VDR_Q6_K_Q8_1_MMQ) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
+ const int j = j0 + threadIdx.y;
+
+#pragma unroll
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
+ const int i = i0 + threadIdx.x;
+
+ const int8_t * sc = ((const int8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/16]);
+
+ sum[j0/nwarps*I/warp_size + i0/warp_size] += vec_dot_q6_K_q8_1_impl_mmq(
+ &x_qs[i*(QR6_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc,
+ x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K], &y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
+ }
+ }
+ }
+}
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ constexpr data_layout input_layout = get_input_data_layout();
+ typedef tile<16, 4, int, input_layout> tile_A;
+ typedef tile<16, 4, int, input_layout> tile_B;
+ typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+ const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+ const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
+ const int k0 = k00 + k01;
+
+ tile_A A[ntx];
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*sram_stride + k0, sram_stride);
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ tile_B B;
+ load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
+
+ const int j = j0 + tile_C::get_j(0);
+ const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C;
+ mma(C, A[n], B);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+ const int8_t * sc = (const int8_t *) (x_sc + i*sram_stride + k00/16);
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*sram_stride] * dB;
+ }
+ }
+ }
+ }
+#elif defined(TURING_MMA_AVAILABLE)
+
+ typedef tile<16, 4, int> tile_A;
+ typedef tile< 8, 4, int> tile_B;
+ typedef tile<16, 8, int> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+
+ y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
+ const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K;
+ const int * y_qs = (const int *) y + 4;
+ const float * y_df = (const float *) y;
+
+ const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
+
+ tile_A A[ntx][8];
+ int scA[ntx][tile_C::ne/2][8];
+ float dA[ntx][tile_C::ne/2];
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
+ const int k0 = k00 + k01;
+
+ load_ldmatrix(A[n][k01/4 + 0], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + 0), sram_stride);
+ load_ldmatrix(A[n][k01/4 + 1], x_qs + (i0 + n*tile_A::I)*sram_stride + (k0 + tile_A::J), sram_stride);
+ }
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 16) {
+ const int k0 = k00 + k01;
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+ const int sc_packed = x_sc[i*sram_stride + k0/16];
+ const int8_t * sc = (const int8_t *) &sc_packed;
+
+#pragma unroll
+ for (int ksc = 0; ksc < sizeof(int); ++ksc) {
+ scA[n][l][k01/4 + ksc] = sc[ksc];
+ }
+ }
+ }
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
+
+ dA[n][l] = x_df[i*sram_stride];
+ }
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
+ float tmp[ntx][tile_C::ne] = {{0.0f}};
+
+#pragma unroll
+ for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
+ tile_B B[2];
+ float dB[tile_C::ne/2];
+
+ // Here load_generic is faster than load_ldmatrix.
+ load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + 0 + k01, MMQ_TILE_Y_K);
+ load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + tile_B::J + k01, MMQ_TILE_Y_K);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne/2; ++l) {
+ const int j = j0 + tile_C::get_j(l);
+
+ dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+ tile_C C[2];
+ mma(C[0], A[n][k01/4 + 0], B[0]);
+ mma(C[1], A[n][k01/4 + 1], B[1]);
+
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ tmp[n][l] += (C[0].x[l]*scA[n][l/2][k01/4 + 0] + C[1].x[l]*scA[n][l/2][k01/4 + 1])*dB[l%2];
+ }
+ }
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp[n][l]*dA[n][l/2];
+ }
+ }
+ }
+#else
+ GGML_UNUSED_VARS(x, y, sum, k00);
+ NO_DEVICE_CODE;
+#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
+}
+
+// ---------------------------------------------------------------------------------------------
+
+// Shared MMA kernel for MXFP4 and NVFP4 on Blackwell.
+// Both quantizations encode values as e2m1 (FP4) and produce one uint32 scale per
+// m16n8k64 MMA call; only the PTX kind (scale_vec::2X ue8m0 vs scale_vec::4X ue4m3)
+// and the per-type stride constant differ.
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_fp4_fp4_mma(
+ const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
+
+ typedef tile<16, 8, int> tile_A;
+ typedef tile<8, 8, int> tile_B;
+ typedef tile<16, 8, float> tile_C;
+
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp / tile_C::I;
+ constexpr int nfrags = MMQ_TILE_NE_K / tile_A::J;
+
+ y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_K);
+
+ const int * x_qs = (const int *) x;
+ const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
+ const int * y_qs = (const int *) y + 4;
+ const uint32_t * y_sc = (const uint32_t *) y;
+
+ // 2 threads per quad supply the packed scale register to the block_scale MMA,
+ // see https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling
+ const int tidx_A = threadIdx.x / 4 + (threadIdx.x % 2) * 8;
+ const int tidx_B = threadIdx.x / 4;
+ const int i0 = (threadIdx.y / ntx) * rows_per_warp;
+
+ tile_A A[ntx][nfrags];
+ uint32_t scaleA[ntx][nfrags];
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int frag = 0; frag < nfrags; ++frag) {
+ const int k0 = k00 + frag * tile_A::J;
+ load_ldmatrix(A[n][frag], x_qs + (i0 + n * tile_A::I) * sram_stride + k0, sram_stride);
+ scaleA[n][frag] = x_sc[(i0 + n * tile_A::I + tidx_A) * sram_stride + k0 / tile_A::J];
+ }
+ }
+
+#pragma unroll
+ for (int j0 = 0; j0 < J; j0 += ntx * tile_C::J) {
+ tile_B B[nfrags];
+ uint32_t scaleB[nfrags];
+
+#pragma unroll
+ for (int frag = 0; frag < nfrags; ++frag) {
+ const int k0 = frag * tile_B::J;
+ load_generic(B[frag], y_qs + j0 * MMQ_TILE_Y_K + k0, MMQ_TILE_Y_K);
+ scaleB[frag] = y_sc[(j0 + tidx_B) * MMQ_TILE_Y_K + frag];
+ }
+
+#pragma unroll
+ for (int n = 0; n < ntx; ++n) {
+#pragma unroll
+ for (int frag = 0; frag < nfrags; ++frag) {
+ tile_C C = {};
+ mma_block_scaled_fp4<type>(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]);
+#pragma unroll
+ for (int l = 0; l < tile_C::ne; ++l) {
+ sum[(j0 / tile_C::J + n) * tile_C::ne + l] += C.x[l];
+ }
+ }
+ }
+ }
+}
+
#include "quantize.cuh"
#include "mmid.cuh"
+#include <cstdint>
+
static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
switch (args.type_x) {
case GGML_TYPE_Q1_0:
mul_mat_q_case<GGML_TYPE_Q1_0>(ctx, args, stream);
break;
+ case GGML_TYPE_Q2_0:
+ mul_mat_q_case<GGML_TYPE_Q2_0>(ctx, args, stream);
+ break;
case GGML_TYPE_Q4_0:
mul_mat_q_case<GGML_TYPE_Q4_0>(ctx, args, stream);
break;
case GGML_TYPE_Q8_0:
mul_mat_q_case<GGML_TYPE_Q8_0>(ctx, args, stream);
break;
- case GGML_TYPE_MXFP4:
- mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
- break;
- case GGML_TYPE_NVFP4:
- mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
- break;
+// -----------------------------------------------------------------------
case GGML_TYPE_Q2_K:
mul_mat_q_case<GGML_TYPE_Q2_K>(ctx, args, stream);
break;
case GGML_TYPE_Q6_K:
mul_mat_q_case<GGML_TYPE_Q6_K>(ctx, args, stream);
break;
+// -----------------------------------------------------------------------
+ case GGML_TYPE_IQ1_S:
+ mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
+ break;
case GGML_TYPE_IQ2_XXS:
mul_mat_q_case<GGML_TYPE_IQ2_XXS>(ctx, args, stream);
break;
case GGML_TYPE_IQ3_S:
mul_mat_q_case<GGML_TYPE_IQ3_S>(ctx, args, stream);
break;
- case GGML_TYPE_IQ1_S:
- mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
- break;
case GGML_TYPE_IQ4_XS:
mul_mat_q_case<GGML_TYPE_IQ4_XS>(ctx, args, stream);
break;
case GGML_TYPE_IQ4_NL:
mul_mat_q_case<GGML_TYPE_IQ4_NL>(ctx, args, stream);
break;
+// -----------------------------------------------------------------------
+ case GGML_TYPE_MXFP4:
+ mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
+ break;
+ case GGML_TYPE_NVFP4:
+ mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
+ break;
default:
GGML_ABORT("fatal error");
break;
const int64_t s03 = src0->nb[3] / ts_src0;
const int64_t s3 = dst->nb[3] / ts_dst;
- const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA)
- || GGML_CUDA_CC_IS_CDNA(cc);
+ const bool fallback = ne01 % 128 != 0;
- // TODO: tighter pool buffer size vs q8 path
const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4);
+ const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
+ const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ;
if (!ids) {
- const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1 +
- get_mmq_x_max_host(cc)*sizeof(block_q8_1_mmq);
+ const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block +
+ ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+ ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+ if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+ src1_scale.alloc(ne13*ne12*ne11);
+ }
{
const int64_t s11 = src1->nb[1] / ts_src1;
const int64_t s12 = src1->nb[2] / ts_src1;
const int64_t s13 = src1->nb[3] / ts_src1;
if (use_native_fp4) {
+ static constexpr size_t align_float8 = 32;
+ const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));
- quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+ quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
ne11, ne12, ne13, stream);
} else {
// Stride depends on quantization format
const int64_t s12 = use_native_fp4 ?
- ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_K * sizeof(int)) : // block_fp4_mmq holds 256 values
+ ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
const int64_t s13 = ne12*s12;
const mmq_args args = {
src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
+ src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- use_stream_k, ne1};
+ ne1};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
return;
}
ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows);
ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
+ // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
+ // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
+ const bool dedup_bcast = ne11 == 1 && n_expert_used > 1;
+
{
GGML_ASSERT(ids->nb[0] == ggml_element_size(ids));
const int si1 = ids->nb[1] / ggml_element_size(ids);
const int sis1 = nb12 / nb11;
ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
- ne02, ne12, n_expert_used, ne11, si1, sis1, stream);
+ ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
CUDA_CHECK(cudaGetLastError());
}
- const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * sizeof(block_q8_1)/QK8_1 +
- get_mmq_x_max_host(cc)*sizeof(block_q8_1_mmq);
+ const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +
+ ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+ ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+ if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+ src1_scale.alloc(ne12*n_expert_used);
+ }
const int64_t ne11_flat = ne12*n_expert_used;
const int64_t ne12_flat = 1;
const int64_t s13 = src1->nb[3] / ts_src1;
if (use_native_fp4) {
- quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
- ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+ static constexpr size_t align_float8 = 32;
+ const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
+ if (dedup_bcast) {
+ quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10,
+ /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+ } else {
+ quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
+ ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+ }
+ } else if (dedup_bcast) {
+ quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+ /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
} else {
quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
CUDA_CHECK(cudaGetLastError());
}
- static_assert(QK_K == 8 * QK_MXFP4, "QK_K needs to be 8 * QK_MXFP4");
- const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_K * sizeof(int)) :
+ static_assert(QK_FP4_MMQ == 8 * QK_MXFP4, "QK_FP4_MMQ needs to be 8 * QK_MXFP4");
+ const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
const int64_t s13 = ne12*s12;
// Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
const mmq_args args = {
src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
+ src1_scale.ptr,
ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
ne02, ne02, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- use_stream_k, ne12};
-
- ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
-}
-
-void ggml_cuda_op_mul_mat_q(
- ggml_backend_cuda_context & ctx,
- const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
- const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
- const int64_t src1_padded_row_size, cudaStream_t stream) {
-
- const int64_t ne00 = src0->ne[0];
-
- const int64_t ne10 = src1->ne[0];
- const int64_t ne11 = src1->ne[1];
- GGML_ASSERT(ne10 % QK8_1 == 0);
-
- const int64_t ne0 = dst->ne[0];
-
- const int64_t row_diff = row_high - row_low;
- const int64_t stride01 = ne00 / ggml_blck_size(src0->type);
-
- const int id = ggml_cuda_get_device();
- const int cc = ggml_cuda_info().devices[id].cc;
-
- // the main device has a larger memory buffer to hold the results from all GPUs
- // nrows_dst == nrows of the matrix that the kernel writes into
- const int64_t nrows_dst = id == ctx.device ? ne0 : row_diff;
-
- // The stream-k decomposition is only faster for recent NVIDIA GPUs.
- // Also its fixup needs to allocate a temporary buffer in the memory pool.
- // There are multiple parallel CUDA streams for src1_ncols != ne11 which would introduce a race condition for this buffer.
- const bool use_stream_k = ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA)
- || GGML_CUDA_CC_IS_CDNA(cc))
- && src1_ncols == ne11;
- const mmq_args args = {
- src0_dd_i, src0->type, (const int *) src1_ddq_i, nullptr, nullptr, dst_dd_i,
- ne00, row_diff, src1_ncols, stride01, ne11, nrows_dst,
- 1, 1, 0, 0, 0,
- 1, 1, 0, 0, 0,
- use_stream_k, src1_ncols};
+ ne12};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
-
- GGML_UNUSED_VARS(src1, dst, src1_ddf_i, src1_padded_row_size);
}
bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) {
switch (type) {
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q8_0:
- case GGML_TYPE_MXFP4:
- case GGML_TYPE_NVFP4:
+// -------------------------------------------------
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q6_K:
+// -------------------------------------------------
+ case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ2_XXS:
case GGML_TYPE_IQ2_XS:
case GGML_TYPE_IQ2_S:
case GGML_TYPE_IQ3_XXS:
case GGML_TYPE_IQ3_S:
- case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ4_XS:
case GGML_TYPE_IQ4_NL:
+// -------------------------------------------------
+ case GGML_TYPE_MXFP4:
+ case GGML_TYPE_NVFP4:
mmq_supported = true;
break;
default:
return false;
}
+ // MMQ tiles require at least 48 KiB per-block shared memory; fall back to BLAS otherwise.
+ {
+ const int id = ggml_cuda_get_device();
+ const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
+ if (smpbo < 48 * 1024) {
+ return false;
+ }
+ }
+
if (turing_mma_available(cc)) {
return true;
}
#pragma once
#include "common.cuh"
-#include "vecdotq.cuh"
-#include "mma.cuh"
#include <climits>
#include <cstdint>
-using namespace ggml_cuda_mma;
-
#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available.
#define MMQ_ITER_K 256
#define MMQ_ITER_K_FP4 512
#define MMQ_NWARPS 8
-typedef void (*load_tiles_mmq_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
-typedef void (*vec_dot_mmq_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
-typedef void (*mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
- float * __restrict__ dst, const int stride, const int i_max, const int j_max);
+typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
+typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
+typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
+ float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);
enum mmq_q8_1_ds_layout {
MMQ_Q8_1_DS_LAYOUT_D4,
MMQ_Q8_1_DS_LAYOUT_D2S6,
};
+static constexpr int QK8_1_MMQ = 4*QK8_1;
+static constexpr int QK_FP4_MMQ = 2*QK8_1_MMQ;
+
struct block_q8_1_mmq {
// The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block.
// The y float data is first grouped as blocks of 128 values.
half d2s6[8]; // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values,
// stored as d0,d1,s1,s2,s3,s4,s5
};
- int8_t qs[4*QK8_1]; // 128 values quantized to 8 bit each
+ int8_t qs[QK8_1_MMQ];
};
// this struct is used for fp4 data types (currently only used for Blackwell)
// nvfp4 has block size 16, each int32 of d4 contains 4 ue4m3 scales
struct block_fp4_mmq {
uint32_t d4[4];
- int8_t qs[4 * 32]; // 256 FP4 values packed as 4-bit pairs (2 per byte)
+ int8_t qs[QK_FP4_MMQ / 2];
};
-static_assert(sizeof(block_q8_1_mmq) == 4*QK8_1 + 4*sizeof(half2), "Unexpected block_q8_1_mmq size");
+static_assert(sizeof(block_q8_1_mmq) == QK8_1_MMQ + 4*sizeof(half2), "Unexpected block_q8_1_mmq size");
static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1), "Unexpected block_q8_1_mmq size");
static_assert(sizeof(block_fp4_mmq) == sizeof(block_q8_1_mmq), "Unexpected block_fp4_mmq size");
static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
switch (type_x) {
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
return MMQ_Q8_1_DS_LAYOUT_D4;
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
int sc;
};
-static int get_mmq_x_max_host(const int cc) {
- return (turing_mma_available(cc) || amd_wmma_available(cc)) ? 128 :
- GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA ?
-#ifdef GGML_CUDA_FORCE_MMQ
- 128 : 64;
-#else
- MMQ_DP4A_MAX_BATCH_SIZE : 64;
-#endif // GGML_CUDA_FORCE_MMQ
-}
-
-static constexpr __device__ int get_mmq_x_max_device() {
-#if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- return 128;
-#else // defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
-#if defined(GGML_USE_HIP)
- return 64;
-#else // defined(GGML_USE_HIP)
-
-#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
-#ifdef GGML_CUDA_FORCE_MMQ
- return 128;
-#else // GGML_CUDA_FORCE_MMQ
- return MMQ_DP4A_MAX_BATCH_SIZE;
-#endif // GGML_CUDA_FORCE_MMQ
-#else // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
- return 64;
-#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
-
-#endif // defined(GGML_USE_HIP)
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-}
-
-static int get_mmq_y_host(const int cc) {
- return GGML_CUDA_CC_IS_AMD(cc) ? (GGML_CUDA_CC_IS_RDNA1(cc) ? 64 : 128) :
- ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64);
-}
-
-static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type) {
-#if defined(BLACKWELL_MMA_AVAILABLE)
-if (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4) {
- return MMQ_ITER_K_FP4;
-}
-#endif // defined(BLACKWELL_MMA_AVAILABLE)
- return MMQ_ITER_K;
-}
-
-static constexpr __device__ int get_mmq_y_device() {
-#if defined(GGML_USE_HIP)
-#if defined(RDNA1)
- return 64;
-#else
- return 128;
-#endif // defined RDNA1
-#else
-#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
- return 128;
-#else
- return 64;
-#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
-#endif // defined(GGML_USE_HIP)
-}
-
// Decouple shared memory tile sizes from WARP_SIZE to allow for different warp sizes.
// The K dimension of the tiles has either,
// 1*MMQ_TILE_NE_K==32 (always for TILE_Y_K) or 2*MMQ_TILE_NE_K==64 (typically for TILE_X_K),
// in terms of 32 bit elements that means K % 2 == 1 for dp4a or K % 8 == 4 for mma.
#define MMQ_TILE_NE_K 32
-#define MMQ_DP4A_TXS_Q4_0 tile_x_sizes{mmq_y*MMQ_TILE_NE_K + mmq_y, mmq_y*MMQ_TILE_NE_K/QI4_0 + mmq_y/QI4_0, 0}
-#define MMQ_DP4A_TXS_Q4_1 tile_x_sizes{mmq_y*MMQ_TILE_NE_K + mmq_y, mmq_y*MMQ_TILE_NE_K/QI4_1 + mmq_y/QI4_1, 0}
-#define MMQ_DP4A_TXS_Q8_0 tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K*2/QI8_0 + mmq_y/(QI8_0/2), 0}
-#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K*4/QI8_0 + mmq_y/(QI8_0/4), 0}
-#define MMQ_DP4A_TXS_Q8_1 tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K*2/QI8_1 + mmq_y/(QI8_1/2), 0}
-#define MMQ_DP4A_TXS_Q2_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K + mmq_y, 0}
-#define MMQ_DP4A_TXS_Q3_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8}
-#define MMQ_DP4A_TXS_Q4_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K + mmq_y, mmq_y*MMQ_TILE_NE_K/QI4_K, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8}
-#define MMQ_DP4A_TXS_Q5_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K/QI5_K + mmq_y/QI5_K, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8}
-#define MMQ_DP4A_TXS_Q6_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K/QI6_K + mmq_y/QI6_K, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8}
-
-static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml_type type, int mmq_y) {
- switch (type) {
- case GGML_TYPE_Q1_0: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0;
- case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1;
- case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1;
- case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1;
- case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16;
- case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K;
- case GGML_TYPE_Q3_K: return MMQ_DP4A_TXS_Q3_K;
- case GGML_TYPE_Q4_K: return MMQ_DP4A_TXS_Q4_K;
- case GGML_TYPE_Q5_K: return MMQ_DP4A_TXS_Q5_K;
- case GGML_TYPE_Q6_K: return MMQ_DP4A_TXS_Q6_K;
- case GGML_TYPE_IQ2_XXS: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_IQ2_XS: return MMQ_DP4A_TXS_Q8_0_16;
- case GGML_TYPE_IQ2_S: return MMQ_DP4A_TXS_Q8_0_16;
- case GGML_TYPE_IQ3_XXS: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_IQ3_S: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_IQ1_S: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_IQ4_XS: return MMQ_DP4A_TXS_Q8_0;
- case GGML_TYPE_IQ4_NL: return MMQ_DP4A_TXS_Q8_0;
- default: return tile_x_sizes{0, 0, 0};
- }
-}
-
-#define MMQ_MMA_TILE_X_K_Q8_0 (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4)
-#define MMQ_MMA_TILE_X_K_FP4 (2*MMQ_TILE_NE_K + 8 + 4) // MXFP4 and NVFP4 Blackwell
-#define MMQ_MMA_TILE_X_K_NVFP4 (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4) // NVFP4 Generic
-#define MMQ_MMA_TILE_X_K_Q8_1 (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4)
-#define MMQ_MMA_TILE_X_K_Q2_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K + 4)
-#define MMQ_MMA_TILE_X_K_Q3_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4)
-#define MMQ_MMA_TILE_X_K_Q6_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI6_K + MMQ_TILE_NE_K/8 + 7)
-
-static_assert(MMQ_MMA_TILE_X_K_Q8_0 % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_Q8_1 % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_Q2_K % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_Q3_K % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_Q6_K % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_FP4 % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_FP4 == MMQ_MMA_TILE_X_K_Q8_1, "Wrong tile size for MXFP4");
-static_assert(MMQ_MMA_TILE_X_K_NVFP4 % 8 == 4, "Wrong padding.");
-
-
-static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) {
- switch (type) {
- case GGML_TYPE_Q1_0: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_Q4_0: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_Q4_1: return MMQ_MMA_TILE_X_K_Q8_1;
- case GGML_TYPE_Q5_0: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_Q5_1: return MMQ_MMA_TILE_X_K_Q8_1;
- case GGML_TYPE_Q8_0: return MMQ_MMA_TILE_X_K_Q8_0;
- // tile sizes are the same for Q8_1 and FP4 for blackwell
- case GGML_TYPE_MXFP4: return MMQ_MMA_TILE_X_K_Q8_1;
-#if defined(BLACKWELL_MMA_AVAILABLE)
- case GGML_TYPE_NVFP4: return MMQ_MMA_TILE_X_K_FP4;
-#else
- case GGML_TYPE_NVFP4: return MMQ_MMA_TILE_X_K_NVFP4;
-#endif // defined(BLACKWELL_MMA_AVAILABLE)
- case GGML_TYPE_Q2_K: return MMQ_MMA_TILE_X_K_Q2_K;
- case GGML_TYPE_Q3_K: return MMQ_MMA_TILE_X_K_Q3_K;
- case GGML_TYPE_Q4_K: return MMQ_MMA_TILE_X_K_Q8_1;
- case GGML_TYPE_Q5_K: return MMQ_MMA_TILE_X_K_Q8_1;
- case GGML_TYPE_Q6_K: return MMQ_MMA_TILE_X_K_Q6_K;
- case GGML_TYPE_IQ2_XXS: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_IQ2_XS: return MMQ_MMA_TILE_X_K_Q3_K;
- case GGML_TYPE_IQ2_S: return MMQ_MMA_TILE_X_K_Q3_K;
- case GGML_TYPE_IQ3_XXS: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_IQ3_S: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_IQ1_S: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_IQ4_XS: return MMQ_MMA_TILE_X_K_Q8_0;
- case GGML_TYPE_IQ4_NL: return MMQ_MMA_TILE_X_K_Q8_0;
- default: return 0;
- }
-}
-
-// block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales)
-#define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1)
-#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K
-
-static int mmq_get_granularity_host(const int mmq_x, const int cc) {
- if (amd_mfma_available(cc) || amd_wmma_available(cc)) {
- return mmq_x >= 128 ? 32 : 16;
- } else if (turing_mma_available(cc) && mmq_x >= 48) {
- return 16;
- } else {
- return 8;
- }
-}
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-static constexpr __device__ int mmq_get_granularity_device(const int mmq_x) {
- return mmq_x >= 128 ? 32 : 16;
-}
-#elif defined(TURING_MMA_AVAILABLE)
-static constexpr __device__ int mmq_get_granularity_device(const int mmq_x) {
- return mmq_x >= 48 ? 16 : 8;
-}
-#else
-static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) {
- return 8;
-}
-#endif // AMD_MFMA_AVAILABLE
-
-#if defined(GGML_USE_HIP)
-static int mmq_get_nwarps_host(const int cc, const int warp_size) {
- return amd_mfma_available(cc) ? 8 : 256/warp_size;
-}
-#else
-static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) {
- return 256/warp_size;
-}
-#endif // (GGML_USE_HIP)
-
-static constexpr __device__ int mmq_get_nwarps_device() {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- return 8;
-#else
- return 256/ggml_cuda_get_physical_warp_size();
-#endif // AMD_MFMA_AVAILABLE
-}
-
-// ------------------------------------------------------------
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q1_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int blocks_per_iter = MMQ_ITER_K / QK1_0;
- constexpr int threads_per_row = blocks_per_iter * QI1_0;
- constexpr int nrows = warp_size / threads_per_row;
- constexpr int scale_entries_per_block = QK1_0 / QK8_1;
- constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block;
-
- const int txi = threadIdx.x % threads_per_row;
- const int kbx = txi / QI1_0;
- const int kqsx = txi % QI1_0;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + kbx;
- const int qs_offset = 4*kqsx;
- const int qs0 = bxi->qs[qs_offset + 0] | (bxi->qs[qs_offset + 1] << 8) |
- (bxi->qs[qs_offset + 2] << 16) | (bxi->qs[qs_offset + 3] << 24);
-
- int unpacked_bytes[8];
-#pragma unroll
- for (int j = 0; j < 8; ++j) {
- const int shift = j * 4;
- const int bits4 = (qs0 >> shift) & 0x0F;
- const int b0 = (bits4 & 0x01) ? 1 : -1;
- const int b1 = (bits4 & 0x02) ? 1 : -1;
- const int b2 = (bits4 & 0x04) ? 1 : -1;
- const int b3 = (bits4 & 0x08) ? 1 : -1;
- unpacked_bytes[j] = (b0 & 0xFF) | ((b1 & 0xFF) << 8) | ((b2 & 0xFF) << 16) | ((b3 & 0xFF) << 24);
- }
-
- const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
-#pragma unroll
- for (int j = 0; j < 8; ++j) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + dst_offset + j] = unpacked_bytes[j];
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j] = unpacked_bytes[j];
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
- }
-
- const int ksx = threadIdx.x % scale_entries_per_row;
- const int scale_block = ksx / scale_entries_per_block;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps) {
- int i = i0 + threadIdx.y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + scale_block;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + ksx] = bxi->d;
-#else
- x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + ksx] = bxi->d;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q4_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_0);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI4_0;
- const int kqsx = txi % QI4_0;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbx;
- const int qs0 = get_int_b2(bxi->qs, kqsx);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI4_0) + kqsx + 0] = __vsubss4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808);
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsubss4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808);
-#else
- x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = bxi->d;
-#else
- x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + kbxd] = bxi->d;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q4_0_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, mmq_y);
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + txs.qs;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_0*VDR_Q4_0_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2);
-
- int u[2*VDR_Q4_0_Q8_1_MMQ];
-
- constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes();
- constexpr int mcpy_int = max_cpy / sizeof(int);
- static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ");
-
- int tmp0[4], tmp1[4];
-
- #pragma unroll
- for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) {
- ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] );
- ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_0 + l0 * mcpy_int]);
- }
-
- u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3];
- u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3];
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMQ>
- (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_0], u,
- x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + k0/(QR4_0*QI4_0)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q4_1(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, mmq_y);
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_1);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI4_1;
- const int kqsx = txi % QI4_1;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbx;
- const int qs0 = get_int_b4(bxi->qs, kqsx);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI4_1) + kqsx + 0] = (qs0 >> 0) & 0x0F0F0F0F;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI4_1) + kqsx + QI4_1] = (qs0 >> 4) & 0x0F0F0F0F;
-#else
- x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_1;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + kbxd] = bxi->dm;
-#else
- x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + kbxd] = bxi->dm;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q4_1_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, mmq_y);
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + txs.qs;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_1*VDR_Q4_1_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
- const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2);
-
- int u[2*VDR_Q4_1_Q8_1_MMQ];
-
- constexpr int max_cpy = ggml_cuda_get_max_cpy_bytes();
- constexpr int mcpy_int = max_cpy / sizeof(int);
- static_assert(VDR_Q4_0_Q8_1_MMQ == 4, "bad VDR_Q4_0_Q8_1_MMQ");
-
- int tmp0[4], tmp1[4];
-
- #pragma unroll
- for (int l0 = 0; l0 < 4 / mcpy_int; ++l0) {
- ggml_cuda_memcpy_1<max_cpy>(tmp0 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + l0 * mcpy_int] );
- ggml_cuda_memcpy_1<max_cpy>(tmp1 + l0 * mcpy_int, &y_qs[j*MMQ_TILE_Y_K + kyqs + QI4_1 + l0 * mcpy_int]);
- }
-
- u[0]=tmp0[0]; u[2]=tmp0[1]; u[4]=tmp0[2]; u[6]=tmp0[3];
- u[1]=tmp1[0]; u[3]=tmp1[1]; u[5]=tmp1[2]; u[7]=tmp1[3];
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q4_1_q8_1_impl<VDR_Q4_1_Q8_1_MMQ>
- (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_1], u,
- x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + k0/(QR4_1*QI4_1)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q5_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_0, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_0);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI5_0;
- const int kqsx = txi % QI5_0;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbx;
-
- const int ql = get_int_b2(bxi->qs, kqsx);
- const int qh = get_int_b2(bxi->qh, 0) >> (4 * kqsx);
-
- int qs0 = (ql >> 0) & 0x0F0F0F0F;
- qs0 |= (qh << 4) & 0x00000010; // 0 -> 4
- qs0 |= (qh << 11) & 0x00001000; // 1 -> 12
- qs0 |= (qh << 18) & 0x00100000; // 2 -> 20
- qs0 |= (qh << 25) & 0x10000000; // 3 -> 28
- qs0 = __vsubss4(qs0, 0x10101010); // subtract 16
-
- int qs1 = (ql >> 4) & 0x0F0F0F0F;
- qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4
- qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12
- qs1 |= (qh << 2) & 0x00100000; // 18 -> 20
- qs1 |= (qh << 9) & 0x10000000; // 19 -> 28
- qs1 = __vsubss4(qs1, 0x10101010); // subtract 16
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI5_0) + kqsx + 0] = qs0;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + 0] = qs0;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_0;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = bxi->d;
-#else
- x_df[i*(MMQ_TILE_NE_K/QI5_0) + i/QI5_0 + kbxd] = bxi->d;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q5_1(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, mmq_y);
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_1);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI5_1;
- const int kqsx = txi % QI5_1;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbx;
-
- const int ql = get_int_b4(bxi->qs, kqsx);
- const int qh = get_int_b4(bxi->qh, 0) >> (4 * kqsx);
-
- int qs0 = (ql >> 0) & 0x0F0F0F0F;
- qs0 |= (qh << 4) & 0x00000010; // 0 -> 4
- qs0 |= (qh << 11) & 0x00001000; // 1 -> 12
- qs0 |= (qh << 18) & 0x00100000; // 2 -> 20
- qs0 |= (qh << 25) & 0x10000000; // 3 -> 28
-
- int qs1 = (ql >> 4) & 0x0F0F0F0F;
- qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4
- qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12
- qs1 |= (qh << 2) & 0x00100000; // 18 -> 20
- qs1 |= (qh << 9) & 0x10000000; // 19 -> 28
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI5_1) + kqsx + 0] = qs0;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + 0] = qs0;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_1;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + kbxd] = bxi->dm;
-#else
- x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + kbxd] = bxi->dm;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q8_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- // MMQ_ITER_K / (4 * QR8_0) == 64 required. but NV has only 32 threads per warp
- constexpr int threads_per_row = 32;
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI8_0;
- const int kqsx = txi % QI8_0;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 0 + txi] = get_int_b2(bxi[0].qs, kqsx);
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx);
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 0 + txi] = get_int_b2(bxi[0].qs, kqsx);
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = bxi->d;
-#else
- x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = bxi->d;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_mxfp4(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_MXFP4, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_MXFP4);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI_MXFP4;
- const int kqsx = txi % QI_MXFP4;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbx;
-
- const int aux_q4 = get_int_b1(bxi->qs, kqsx);
- const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4);
- const int k0 = kbx * (2 * QI_MXFP4) + kqsx;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + k0 + 0] = v.x;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + k0 + QI_MXFP4] = v.y;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_MXFP4] = v.y;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI_MXFP4;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_1 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f;
-#else
- x_df[i*(MMQ_TILE_NE_K/QI_MXFP4) + i/QI_MXFP4 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check>
-static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restrict__ x,
- int * __restrict__ x_tile,
- const int kbx0,
- const int i_max,
- const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- int * x_qs = (int *) x_tile;
- uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
-
- const int txi = threadIdx.x;
-
- constexpr int iter_k = get_iter_k(GGML_TYPE_MXFP4);
-
- constexpr int threads_per_row = iter_k / QK_MXFP4; // each thread processes 1 block
- constexpr int rows_per_warp = warp_size / threads_per_row;
- const int kbx = txi % threads_per_row;
- const int row_in_warp = txi / threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) {
- int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
-
- if constexpr (need_check) {
- i = min(i, i_max);
- }
-
- const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx;
-
- // quantize_mxfp4_mmq permutes nibbles to match the quantized format
- const int k0 = kbx * 4;
- memcpy(x_qs + i * MMQ_MMA_TILE_X_K_FP4 + k0, bxi->qs, 16);
-
- // Load E8M0 scales: pack 2 consecutive scales into one uint32
- if (kbx % 2 == 0) {
- uint32_t e = bxi->e;
- e |= ((bxi + 1)->e << 8);
- x_sc[i * MMQ_MMA_TILE_X_K_FP4 + kbx / 2] = e;
- }
- }
-}
-
-#ifdef BLACKWELL_MMA_AVAILABLE
-template <int mmq_y, bool need_check>
-static __device__ __forceinline__ void load_tiles_nvfp4_nvfp4(const char * __restrict__ x,
- int * __restrict__ x_tile,
- const int kbx0,
- const int i_max,
- const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int iter_k = get_iter_k(GGML_TYPE_NVFP4);
- constexpr int threads_per_row = iter_k / QK_NVFP4; // each thread processes 1 block
- constexpr int rows_per_warp = warp_size / threads_per_row;
-
- uint32_t * x_u32 = (uint32_t *) x_tile;
-
- const int txi = threadIdx.x;
- const int kbx = txi % threads_per_row;
- const int row_in_warp = txi / threads_per_row;
-
- const block_nvfp4 * bxi_base = (const block_nvfp4 *) x + kbx0 + kbx;
- uint32_t * x_u32_scale = x_u32 + 64 + kbx;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) {
- int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
-
- if constexpr (need_check) {
- i = min(i, i_max);
- }
-
- const block_nvfp4 * bxi = bxi_base + i * stride;
- const int row_base = i * MMQ_MMA_TILE_X_K_FP4;
- const int q_base = row_base + 8 * kbx;
-
- const uint32_t * src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
-
-#pragma unroll
- for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) {
- x_u32[q_base + 2 * sub + 0] = src_qs[2 * sub + 0];
- x_u32[q_base + 2 * sub + 1] = src_qs[2 * sub + 1];
- }
-
- x_u32_scale[row_base] = get_int_b4(bxi->d, 0);
- }
-}
-
-// Shared MMA kernel for MXFP4 and NVFP4 on Blackwell.
-// Both quantizations encode values as e2m1 (FP4) and produce one uint32 scale per
-// m16n8k64 MMA call; only the PTX kind (scale_vec::2X ue8m0 vs scale_vec::4X ue4m3)
-// and the per-type stride constant differ.
-template <int mmq_x, int mmq_y, ggml_type type>
-static __device__ __forceinline__ void vec_dot_fp4_fp4_mma(const int * __restrict__ x,
- const int * __restrict__ y,
- float * __restrict__ sum,
- const int k00) {
- static_assert(type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4,
- "vec_dot_fp4_fp4_mma: type must be MXFP4 or NVFP4");
-
- typedef tile<16, 8, int> tile_A;
- typedef tile<8, 8, int> tile_B;
- typedef tile<16, 8, float> tile_C;
-
- constexpr int stride = MMQ_MMA_TILE_X_K_FP4;
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = 2 * granularity;
- constexpr int ntx = rows_per_warp / tile_C::I;
- constexpr int nfrags = MMQ_TILE_NE_K / tile_A::J;
-
- y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
- const int * y_qs = (const int *) y + 4;
- const uint32_t * y_sc = (const uint32_t *) y;
-
- // 2 threads per quad supply the packed scale register to the block_scale MMA,
- // see https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling
- const int tidx_A = threadIdx.x / 4 + (threadIdx.x % 2) * 8;
- const int tidx_B = threadIdx.x / 4;
- const int i0 = (threadIdx.y / ntx) * rows_per_warp;
-
- tile_A A[ntx][nfrags];
- uint32_t scaleA[ntx][nfrags];
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int frag = 0; frag < nfrags; ++frag) {
- const int k0 = k00 + frag * tile_A::J;
- load_ldmatrix(A[n][frag], x_qs + (i0 + n * tile_A::I) * stride + k0, stride);
- scaleA[n][frag] = x_sc[(i0 + n * tile_A::I + tidx_A) * stride + k0 / tile_A::J];
- }
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) {
- tile_B B[nfrags];
- uint32_t scaleB[nfrags];
-
-#pragma unroll
- for (int frag = 0; frag < nfrags; ++frag) {
- const int k0 = frag * tile_B::J;
- load_generic(B[frag], y_qs + j0 * MMQ_TILE_Y_K + k0, MMQ_TILE_Y_K);
- scaleB[frag] = y_sc[(j0 + tidx_B) * MMQ_TILE_Y_K + frag];
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int frag = 0; frag < nfrags; ++frag) {
- tile_C C = {};
- mma_block_scaled_fp4<type>(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]);
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- sum[(j0 / tile_C::J + n) * tile_C::ne + l] += C.x[l];
- }
- }
- }
- }
-}
-#endif // BLACKWELL_MMA_AVAILABLE
-
-
-template <int mmq_y, bool need_check>
-static __device__ __forceinline__ void load_tiles_nvfp4(const char * __restrict__ x,
- int * __restrict__ x_tile,
- const int kb0,
- const int i_max,
- const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_NVFP4, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / QK_NVFP4;
- constexpr int rows_per_warp = warp_size / threads_per_row;
- const int kbx = threadIdx.x % threads_per_row;
- const int row_in_warp = threadIdx.x / threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) {
- int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;
-
- if constexpr (need_check) {
- i = min(i, i_max);
- }
-
- const block_nvfp4 * bxi = (const block_nvfp4 *) x + kb0 + i * stride + kbx;
- const uint32_t * __restrict__ src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
- const int kqs = 16 * kbx;
- const int ksc = 4 * kbx;
-
-#pragma unroll
- for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) {
- const int2 q0 = get_int_from_table_16(src_qs[2 * sub + 0], kvalues_mxfp4);
- const int2 q1 = get_int_from_table_16(src_qs[2 * sub + 1], kvalues_mxfp4);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 0] = q0.x;
- x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 1] = q1.x;
- x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 2] = q0.y;
- x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 3] = q1.y;
- x_df[i * MMQ_MMA_TILE_X_K_NVFP4 + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]);
-#else
- x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 0] = q0.x;
- x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 1] = q1.x;
- x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 2] = q0.y;
- x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 3] = q1.y;
- x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI_NVFP4) + i / (QK_NVFP4_SUB / QI_NVFP4) + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q8_0_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, mmq_y);
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + txs.qs;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl<float, VDR_Q8_0_Q8_1_MMQ>
- (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K],
- x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + k0/QI8_0], y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)]);
- }
- }
- }
-}
-
-template <int mmq_x, int mmq_y, mmq_q8_1_ds_layout ds_layout>
-static __device__ __forceinline__ void vec_dot_q8_0_q8_1_mma(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr data_layout input_layout = get_input_data_layout();
- typedef tile<16, 8, int, input_layout> tile_A;
- typedef tile<16, 8, int, input_layout> tile_B;
- typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
- const half2 * y_ds = (const half2 *) y;
-
- const int i0 = (threadIdx.y / ntx) * rows_per_warp;
-
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
- const int k0 = k00 + k01;
-
- tile_A A[ntx];
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_0 + k0, MMQ_MMA_TILE_X_K_Q8_0);
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- tile_B B;
- load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
-
- float dB;
- const int j = j0 + tile_C::get_j(0);
- if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
- dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
- } else {
- dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C;
- mma(C, A[n], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- const int i = i0 + n*tile_A::I + tile_C::get_i(l);
- const float dA = x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + k0/QI8_0];
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA*dB;
- }
- }
- }
- }
-#else
- typedef tile<16, 8, int> tile_A;
- typedef tile< 8, 8, int> tile_B;
- typedef tile<16, 8, int> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = 2 * granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
- const half2 * y_ds = (const half2 *) y;
-
- tile_A A[ntx][MMQ_TILE_NE_K/QI8_0];
- float dA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_0];
-
- const int i0 = (threadIdx.y/ntx)*rows_per_warp;
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
- const int k0 = k00 + k01;
-
- load_ldmatrix(A[n][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_0 + k0, MMQ_MMA_TILE_X_K_Q8_0);
- }
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
- const int k0 = k00 + k01;
-
- dA[n][l][k01/QI8_0] = x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + k0/QI8_0];
- }
- }
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
- tile_B B;
- float dB[tile_C::ne/2];
-
- load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int j = j0 + tile_C::get_j(l);
-
- if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
- dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
- } else {
- dB[l] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C;
- mma(C, A[n][k01/QI8_0], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l/2][k01/QI8_0]*dB[l%2];
- }
- }
- }
- }
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-}
-
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q8_1_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, mmq_y);
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + txs.qs;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q8_1_q8_1_impl<QR5_1*VDR_Q5_1_Q8_1_MMQ>
- (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
- x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + k0/QI8_1], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q8_1_q8_1_mma(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr data_layout input_layout = get_input_data_layout();
- typedef tile<16, 8, int, input_layout> tile_A;
- typedef tile<16, 8, int, input_layout> tile_B;
- typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_dm = (const half2 *) y;
-
- const int i0 = (threadIdx.y / ntx) * rows_per_warp;
-
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
- const int k0 = k00 + k01;
-
- tile_A A[ntx];
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_1 + k0, MMQ_MMA_TILE_X_K_Q8_1);
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- tile_B B;
- load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
-
- const int j = j0 + tile_C::get_j(0);
- const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C;
- mma(C, A[n], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- const int i = i0 + n*tile_A::I + tile_C::get_i(l);
- float2 dmA = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + k0/QI8_1]);
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.x*dsB.x*C.x[l];
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.y*dsB.y;
- }
- }
- }
- }
-#else
- typedef tile<16, 8, int> tile_A;
- typedef tile< 8, 8, int> tile_B;
- typedef tile<16, 8, int> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = 2 * granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_dm = (const half2 *) y;
-
- tile_A A[ntx][MMQ_TILE_NE_K/QI8_1];
- float2 dmA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_1];
-
- const int i0 = (threadIdx.y/ntx)*rows_per_warp;
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
- const int k0 = k00 + k01;
-
- load_ldmatrix(A[n][k01/QI8_1], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_1 + k0, MMQ_MMA_TILE_X_K_Q8_1);
- }
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int i = i0 + n*tile_A::I + tile_C::get_i(2*l);
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
- const int k0 = k00 + k01;
-
- dmA[n][l][k01/QI8_1] = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + k0/QI8_1]);
- }
- }
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
- tile_B B;
- float2 dsB[tile_C::ne/2];
-
- load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int j = j0 + tile_C::get_j(l);
-
- dsB[l] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C;
- mma(C, A[n][k01/QI8_1], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].x*dsB[l%2].x*C.x[l];
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].y*dsB[l%2].y;
- }
- }
- }
- }
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-}
-
-// Used for NVFP4, Q3_K, IQ2_S, and IQ2_XS
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q8_0_16_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = MMQ_DP4A_TXS_Q8_0_16;
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + txs.qs;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q8_0_16_q8_1_impl<QI8_0>(
- &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0],
- &y_qs[j*MMQ_TILE_Y_K + k01],
- &x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + k0/(QI8_0/2)],
- y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-// Used for Q3_K, IQ2_S, and IQ2_XS:
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q8_0_16_q8_1_mma(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr data_layout input_layout = get_input_data_layout();
- typedef tile<16, 4, int, input_layout> tile_A;
- typedef tile<16, 4, int, input_layout> tile_B;
- typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
- const int i0 = (threadIdx.y / ntx) * rows_per_warp;
-
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
- const int k0 = k00 + k01;
-
- tile_A A[ntx];
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q3_K + k0, MMQ_MMA_TILE_X_K_Q3_K);
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- tile_B B;
- load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
-
- const int j = j0 + tile_C::get_j(0);
- const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C;
- mma(C, A[n], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(l);
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*MMQ_MMA_TILE_X_K_Q3_K + k0/4] * dB;
- }
- }
- }
- }
-#elif defined(TURING_MMA_AVAILABLE)
-
- typedef tile<16, 4, int> tile_A;
- typedef tile<16, 8, int> tile_A_8;
- typedef tile< 8, 4, int> tile_B;
- typedef tile<16, 8, int> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = 2 * granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
- const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
-
- tile_A A[ntx][8];
- float dA[ntx][tile_C::ne/2][8];
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
- const int k0 = k00 + k01;
-
- load_ldmatrix(((tile_A_8 *) A[n])[k01/8], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q3_K + k0, MMQ_MMA_TILE_X_K_Q3_K);
- }
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
- const int k0 = k00 + k01;
-
- dA[n][l][k01/4] = x_df[i*MMQ_MMA_TILE_X_K_Q3_K + k0/4];
- }
- }
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) {
- tile_B B[2];
- float dB[tile_C::ne/2];
-
- // Here load_generic is faster than load_ldmatrix.
- load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K);
- load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int j = j0 + tile_C::get_j(l);
-
- dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C[2];
- mma(C[0], A[n][k01/4 + 0], B[0]);
- mma(C[1], A[n][k01/4 + 1], B[1]);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += dB[l%2]*(C[0].x[l]*dA[n][l/2][k01/4 + 0] + C[1].x[l]*dA[n][l/2][k01/4 + 1]);
- }
- }
- }
- }
-#else
- GGML_UNUSED_VARS(x, y, sum, k00);
- NO_DEVICE_CODE;
-#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q2_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, mmq_y);
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR2_K);
- constexpr int nrows = ggml_cuda_get_physical_warp_size() / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q2_K * bxi = (const block_q2_K *) x + kbx0 + i*stride;
-
- const int x_ql_0 = get_int_b2(bxi->qs, kqsx);
-
-#pragma unroll
- for (int l = 0; l < QR2_K; ++l) {
- const int k = (kqsx/8)*32 + l*8 + kqsx % 8;
-
- const int x_qs_k = (x_ql_0 >> (2*l)) & 0x03030303;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q2_K + k] = x_qs_k;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- const int sc_m = bxi->scales[kqsx];
-#ifdef FAST_FP16_AVAILABLE
- const half2 x_dm_ik = __hmul2(bxi->dm, make_half2(sc_m & 0x0F, sc_m >> 4));
-#else
- const float2 bxi_dmf = __half22float2(bxi->dm);
- const half2 x_dm_ik = make_half2(bxi_dmf.x*(sc_m & 0x0F), bxi_dmf.y*(sc_m >> 4));
-#endif // FAST_FP16_AVAILABLE
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + kqsx] = x_dm_ik;
-#else
- x_dm[i*(MMQ_TILE_NE_K + 1) + kqsx] = x_dm_ik;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q2_K_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, mmq_y);
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + txs.qs;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
- float2 y_df[mmq_x/nwarps];
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
- y_df[j0/nwarps] = __half22float2(y_ds[j*MMQ_TILE_Y_K]);
- }
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K/2; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- constexpr int ns = 2;
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>(
- &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
- &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y,
- &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
- }
- }
- }
-
- // Some compilers fail to unroll the loop over k01 if there is a conditional statement for ns in the inner loop.
- // As a workaround 2 separate loops are used instead.
-#pragma unroll
- for (int k01 = MMQ_TILE_NE_K/2; k01 < MMQ_TILE_NE_K; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- constexpr int ns = 1;
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq<ns>(
- &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01],
- &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y,
- &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
- }
- }
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q2_K_q8_1_mma(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr data_layout input_layout = get_input_data_layout();
- typedef tile<16, 4, int, input_layout> tile_A;
- typedef tile<16, 4, int, input_layout> tile_B;
- typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
- const int i0 = (threadIdx.y / ntx) * rows_per_warp;
-
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
- const int k0 = k00 + k01;
-
- tile_A A[ntx];
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q2_K + k0, MMQ_MMA_TILE_X_K_Q2_K);
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- tile_B B;
- load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
-
- const int j = j0 + tile_C::get_j(0);
- const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y;
- const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0
- : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y
- : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x);
-
- tile_C Cm;
- if (k01 >= MMQ_TILE_NE_K * 3/4) {
- tile_A A1;
-#pragma unroll
- for (int l = 0; l < tile_A::ne; ++l) {
- A1.x[l] = 0x01010101;
- }
- mma(Cm, A1, B);
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C Cd;
- mma(Cd, A[n], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(l);
- const float2 dm = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + k0/4]);
- float tmp = Cd.x[l]*dm.x;
- if (k01 >= MMQ_TILE_NE_K * 3/4) {
- tmp -= Cm.x[l]*dm.y;
- }
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB;
- sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB;
- }
- }
- }
- }
-#elif defined(TURING_MMA_AVAILABLE)
-
- typedef tile<16, 4, int> tile_A;
- typedef tile<16, 8, int> tile_A_8;
- typedef tile< 8, 4, int> tile_B;
- typedef tile<16, 8, int> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = 2 * granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
- const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
-
- tile_A A[ntx][8];
- float dA[ntx][tile_C::ne/2][8];
- float mA[ntx][tile_C::ne/2][8];
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
- const int k0 = k00 + k01;
-
- load_ldmatrix(((tile_A_8 *) A[n])[k01/QI8_1], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q2_K + k0, MMQ_MMA_TILE_X_K_Q2_K);
- }
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1/2) {
- const int k0 = k00 + k01;
-
- const float2 dm = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + k0/(QI8_1/2)]);
-
- dA[n][l][k01/(QI8_1/2)] = dm.x;
- mA[n][l][k01/(QI8_1/2)] = dm.y;
- }
- }
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- float2 dB[tile_C::ne/2];
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int j = j0 + tile_C::get_j(l);
-
- dB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K]);
- }
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) {
- tile_B B[2];
-
- // Here load_generic is faster than load_ldmatrix.
- load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K);
- load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K);
-
- tile_C Cm[2];
- if (k01 >= MMQ_TILE_NE_K * 3/4) {
- tile_A A1;
- A1.x[0] = 0x01010101;
- A1.x[1] = 0x01010101;
- mma(Cm[0], A1, B[0]);
- mma(Cm[1], A1, B[1]);
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C Cd[2];
-
- mma(Cd[0], A[n][k01/4 + 0], B[0]);
- mma(Cd[1], A[n][k01/4 + 1], B[1]);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- float tmp = Cd[0].x[l]*dA[n][l/2][k01/4 + 0] + Cd[1].x[l]*dA[n][l/2][k01/4 + 1];
- if (k01 >= MMQ_TILE_NE_K * 3/4) {
- tmp -= Cm[0].x[l]*mA[n][l/2][k01/4 + 0] + Cm[1].x[l]*mA[n][l/2][k01/4 + 1];
- }
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*(k01 < MMQ_TILE_NE_K/2 ? dB[l%2].x : dB[l%2].y);
- }
- }
- }
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K * 3/4; k01 += QI8_1) {
- float2 sB[tile_C::ne/2];
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int j = j0 + tile_C::get_j(l);
-
- sB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]);
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 0]*sB[l%2].x;
- sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 1]*sB[l%2].y;
- }
- }
- }
- }
-#else
- GGML_UNUSED_VARS(x, y, sum, k00);
- NO_DEVICE_CODE;
-#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q3_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
- int * x_sc = (int *) (x_df + txs.dm);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR3_K);
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
-
- const int x_ql_0 = get_int_b2(bxi->qs, kqsx);
- const int x_qh_0 = get_int_b2(bxi->hmask, kqsx % (QI3_K/2)) >> (4 * (kqsx / (QI3_K/2)));
-
-#pragma unroll
- for (int l = 0; l < QR3_K; ++l) {
- const int k = (kqsx/8)*32 + l*8 + kqsx % 8;
-
- const int x_ql_k = (x_ql_0 >> (2*l)) & 0x03030303;
- const int x_qh_k = ((x_qh_0 >> l) << 2) & 0x04040404;
-
- const int x_qs_k = __vsubss4(x_ql_k | x_qh_k, 0x04040404);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k] = x_qs_k;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
- }
-
- constexpr int rows_per_warp = warp_size / 4;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) {
- int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/4;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
-
- const int ksc = threadIdx.x % 4;
-
- const int ksc_low = ksc % (QI3_K/8);
- const int shift_low = 4 * (ksc / (QI3_K/8));
- const int sc_low = (get_int_b2(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F;
-
- const int ksc_high = QI3_K/8;
- const int shift_high = 2 * ksc;
- const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030;
-
- const int sc = __vsubss4(sc_low | sc_high, 0x20202020);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- const int8_t * sc8 = (const int8_t *) ≻
- const float d = bxi->d;
-
-#pragma unroll
- for (int l = 0; l < int(sizeof(int)); ++l) {
- x_df[i*MMQ_MMA_TILE_X_K_Q3_K + sizeof(int)*ksc + l] = d*sc8[l];
- }
-#else
- x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = sc;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
-#if !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE))
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) {
- int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride;
-
- x_df[i] = bxi->d;
- }
-#endif // !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)) || defined(AMD_WMMA_AVAILABLE)
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q3_K_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, mmq_y);
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + txs.qs;
- const int * x_sc = (const int *) x_df + txs.dm;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- const int8_t * scales = ((const int8_t *) (x_sc + i*(MMQ_TILE_NE_K/8) + i/8)) + k0/4;
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q3_K_q8_1_impl_mmq(
- &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], scales,
- x_df[i], y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, const int ksc) {
- // scale arrangement after the following two lines:
- // - ksc == 0: sc0, sc1, sc2, sc3
- // - ksc == 1: sc4, sc5, sc6, sc7
- // - ksc == 2: m0, m1, m2, m3
- // - ksc == 3: m4, m5, m6, m7
- return ((scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F) | // lower 4 bits
- ((scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030); // upper 2 bits
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q4_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, mmq_y);
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + txs.qs);
- int * x_sc = (int *) (x_dm + txs.dm);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_K);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
- const int qs0 = get_int_b4(bxi->qs, txi);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 16*(txi/8) + txi % 8 + 0] = (qs0 >> 0) & 0x0F0F0F0F;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 16*(txi/8) + txi % 8 + 8] = (qs0 >> 4) & 0x0F0F0F0F;
-#else
- x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr int rows_per_warp = warp_size / 2;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- // Need if on AMD instead of % because warp_size == 64
- // This causes double work and throughput loss (MI300X)
- // H100 loses about 100 t/s with 'if' condition over '%'
- int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2;
- if (i < mmq_y) {
-#else
- int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % mmq_y;
- {
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
-
- const int * scales = (const int *) bxi->scales;
- const int ksc = threadIdx.x % 2;
-
- const int sc32 = unpack_scales_q45_K(scales, ksc + 0);
- const int m32 = unpack_scales_q45_K(scales, ksc + 2);
-
- const uint8_t * sc8 = (const uint8_t *) &sc32;
- const uint8_t * m8 = (const uint8_t *) &m32;
-
- const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
-
- #pragma unroll
- for (int l = 0; l < sizeof(int); ++l) {
- x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
- }
- }
- }
-#else
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) {
- int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
-
- x_dm[i] = bxi->dm;
- }
- constexpr int rows_per_warp = warp_size / 4;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) {
- int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
-
- const int * scales = (const int *) bxi->scales;
-
- const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8);
- const int scales8 = unpack_scales_q45_K(scales, ksc);
-
- x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8;
- }
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q4_K_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, mmq_y);
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + txs.qs;
- const int * x_sc = (const int *) x_dm + txs.dm;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_K*VDR_Q4_K_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- const uint8_t * sc = (const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/32] + 2*(k01/16);
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q4_K_q8_1_impl_mmq(
- &x_qs[i*(MMQ_TILE_NE_K + 1) + k0/2], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8,
- x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q5_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, mmq_y);
- int * x_qs = (int *) x_tile;
- half2 * x_dm = (half2 *) (x_qs + txs.qs);
- int * x_sc = (int *) (x_dm + txs.dm);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_K);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
- const int ky = QR5_K*txi;
-
- const int ql = get_int_b4(bxi->qs, txi);
- const int ql0 = (ql >> 0) & 0x0F0F0F0F;
- const int ql1 = (ql >> 4) & 0x0F0F0F0F;
-
- const int qh = get_int_b4(bxi->qh, txi % (QI5_K/4));
- const int qh0 = ((qh >> (2 * (txi / (QI5_K/4)) + 0)) << 4) & 0x10101010;
- const int qh1 = ((qh >> (2 * (txi / (QI5_K/4)) + 1)) << 4) & 0x10101010;
-
- const int kq0 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + 0;
- const int kq1 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + QI5_K/4;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kq0] = ql0 | qh0;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kq1] = ql1 | qh1;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = ql0 | qh0;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = ql1 | qh1;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr int rows_per_warp = warp_size / 2;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) {
-#if defined(AMD_MFMA_AVAILABLE)
- // Need if on AMD instead of % because warp_size == 64
- // This causes double work and throughput loss (MI300X)
- // H100 loses about 100 t/s with 'if' condition over '%'
- int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2;
- if (i < mmq_y) {
-#else
- int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % mmq_y;
- {
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
-
- const int * scales = (const int *) bxi->scales;
- const int ksc = threadIdx.x % 2;
-
- const int sc32 = unpack_scales_q45_K(scales, ksc + 0);
- const int m32 = unpack_scales_q45_K(scales, ksc + 2);
-
- const uint8_t * sc8 = (const uint8_t *) &sc32;
- const uint8_t * m8 = (const uint8_t *) &m32;
-
- const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
-
-#pragma unroll
- for (int l = 0; l < int(sizeof(int)); ++l) {
- x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
- }
- }
- }
-#else
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) {
- int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
-
- x_dm[i] = bxi->dm;
- }
-
- constexpr int rows_per_warp = warp_size / 4;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) {
- int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
-
- const int * scales = (const int *) bxi->scales;
-
- const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8);
- const int scales8 = unpack_scales_q45_K(scales, ksc);
-
- x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8;
- }
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q5_K_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, mmq_y);
- const int * x_qs = (const int *) x;
- const half2 * x_dm = (const half2 *) x_qs + txs.qs;
- const int * x_sc = (const int *) x_dm + txs.dm;
- const int * y_qs = (const int *) y + 4;
- const half2 * y_ds = (const half2 *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR5_K*VDR_Q5_K_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- const uint8_t * sc = ((const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k00/32]) + 2*(k01/16);
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q5_K_q8_1_impl_mmq(
- &x_qs[i*(QR5_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8,
- x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_q6_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
- int * x_sc = (int *) (x_df + MMQ_TILE_NE_K/QI6_K);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
- int * x_sc = (int *) (x_df + txs.dm);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR6_K);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
-
- const int ql = get_int_b2(bxi->ql, txi);
- const int ql0 = (ql >> 0) & 0x0F0F0F0F;
- const int ql1 = (ql >> 4) & 0x0F0F0F0F;
-
- const int qh = get_int_b2(bxi->qh, (QI6_K/4) * (txi / (QI6_K/2)) + txi % (QI6_K/4));
- const int qh0 = ((qh >> ((txi & 0x08) >> 2)) << 4) & 0x30303030;
- const int qh1 = (qh >> ((txi & 0x08) >> 2)) & 0x30303030;
-
- const int kq0 = 2*txi - txi % (QI6_K/2) + 0;
- const int kq1 = 2*txi - txi % (QI6_K/2) + QI6_K/2;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q6_K + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
- x_qs[i*MMQ_MMA_TILE_X_K_Q6_K + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) {
- int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q6_K] = bxi->d;
-#else
- x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K] = bxi->d;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int rows_per_warp = warp_size / 4;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) {
- int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % mmq_y;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_sc[i*MMQ_MMA_TILE_X_K_Q6_K + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8));
-#else
- x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + threadIdx.x%(MMQ_TILE_NE_K/8)] = get_int_b2(bxi->scales, threadIdx.x%(QI6_K/8));
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q6_K_q8_1_dp4a(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, mmq_y);
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + txs.qs;
- const int * x_sc = (const int *) x_df + txs.dm;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
-// #pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR6_K*VDR_Q6_K_Q8_1_MMQ) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
- const int j = j0 + threadIdx.y;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
- const int i = i0 + threadIdx.x;
-
- const int8_t * sc = ((const int8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/16]);
-
- sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q6_K_q8_1_impl_mmq(
- &x_qs[i*(QR6_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc,
- x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K], &y_df[j*MMQ_TILE_Y_K + k01/QI8_1]);
- }
- }
- }
-}
-
-template <int mmq_x, int mmq_y>
-static __device__ __forceinline__ void vec_dot_q6_K_q8_1_mma(
- const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
-#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr data_layout input_layout = get_input_data_layout();
- typedef tile<16, 4, int, input_layout> tile_A;
- typedef tile<16, 4, int, input_layout> tile_B;
- typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
- const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
- const int i0 = (threadIdx.y / ntx) * rows_per_warp;
-
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
- const int k0 = k00 + k01;
-
- tile_A A[ntx];
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- load_ldmatrix(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + k0, MMQ_MMA_TILE_X_K_Q6_K);
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- tile_B B;
- load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K);
-
- const int j = j0 + tile_C::get_j(0);
- const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C;
- mma(C, A[n], B);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(l);
- const int8_t * sc = (const int8_t *) (x_sc + i*MMQ_MMA_TILE_X_K_Q6_K + k00/16);
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*MMQ_MMA_TILE_X_K_Q6_K] * dB;
- }
- }
- }
- }
-#elif defined(TURING_MMA_AVAILABLE)
-
- typedef tile<16, 4, int> tile_A;
- typedef tile< 8, 4, int> tile_B;
- typedef tile<16, 8, int> tile_C;
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int rows_per_warp = 2 * granularity;
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
-
- y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K);
-
- const int * x_qs = (const int *) x;
- const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2;
- const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K;
- const int * y_qs = (const int *) y + 4;
- const float * y_df = (const float *) y;
-
- const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I);
-
- tile_A A[ntx][8];
- int scA[ntx][tile_C::ne/2][8];
- float dA[ntx][tile_C::ne/2];
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
- const int k0 = k00 + k01;
-
- load_ldmatrix(A[n][k01/4 + 0], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + (k0 + 0), MMQ_MMA_TILE_X_K_Q6_K);
- load_ldmatrix(A[n][k01/4 + 1], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + (k0 + tile_A::J), MMQ_MMA_TILE_X_K_Q6_K);
- }
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 16) {
- const int k0 = k00 + k01;
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
-
- const int sc_packed = x_sc[i*MMQ_MMA_TILE_X_K_Q6_K + k0/16];
- const int8_t * sc = (const int8_t *) &sc_packed;
-
-#pragma unroll
- for (int ksc = 0; ksc < sizeof(int); ++ksc) {
- scA[n][l][k01/4 + ksc] = sc[ksc];
- }
- }
- }
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int i = i0 + n*tile_C::I + tile_C::get_i(2*l);
-
- dA[n][l] = x_df[i*MMQ_MMA_TILE_X_K_Q6_K];
- }
- }
-
-#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
- float tmp[ntx][tile_C::ne] = {{0.0f}};
-
-#pragma unroll
- for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
- tile_B B[2];
- float dB[tile_C::ne/2];
-
- // Here load_generic is faster than load_ldmatrix.
- load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + 0 + k01, MMQ_TILE_Y_K);
- load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + tile_B::J + k01, MMQ_TILE_Y_K);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne/2; ++l) {
- const int j = j0 + tile_C::get_j(l);
-
- dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1];
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
- tile_C C[2];
- mma(C[0], A[n][k01/4 + 0], B[0]);
- mma(C[1], A[n][k01/4 + 1], B[1]);
-
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- tmp[n][l] += (C[0].x[l]*scA[n][l/2][k01/4 + 0] + C[1].x[l]*scA[n][l/2][k01/4 + 1])*dB[l%2];
- }
- }
- }
-
-#pragma unroll
- for (int n = 0; n < ntx; ++n) {
-#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
- sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp[n][l]*dA[n][l/2];
- }
- }
- }
-#else
- GGML_UNUSED_VARS(x, y, sum, k00);
- NO_DEVICE_CODE;
-#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq4_nl(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_NL, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_NL);
- constexpr int nrows = warp_size / threads_per_row;
- const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
- const int kbx = txi / QI4_NL;
- const int kqsx = txi % QI4_NL;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbx;
-
- const int aux_q4 = get_int_b2(bxi->qs, kqsx);
- const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
- const int k0 = kbx * (2 * QI4_NL) + kqsx;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + QI4_NL] = v.y;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI4_NL] = v.y;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_NL;
- constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row;
- const int kbxd = threadIdx.x % blocks_per_tile_x_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbxd;
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = __half2float(bxi->d);
-#else
- x_df[i*(MMQ_TILE_NE_K/QI4_NL) + i/QI4_NL + kbxd] = __half2float(bxi->d);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq2_xxs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XXS, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XXS)) / 2;
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_iq2_xxs * bxi = (const block_iq2_xxs *) x + kbx0 + i*stride;
-
- const int q2 = get_int_b2(bxi->qs, 2*kqsx+0);
- const uint8_t * aux8 = (const uint8_t *) &q2;
- const uint32_t aux32 = get_int_b2(bxi->qs, 2*kqsx+1);
-
-#pragma unroll
- for (int l = 0; l < QR2_XXS; ++l) {
- const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[l]];
- const uint32_t signs = unpack_ksigns(aux32 >> (7 * l));
-
- const int signs0 = __vcmpne4(signs & 0x08040201, 0);
- const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0);
-
- const int signs1 = __vcmpne4(signs & 0x80402010, 0);
- const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 0)] = grid0;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 1)] = grid1;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid0;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid1;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- const int ls = aux32 >> 27 | 1; // (scale * 2 + 1)
- const float d = bxi->d;
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4
-#else
- x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq2_xs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = MMQ_DP4A_TXS_Q8_0_16;
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XS)) / 2;
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_iq2_xs * bxi = (const block_iq2_xs *) x + kbx0 + i*stride;
-
- const int2 q2_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
- const uint16_t * q2 = (const uint16_t *) &q2_packed;
-
- #pragma unroll
- for (int l = 0; l < QR2_XS; ++l) {
- const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l] & 0x1FF];
- const uint32_t signs = unpack_ksigns(q2[l] >> 9);
-
- const int signs0 = __vcmpne4(signs & 0x08040201, 0);
- const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
-
- const int signs1 = __vcmpne4(signs & 0x80402010, 0);
- const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 0)] = grid_l;
- x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 1)] = grid_h;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- const int ls = bxi->scales[kqsx];
- const float d = bxi->d;
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
- x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
-#else
- x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
- x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq2_s(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_S, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_S)) / 2;
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_iq2_s * bxi = (const block_iq2_s *) x + kbx0 + i*stride;
-
- const int qs_packed = get_int_b2(bxi->qs, kqsx);
- const uint8_t * qs = (const uint8_t *) &qs_packed;
-
- const int qh = bxi->qh[kqsx];
-
- const int signs_packed_32 = get_int_b2(bxi->qs, QK_K/32 + kqsx);
- const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32;
-
-#pragma unroll
- for (int l = 0; l < QR2_S; ++l) {
- const int * grid_pos = (const int *)(iq2s_grid + (qs[l] | ((qh << (8-2*l)) & 0x300)));
-
- const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000);
- const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000);
-
- const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0);
- const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 0)] = grid_l;
- x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 1)] = grid_h;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-
- const int ls = bxi->scales[kqsx];
- const float d = bxi->d;
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
- x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
-#else
- x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4;
- x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq3_xxs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_XXS, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_XXS)) / 2;
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
-
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
- }
-
- const block_iq3_xxs * bxi = (const block_iq3_xxs *) x + kbx0 + i*stride;
-
- const int2 q3_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
- const uint8_t * q3 = (const uint8_t *) &q3_packed;
- const uint32_t aux32 = get_int_b2(bxi->qs, QK_K/16 + kqsx);
-
-#pragma unroll
- for (int l = 0; l < QR3_XXS; ++l) {
- const int2 grid_pos = make_int2(iq3xxs_grid[q3[2*l+0]], iq3xxs_grid[q3[2*l+1]]);
- const uint32_t signs = unpack_ksigns(aux32 >> (7*l));
-
- const int signs0 = __vcmpne4(signs & 0x08040201, 0);
- const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
-
- const int signs1 = __vcmpne4(signs & 0x80402010, 0);
- const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+// block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales)
+#define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1)
+#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 0)] = grid_l;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 1)] = grid_h;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
+enum ggml_cuda_mmq_sram_layout {
+ GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0,
+ GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1,
+ GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K,
+ GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K,
+ GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K,
+ GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, // MXFP4 and NVFP4 on Blackwell.
+ GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, // Generic NVFP4
+};
- const int ls = aux32 >> 28;
- const float d = bxi->d;
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = (ls*d + d/2)/2;
+static constexpr __host__ __device__ int ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_sram_layout sram_layout) {
+ switch (sram_layout) {
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0:
+ return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4;
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1:
+ return 2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_1 + 4;
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K:
+ return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K + 4;
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K:
+ return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4;
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K:
+ return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI6_K + MMQ_TILE_NE_K/8 + 7;
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_FP4:
+ return 2*MMQ_TILE_NE_K + 8 + 4;
+ case GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4:
+ return 2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4;
+ default:
+ return -1;
+ }
+}
+
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0) % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1) % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K) % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K) % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K) % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) % 8 == 4, "Wrong padding.");
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4) % 8 == 4, "Wrong padding.");
+
+static_assert(ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_FP4) == ggml_cuda_mmq_get_sram_stride(GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1), "Wrong tile size for MXFP4");
+
+// Config options for the MMQ kernel.
+// Should not affect results, only speed/register pressure/shared memory use.
+struct ggml_cuda_mmq_config {
+ ggml_type type; // src0->type
+ int nthreads; // Number of threads per CUDA block.
+ int occupancy; // Targeted occupancy for the MMA kernel.
+ int I; // SRAM tile width in src0->ne[1]/dst->ne[0] direction.
+ int J; // SRAM tile width in src1->ne[1]/dst->ne[1] direction.
+ ggml_cuda_mmq_sram_layout sram_layout; // SRAM tile length in src0->ne[0]/src1->ne[0] direction (physical 32 bit elements).
+ int K_vram; // VRAM tile length in src0->ne[0]/src1->ne[0] direction (logical elements).
+ bool stream_k; // Whether or not to use stream-k decomposition.
+ bool fallback; // Whether a fallback for out-of-bounds check in src0->ne[1] direction is needed.
+
+ constexpr __host__ __device__ ggml_cuda_mmq_config(
+ ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool stream_k, bool fallback) :
+ type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), stream_k(stream_k), fallback(fallback) {}
+
+ constexpr __device__ int rows_per_warp() const {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ return 16;
#else
- x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = (ls*d + d/2)/2;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ return J >= 48 && J % 16 == 0 ? 32 : 16;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
}
-}
-
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq3_s(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-
- constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_S)) / 2;
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
-
- if (need_check) {
- i = min(i, i_max);
+ // TODO transition all combinations of GPUs and quantizations to the MMA data layout.
+ __host__ int use_mma_data_layout(const int cc) const {
+ if (amd_mfma_available(cc) || amd_wmma_available(cc) || turing_mma_available(cc)) {
+ return true;
}
+ return false;
+ }
- const block_iq3_s * bxi = (const block_iq3_s *) x + kbx0 + i*stride;
-
- const int2 qs_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1));
- const uint8_t * qs = (const uint8_t *) &qs_packed;
+ constexpr __device__ bool use_mma_data_layout() const {
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+ return true;
+#else
+ return false;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
+ }
- const int qh = bxi->qh[kqsx];
+};
- const int signs_packed_32 = get_int_b2(bxi->signs, kqsx);
- const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32;
+#define CASE(type_, nthreads_, occupancy_, I_, J_, sram_layout_, K_vram_, stream_k_, fallback_) \
+ if (type == (type_) && J == (J_) && fallback == (fallback_)) { \
+ static_assert((nthreads_) % 32 == 0 && (nthreads_) <= 512, "bad nthreads"); \
+ static_assert( (occupancy_) <= 8, "bad occupancy"); \
+ static_assert((I_) % 32 == 0, "bad I"); \
+ static_assert((J_) % 8 == 0, "bad J"); \
+ static_assert((K_vram_) % 256 == 0, "bad K_vram"); \
+ return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
+ } \
-#pragma unroll
- for (int l = 0; l < QR3_S; ++l) {
- const int2 grid_pos = make_int2(
- iq3s_grid[qs[2*l+0] | ((qh << (8 - 2*l)) & 0x100)],
- iq3s_grid[qs[2*l+1] | ((qh << (7 - 2*l)) & 0x100)]);
+#include "mmq-config-pascal.cuh"
+#include "mmq-config-ampere.cuh"
+#include "mmq-config-blackwell.cuh"
- const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000);
- const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000);
+#include "mmq-config-cdna.cuh"
+#include "mmq-config-rdna2.cuh"
+#include "mmq-config-rdna3.cuh"
+#include "mmq-config-rdna3-5.cuh"
+#include "mmq-config-rdna4.cuh"
- const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0);
- const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1);
+#undef CASE
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l+0)] = grid_l;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l+1)] = grid_h;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid_l;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid_h;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) {
+ if (GGML_CUDA_CC_IS_AMD(cc)) {
+ if (GGML_CUDA_CC_IS_CDNA(cc)) {
+ return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
}
-
- const int ls = 1 + 2*((bxi->scales[kqsx/2] >> (((2*kqsx) << 1) & 0x04)) & 0x0F);
- const float d = bxi->d;
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = ls*d;
-#else
- x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = ls*d;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ if (GGML_CUDA_CC_IS_RDNA4(cc)) {
+ return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
+ }
+ if (GGML_CUDA_CC_IS_RDNA3_5(cc)) {
+ return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
+ }
+ if (GGML_CUDA_CC_IS_RDNA3(cc)) { // covers RDNA 3.0
+ return ggml_cuda_mmq_get_config_rdna3(type, J, fallback);
+ }
+ return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
+ }
+ if (blackwell_mma_available(cc)) {
+ return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
+ }
+ if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) {
+ return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
}
+ return ggml_cuda_mmq_get_config_pascal(type, J, fallback);
}
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq1_s(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- half2 * x_ds = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
+static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) {
+#ifdef GGML_USE_HIP
+#ifdef CDNA
+ return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
+#elif defined(RDNA4)
+ return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
+#elif defined(RDNA3_5)
+ return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
+#elif defined(RDNA3)
+ return ggml_cuda_mmq_get_config_rdna3(type, J, fallback);
+#else
+ return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
+#endif // CDNA
+#else
+#ifdef BLACKWELL_MMA_AVAILABLE
+ return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
+#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
+ return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, mmq_y);
- int * x_qs = (int *) x_tile;
- half2 * x_ds = (half2 *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ return ggml_cuda_mmq_get_config_pascal(type, J, fallback);
+#endif // BLACKWELL_MMA_AVAILABLE
+#endif // GGML_USE_HIP
+ GGML_UNUSED_VARS(type, J, fallback);
+}
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR1_S);
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
+static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).type;
+}
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) {
- int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row;
+static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).type;
+}
- if (need_check) {
- i = min(i, i_max);
- }
+static __host__ int ggml_cuda_mmq_get_nthreads(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).nthreads;
+}
- const block_iq1_s * bxi = (const block_iq1_s *) x + kbx0 + i*stride;
+static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).nthreads;
+}
- const int qs_packed = get_int_b2(bxi->qs, kqsx);
- const uint8_t * qs = (const uint8_t *) &qs_packed;
+static __host__ int ggml_cuda_mmq_get_occupancy(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).occupancy;
+}
- const int qh = bxi->qh[kqsx];
+static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).occupancy;
+}
- #pragma unroll
- for (int l = 0; l < QR1_S/2; ++l) {
- const int grid = iq1s_grid_gpu[qs[l] | (((qh >> (3*l)) & 0x07) << 8)];
+static __host__ int ggml_cuda_mmq_get_I(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).I;
+}
- const int grid0 = (grid >> 0) & 0x0F0F0F0F;
- const int grid1 = (grid >> 4) & 0x0F0F0F0F;
+static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).I;
+}
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 8*kqsx + (2*l+0)] = grid0;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 8*kqsx + (2*l+1)] = grid1;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid0;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid1;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
+static __host__ int ggml_cuda_mmq_get_J(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).J;
+}
- const float d1q = __half2float(bxi->d) * (((qh >> 11) & 0x0E) + 1);
- const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000);
+static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).J;
+}
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_ds[i*MMQ_MMA_TILE_X_K_Q8_1 + kqsx] = make_half2(d1q, d1q*delta);
-#else
- x_ds[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = make_half2(d1q, d1q*delta);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
+static __host__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).sram_layout;
}
-template <int mmq_y, bool need_check> static __device__ __forceinline__ void load_tiles_iq4_xs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
- constexpr int nwarps = mmq_get_nwarps_device();
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).sram_layout;
+}
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
-#else
- constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_XS, mmq_y);
- int * x_qs = (int *) x_tile;
- float * x_df = (float *) (x_qs + txs.qs);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+static __host__ int ggml_cuda_mmq_get_K_vram(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).K_vram;
+}
- constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_XS);
- constexpr int nrows = warp_size / threads_per_row;
- const int kqsx = threadIdx.x % threads_per_row;
+static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).K_vram;
+}
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) {
- int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row);
+static __host__ bool ggml_cuda_mmq_get_stream_k(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).stream_k;
+}
- if (need_check) {
- i = min(i, i_max);
- }
+static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).stream_k;
+}
- const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+static __host__ int ggml_cuda_mmq_get_fallback(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, cc).fallback;
+}
- const int aux_q4 = get_int_b4(bxi->qs, kqsx);
- const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
- const int k0 = 8 * (kqsx / 4) + kqsx % 4;
+static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).fallback;
+}
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x;
- x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 4] = v.y;
-#else
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x;
- x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 4] = v.y;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- }
+// ---------------------------------------------------------------------------------------------
- constexpr int rows_per_warp = warp_size / 8;
-#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) {
- int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / (MMQ_TILE_NE_K/4);
+static __host__ int ggml_cuda_mmq_get_sram_stride(const ggml_type type, const int J, const bool fallback, const int cc) {
+ return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, cc));
+}
+
+static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback));
+}
- if (need_check) {
- i = min(i, i_max);
+static __host__ int ggml_cuda_mmq_get_J_max(const ggml_type type, const bool fallback, const int cc, const int64_t ne11) {
+ int ret = std::min(ne11, int64_t(512));
+ ret -= ret % 8;
+ for (;ret > 0; ret -= 8) {
+ if (ggml_cuda_mmq_get_config(type, ret, fallback, cc).type != GGML_TYPE_COUNT) {
+ return ret;
}
+ }
+ return ret;
+}
- const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+static constexpr __device__ int ggml_cuda_mmq_get_rows_per_warp(ggml_type type, int J, bool fallback) {
+ return ggml_cuda_mmq_get_config(type, J, fallback).rows_per_warp();
+}
- const float d = __half2float(bxi->d);
+#define MMQ_DP4A_TXS_Q4_0 tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_0 + I/QI4_0, 0}
+#define MMQ_DP4A_TXS_Q4_1 tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_1 + I/QI4_1, 0}
+#define MMQ_DP4A_TXS_Q8_0 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_0 + I/(QI8_0/2), 0}
+#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*4/QI8_0 + I/(QI8_0/4), 0}
+#define MMQ_DP4A_TXS_Q8_1 tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K*2/QI8_1 + I/(QI8_1/2), 0}
+#define MMQ_DP4A_TXS_Q2_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K + I, 0}
+#define MMQ_DP4A_TXS_Q3_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I, I*MMQ_TILE_NE_K/8 + I/8}
+#define MMQ_DP4A_TXS_Q4_K tile_x_sizes{I*MMQ_TILE_NE_K + I, I*MMQ_TILE_NE_K/QI4_K, I*MMQ_TILE_NE_K/8 + I/8}
+#define MMQ_DP4A_TXS_Q5_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI5_K + I/QI5_K, I*MMQ_TILE_NE_K/8 + I/8}
+#define MMQ_DP4A_TXS_Q6_K tile_x_sizes{I*MMQ_TILE_NE_K*2 + I, I*MMQ_TILE_NE_K/QI6_K + I/QI6_K, I*MMQ_TILE_NE_K/8 + I/8}
- const int ls = ((bxi->scales_l[(threadIdx.x % 8)/2] >> (4*(threadIdx.x % 2))) & 0x0F)
- | (((bxi->scales_h >> (2*(threadIdx.x % 8))) & 0x03) << 4);
+static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml_type type, int I) {
+ switch (type) {
+ case GGML_TYPE_Q1_0: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_Q2_0: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0;
+ case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1;
+ case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1;
+ case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1;
+ case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16;
+ case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K;
+ case GGML_TYPE_Q3_K: return MMQ_DP4A_TXS_Q3_K;
+ case GGML_TYPE_Q4_K: return MMQ_DP4A_TXS_Q4_K;
+ case GGML_TYPE_Q5_K: return MMQ_DP4A_TXS_Q5_K;
+ case GGML_TYPE_Q6_K: return MMQ_DP4A_TXS_Q6_K;
+ case GGML_TYPE_IQ2_XXS: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_IQ2_XS: return MMQ_DP4A_TXS_Q8_0_16;
+ case GGML_TYPE_IQ2_S: return MMQ_DP4A_TXS_Q8_0_16;
+ case GGML_TYPE_IQ3_XXS: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_IQ3_S: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_IQ1_S: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_IQ4_XS: return MMQ_DP4A_TXS_Q8_0;
+ case GGML_TYPE_IQ4_NL: return MMQ_DP4A_TXS_Q8_0;
+ default: return tile_x_sizes{0, 0, 0};
+ }
+}
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + threadIdx.x % 8] = d * (ls - 32);
-#else
- x_df[i*(MMQ_TILE_NE_K/4) + i/4 + threadIdx.x % 8] = d * (ls - 32);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+// FIXME temporary until all combinations of data types and GPUs can use the MMA data layout
+static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config & config, const int cc) {
+ if (config.use_mma_data_layout(cc)) {
+ return config.I * ggml_cuda_mmq_get_sram_stride(config.sram_layout) * 4;
}
+ const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(config.type, config.I);
+ return (txs.qs + txs.dm + txs.sc) * 4;
}
-template<int mmq_x, int mmq_y, bool need_check>
-static __device__ __forceinline__ void mmq_write_back_dp4a(
+// ------------------------------------------------------------
+
+#include "mmq-load-tiles.cuh"
+#include "mmq-vec-dot.cuh"
+
+template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
- const int stride, const int i_max, const int j_max) {
- constexpr int nwarps = mmq_get_nwarps_device();
+ const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ const bool y_scale_used = y_scale != nullptr;
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
if (j > j_max) {
}
#pragma unroll
- for (int i0 = 0; i0 < mmq_y; i0 += warp_size) {
+ for (int i0 = 0; i0 < I; i0 += warp_size) {
const int i = i0 + threadIdx.x;
- if (need_check && i > i_max) {
+ if (fallback && i > i_max) {
continue;
}
- dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (mmq_y/warp_size) + i0/warp_size];
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ if (y_scale_used) {
+ dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+ } else {
+ dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+ }
+ } else {
+ dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+ GGML_UNUSED(y_scale_used);
+ }
}
}
}
-template<ggml_type type, int mmq_x, int mmq_y, bool need_check>
-static __device__ __forceinline__ void mmq_write_back_mma(
- const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
- const int stride, const int i_max, const int j_max) {
-
- constexpr int granularity = mmq_get_granularity_device(mmq_x);
- constexpr int nwarps = mmq_get_nwarps_device();
+template<ggml_type type, int J, bool fallback>
+static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
+ const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
+ const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr int tileC_IJ = mmq_get_granularity_device(0);
- typedef tile<tileC_IJ, tileC_IJ, int, DATA_LAYOUT_J_MAJOR> tile_C;
- constexpr int rows_per_warp = granularity;
+ typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
#else
- typedef tile<16, 8, int> tile_C;
- constexpr int rows_per_warp = 2 * granularity;
-#endif // defined(AMD_MFMA_AVAILABLE)
- constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
+ typedef tile<16, 8, int> tile_C;
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
+ constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);
-#if defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- static_assert(nwarps*tile_C::I == mmq_y, "nwarps*tile_C::I != mmq_y");
-#else
- GGML_UNUSED(nwarps);
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+
+ const bool y_scale_used = y_scale != nullptr;
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) {
+ for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
#pragma unroll
for (int n = 0; n < ntx; ++n) {
#pragma unroll
const int i = i0 + n*tile_C::I + tile_C::get_i(l);
- if (need_check && i > i_max) {
+ if (fallback && i > i_max) {
continue;
}
- dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ if (y_scale_used) {
+ dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
+ } else {
+ dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+ }
+ } else {
+ dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+ GGML_UNUSED(y_scale_used);
+ }
}
}
}
// -------------------------------------------------------------------------------------------------------------------------------------
-template <int mmq_x, int mmq_y, bool need_check, ggml_type type>
-struct mmq_type_traits;
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q1_0> {
- static constexpr int vdr = VDR_Q1_0_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q1_0<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q4_0> {
- static constexpr int vdr = VDR_Q4_0_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_0<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_DS4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q4_1> {
- static constexpr int vdr = VDR_Q4_1_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_1<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_1_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q5_0> {
- static constexpr int vdr = VDR_Q5_0_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q5_0<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q5_1> {
- static constexpr int vdr = VDR_Q5_1_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q5_1<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_1_q8_1_dp4a<mmq_x, mmq_y>;
-};
+// TODO remove this struct and use ggml_cuda_mmq_sram_layout instead.
+struct ggml_cuda_mmq_util_funcs {
+ int vdr;
+ ggml_cuda_mmq_load_tiles_t load_tiles;
+ ggml_cuda_mmq_vec_dot_t vec_dot;
+ ggml_cuda_mmq_write_back_t write_back;
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q8_0> {
- static constexpr int vdr = VDR_Q8_0_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q8_0<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
+ constexpr __host__ __device__ ggml_cuda_mmq_util_funcs(
+ int vdr, ggml_cuda_mmq_load_tiles_t load_tiles, ggml_cuda_mmq_vec_dot_t vec_dot, ggml_cuda_mmq_write_back_t write_back) :
+ vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {}
};
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_MXFP4> {
- static constexpr int vdr = VDR_MXFP4_Q8_1_MMQ;
-#ifdef BLACKWELL_MMA_AVAILABLE
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_mxfp4_fp4<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_fp4_fp4_mma<mmq_x, mmq_y, GGML_TYPE_MXFP4>;
-#else
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_mxfp4<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
-#endif // BLACKWELL_MMA_AVAILABLE
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
+template <ggml_type type, int J, bool fallback>
+static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+
+ if (!ggml_cuda_mmq_get_config(type, J, fallback).use_mma_data_layout()) {
+ switch (type) {
+ case GGML_TYPE_Q1_0:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q1_0_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q2_0:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q2_0_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q4_0:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q4_0_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q4_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q4_1:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q4_1_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q4_1_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q5_0:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q5_0_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q5_1:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q5_1_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q8_0:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q8_0_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+// ---------------------------------------------------------------------------------------------
+ case GGML_TYPE_Q2_K:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q2_K_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q2_K_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q3_K:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q3_K_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q3_K_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q4_K:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q4_K_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q5_K:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q5_K_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_Q6_K:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_Q6_K_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+// ---------------------------------------------------------------------------------------------
+ case GGML_TYPE_IQ1_S:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ1_S_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ2_XXS:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ2_XXS_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ2_XS:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ2_XS_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ2_S:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ2_S_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ3_XXS:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ3_XXS_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ3_S:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ3_S_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ4_XS:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ4_XS_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_IQ4_NL:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_IQ4_NL_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+// ---------------------------------------------------------------------------------------------
+ case GGML_TYPE_MXFP4:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_MXFP4_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ case GGML_TYPE_NVFP4:
+ return ggml_cuda_mmq_util_funcs(
+ VDR_NVFP4_Q8_1_MMQ,
+ ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
+ ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
+ default:
+ return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
+ }
+ }
+
+// ---------------------------------------------------------------------------------------------
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_NVFP4> {
- static constexpr int vdr = VDR_NVFP4_Q8_1_MMQ;
#ifdef BLACKWELL_MMA_AVAILABLE
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_nvfp4_nvfp4<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_fp4_fp4_mma<mmq_x, mmq_y, GGML_TYPE_NVFP4>;
-#else
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_nvfp4<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma<mmq_x, mmq_y>;
+ switch (type) {
+ case GGML_TYPE_MXFP4:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_NVFP4:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ default:
+ break;
+ }
#endif // BLACKWELL_MMA_AVAILABLE
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q2_K> {
- static constexpr int vdr = VDR_Q2_K_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q2_K<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q2_K_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q2_K_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q3_K> {
- static constexpr int vdr = VDR_Q3_K_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q3_K<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q3_K_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q4_K> {
- static constexpr int vdr = VDR_Q4_K_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_K<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_K_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q5_K> {
- static constexpr int vdr = VDR_Q5_K_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q5_K<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q5_K_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_Q6_K> {
- static constexpr int vdr = VDR_Q6_K_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_q6_K<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q6_K_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q6_K_q8_1_dp4a<mmq_x, mmq_y>;
-};
-
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ2_XXS> {
- static constexpr int vdr = VDR_IQ2_XXS_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq2_xxs<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ2_XS> {
- static constexpr int vdr = VDR_IQ2_XS_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq2_xs<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a<mmq_x, mmq_y>;
-};
+// ---------------------------------------------------------------------------------------------
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ2_S> {
- static constexpr int vdr = VDR_IQ2_S_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq2_s<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a<mmq_x, mmq_y>;
-};
+ switch (type) {
+ case GGML_TYPE_Q1_0:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q2_0:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q4_0:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_DS4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q4_1:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q5_0:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q5_1:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q8_0:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+// ---------------------------------------------------------------------------------------------
+ case GGML_TYPE_Q2_K:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q3_K:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q4_K:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q5_K:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_Q6_K:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+// ---------------------------------------------------------------------------------------------
+ case GGML_TYPE_IQ1_S:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ2_XXS:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ2_XS:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ2_S:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ3_XXS:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ3_S:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ4_XS:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_IQ4_NL:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+// ---------------------------------------------------------------------------------------------
+ case GGML_TYPE_MXFP4:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ case GGML_TYPE_NVFP4:
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ default:
+ return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
+ }
+}
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ3_XXS> {
- static constexpr int vdr = VDR_IQ3_XXS_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq3_xxs<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
+template <ggml_type type, int J, bool fallback>
+static constexpr __device__ int ggml_cuda_mmq_get_vdr() {
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().vdr;
+}
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ3_S> {
- static constexpr int vdr = VDR_IQ3_S_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq3_s<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
+template <ggml_type type, int J, bool fallback>
+static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() {
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().load_tiles;
+}
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ1_S> {
- static constexpr int vdr = VDR_IQ1_S_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq1_s<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma<mmq_x, mmq_y>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_1_q8_1_dp4a<mmq_x, mmq_y>;
-};
+template <ggml_type type, int J, bool fallback>
+static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() {
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().vec_dot;
+}
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ4_NL> {
- static constexpr int vdr = VDR_IQ4_NL_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq4_nl<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
+template <ggml_type type, int J, bool fallback>
+static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() {
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().write_back;
+}
-template <int mmq_x, int mmq_y, bool need_check>
-struct mmq_type_traits<mmq_x, mmq_y, need_check, GGML_TYPE_IQ4_XS> {
- static constexpr int vdr = VDR_IQ4_XS_Q8_1_MMQ;
- static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq4_xs<mmq_y, need_check>;
- static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, MMQ_Q8_1_DS_LAYOUT_D4>;
- static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y>;
-};
+// ---------------------------------------------------------------------------------------------
-template <ggml_type type, int mmq_x, bool need_check, bool fixup>
+template <ggml_type type, int J, bool fallback, bool fixup>
static __device__ __forceinline__ void mul_mat_q_process_tile(
const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+ const float * __restrict__ y_scale,
const int stride_row_x, const int ncols_y, const int stride_col_dst,
const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = mmq_get_nwarps_device();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
constexpr int qk = ggml_cuda_type_traits<type>::qk;
- constexpr int mmq_y = get_mmq_y_device();
- constexpr load_tiles_mmq_t load_tiles = mmq_type_traits<mmq_x, mmq_y, need_check, type>::load_tiles;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback>();
+ constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot<type, J, fallback>();
+ constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback>();
extern __shared__ int data_mul_mat_q[];
- int * tile_y = data_mul_mat_q + mmq_x;
- int * tile_x = tile_y + GGML_PAD(mmq_x*MMQ_TILE_Y_K, nwarps*warp_size);
-
-#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
- constexpr vec_dot_mmq_t vec_dot = mmq_type_traits<mmq_x, mmq_y, need_check, type>::vec_dot_mma;
- constexpr mmq_write_back_t write_back = mmq_write_back_mma<type, mmq_x, mmq_y, need_check>;
-#else
- constexpr vec_dot_mmq_t vec_dot = mmq_type_traits<mmq_x, mmq_y, need_check, type>::vec_dot_dp4a;
- constexpr mmq_write_back_t write_back = mmq_write_back_dp4a<mmq_x, mmq_y, need_check>;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ int * tile_y = data_mul_mat_q + J;
+ int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size);
#if defined(BLACKWELL_MMA_AVAILABLE)
// FP4 tile stores 8 blocks
- constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_K : 4 * QK8_1;
+ constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_FP4_MMQ : QK8_1_MMQ;
#else
- constexpr int ne_block = 4 * QK8_1;
+ constexpr int ne_block = QK8_1_MMQ;
#endif // defined(BLACKWELL_MMA_AVAILABLE)
- constexpr int ITER_K = get_iter_k(type);
+ constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback);
constexpr int blocks_per_iter = ITER_K / qk;
- float sum[mmq_x*mmq_y / (nwarps*warp_size)] = {0.0f};
+ float sum[J*I / (nwarps*warp_size)] = {0.0f};
constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
{
const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
#pragma unroll
- for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
+ for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
int l = l0 + threadIdx.y*warp_size + threadIdx.x;
tile_y[l] = by0[l];
{
const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz);
#pragma unroll
- for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
+ for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
int l = l0 + threadIdx.y*warp_size + threadIdx.x;
tile_y[l] = by0[l];
}
if (fixup) {
- write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(mmq_x*mmq_y), mmq_y, mmq_y, mmq_x);
+ write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
} else {
- write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j);
+ write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
}
}
// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
-template <ggml_type type, int mmq_x, bool need_check>
-#if defined(GGML_USE_HIP)
-#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
- __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2)
-#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
-#else
-#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
- __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1)
-#else
- __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2)
-#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
-#endif // defined(GGML_USE_HIP)
+template <ggml_type type, int J, bool fallback>
+__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback))
static __global__ void mul_mat_q(
const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+ const float * __restrict__ y_scale,
const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
const uint3 ntx) {
// Skip unused template specializations for faster compilation:
- if (mmq_x > get_mmq_x_max_device() || mmq_x % mmq_get_granularity_device(mmq_x) != 0) {
+ if (ggml_cuda_mmq_get_config(type, J, fallback).type == GGML_TYPE_COUNT) {
NO_DEVICE_CODE;
return;
}
- constexpr int nwarps = mmq_get_nwarps_device();
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int qk = ggml_cuda_type_traits<type>::qk;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
- constexpr int qk = ggml_cuda_type_traits<type>::qk;
- constexpr int mmq_y = get_mmq_y_device();
-
- const uint32_t nty = (nrows_x + mmq_y - 1) / mmq_y; // Number of tiles y
+ const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y
// Initialize the ids for writing back data with just the index.
// For regular matrix multiplications this is never changed.
// For MoE the correct indices are loaded from ids_dst.
extern __shared__ int ids_dst_shared[]; // Stored at beginning of shared memory.
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) {
+ for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
- if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) {
+ if (j0 + nwarps*warp_size > J && j >= J) {
break;
}
}
__syncthreads();
- // On non-CDNA AMD or old CUDA the performance with stream-k was worse, use conventional tiling instead:
-#if (defined(GGML_USE_HIP) && !defined(CDNA)) || __CUDA_ARCH__ < GGML_CUDA_CC_VOLTA
- {
+ if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback)) {
const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
const int wt = tmp2.x;
const int zt = tmp2.y;
int col_low = 0;
int col_high = ncols_dst;
int col_diff = ncols_dst;
- int offset_y = wt*stride_sample_y + zt*stride_channel_y;
- int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst;
+ int offset_y = wt*stride_sample_y + zt*stride_channel_y;
+ int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+ int offset_y_scale;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+ } else {
+ GGML_UNUSED(offset_y_scale);
+ }
if (ids_dst) {
col_low = expert_bounds[zt + 0];
offset_y = 0;
offset_dst = 0;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = 0;
+ }
- if (jt*mmq_x >= col_diff) {
+ if (jt*J >= col_diff) {
return;
}
// __syncthreads(); // There is no previous tile that could cause a race condition.
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) {
+ for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
- if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) {
+ if (j0 + nwarps*warp_size > J && j >= J) {
break;
}
- ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j];
+ ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
}
__syncthreads();
}
- offset_y += (col_low + jt*mmq_x)*(sizeof(block_q8_1_mmq)/sizeof(int));
- offset_dst += it*mmq_y;
+ offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
+ offset_dst += it*I;
+ const float * y_scale_tile = nullptr;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale += col_low + jt*J;
+ y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+ }
- const int tile_x_max_i = nrows_x - it*mmq_y - 1;
- const int tile_y_max_j = col_diff - jt*mmq_x - 1;
+ const int tile_x_max_i = nrows_x - it*I - 1;
+ const int tile_y_max_j = col_diff - jt*J - 1;
- const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x;
+ const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false;
- mul_mat_q_process_tile<type, mmq_x, need_check, fixup>
- (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+ mul_mat_q_process_tile<type, J, fallback, fixup>
+ (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+ stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
return;
}
-#endif // (defined(GGML_USE_HIP) && !defined(CDNA4) && !defined(CDNA3)) || __CUDA_ARCH__ < GGML_CUDA_CC_VOLTA
- constexpr int ITER_K = get_iter_k(type);
+ constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback);
constexpr int blocks_per_iter = ITER_K / qk;
// kbc == k block continuous, current index in continuous ijk space.
int col_low = 0;
int col_high = ncols_dst;
int col_diff = ncols_dst;
- int offset_y = wt*stride_sample_y + zt*stride_channel_y;
- int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst;
+ int offset_y = wt*stride_sample_y + zt*stride_channel_y;
+ int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+ int offset_y_scale;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+ } else {
+ GGML_UNUSED(offset_y_scale);
+ }
if (ids_dst) {
col_low = expert_bounds[zt + 0];
offset_y = 0;
offset_dst = 0;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = 0;
+ }
- if (jt*mmq_x >= col_diff) {
+ if (jt*J >= col_diff) {
kbc += blocks_per_ne00.z;
kbc -= fastmodulo(kbc, blocks_per_ne00);
__syncthreads();
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) {
+ for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
- if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) {
+ if (j0 + nwarps*warp_size > J && j >= J) {
break;
}
- ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j];
+ ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
}
__syncthreads();
}
- offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int));
- offset_dst += it*mmq_y;
+ offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
+ offset_dst += it*I;
+ const float * y_scale_tile = nullptr;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale += col_low + jt * J;
+ y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+ }
- const int tile_x_max_i = nrows_x - it*mmq_y - 1;
- const int tile_y_max_j = col_diff - jt*mmq_x - 1;
+ const int tile_x_max_i = nrows_x - it*I - 1;
+ const int tile_y_max_j = col_diff - jt*J - 1;
- const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x;
+ const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
- mul_mat_q_process_tile<type, mmq_x, need_check, fixup>
- (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+ mul_mat_q_process_tile<type, J, fallback, fixup>
+ (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+ stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
kbc += blocks_per_ne00.z;
int col_low = 0;
int col_high = ncols_dst;
int col_diff = ncols_dst;
- int offset_y = wt*stride_sample_y + zt*stride_channel_y;
- int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst;
+ int offset_y = wt*stride_sample_y + zt*stride_channel_y;
+ int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+ int offset_y_scale;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+ } else {
+ GGML_UNUSED(offset_y_scale);
+ }
if (ids_dst) {
col_low = expert_bounds[zt + 0];
offset_y = 0;
offset_dst = 0;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale = 0;
+ }
- if (jt*mmq_x >= col_diff) {
+ if (jt*J >= col_diff) {
return;
}
// The memory layout for the fixup buffer is always contiguous, therefore reset ids:
__syncthreads();
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) {
+ for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) {
const int j = j0 + threadIdx.y*warp_size + threadIdx.x;
- if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) {
+ if (j0 + nwarps*warp_size > J && j >= J) {
break;
}
__syncthreads();
}
- offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int));
- offset_dst += it*mmq_y;
+ offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
+ offset_dst += it*I;
+ const float * y_scale_tile = nullptr;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ offset_y_scale += col_low + jt * J;
+ y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+ }
- const int tile_x_max_i = nrows_x - it*mmq_y - 1;
- const int tile_y_max_j = col_diff - jt*mmq_x - 1;
+ const int tile_x_max_i = nrows_x - it*I - 1;
+ const int tile_y_max_j = col_diff - jt*J - 1;
- const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x;
+ const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
- mul_mat_q_process_tile<type, mmq_x, need_check, fixup>
- (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+ mul_mat_q_process_tile<type, J, fallback, fixup>
+ (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+ stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
}
-template <ggml_type type, int mmq_x, bool need_check>
-__launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device()/2, 1)
+template <ggml_type type, int J, bool fallback>
+__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback)/2, 1)
static __global__ void mul_mat_q_stream_k_fixup(
const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst,
const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y,
const int stride_sample_dst, const uint3 ntx) {
- constexpr int mmq_y = get_mmq_y_device();
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback) / 2) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
constexpr int qk = ggml_cuda_type_traits<type>::qk;
- constexpr int ITER_K = get_iter_k(type);
+ constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback);
constexpr int blocks_per_iter = ITER_K / qk;
- constexpr int nwarps = mmq_get_nwarps_device()/2;
- constexpr int warp_size = ggml_cuda_get_physical_warp_size();
-
- float sum[mmq_x / nwarps] = {0.0f};
+ float sum[J / nwarps] = {0.0f};
const int i = blockIdx.y*warp_size + threadIdx.x;
- const int nty = (nrows_x + mmq_y - 1) / mmq_y;
+ const int nty = (nrows_x + I - 1) / I;
const int bidx0 = blockIdx.x;
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
- sum[j0/nwarps] += tmp_last_tile[bidx*(mmq_x*mmq_y) + j*mmq_y + i];
+ sum[j0/nwarps] += tmp_last_tile[bidx*(J*I) + j*I + i];
}
// If this block started in a previous tile we are done and don't need to combine additional partial results.
const int it = tmp2.x;
if (!ids_dst) {
- const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst + it*mmq_y;
+ const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst + it*I;
dst += offset_dst;
- const int i_max = nrows_x - it*mmq_y - 1;
- const int j_max = ncols_dst - jt*mmq_x - 1;
- if (need_check && i > i_max) {
+ const int i_max = nrows_x - it*I - 1;
+ const int j_max = ncols_dst - jt*J - 1;
+ if (fallback && i > i_max) {
return;
}
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
if (j > j_max) {
return;
}
- __shared__ int ids_dst_shared[mmq_x];
+ __shared__ int ids_dst_shared[J];
const int col_low = expert_bounds[zt + 0];
const int col_high = expert_bounds[zt + 1];
const int col_diff = col_high - col_low;
- for (int j = threadIdx.y*warp_size + threadIdx.x; j < mmq_x; j += nwarps*warp_size) {
- ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j];
+ for (int j = threadIdx.y*warp_size + threadIdx.x; j < J; j += nwarps*warp_size) {
+ ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
}
__syncthreads();
- const int offset_dst = it*mmq_y;
+ const int offset_dst = it*I;
dst += offset_dst;
- const int i_max = nrows_x - it*mmq_y - 1;
- const int j_max = col_diff - jt*mmq_x - 1;
- if (need_check && i > i_max) {
+ const int i_max = nrows_x - it*I - 1;
+ const int j_max = col_diff - jt*J - 1;
+ if (fallback && i > i_max) {
return;
}
#pragma unroll
- for (int j0 = 0; j0 < mmq_x; j0 += nwarps) {
+ for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
if (j > j_max) {
struct mmq_args {
const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
+ const float * y_scale;
int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
- bool use_stream_k; int64_t ncols_max;
+ int64_t ncols_max;
};
-template<ggml_type type>
-static size_t mmq_get_nbytes_shared(const int mmq_x, const int mmq_y, const int cc, const int warp_size, const int nwarps) {
- const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, mmq_y);
- const int mmq_tile_x_k = mmq_get_mma_tile_x_k(type);
- const size_t nbs_ids = mmq_x*sizeof(int);
- const size_t nbs_x = (turing_mma_available(cc) || amd_mfma_available(cc) || amd_wmma_available(cc)) ? mmq_y*mmq_tile_x_k*sizeof(int) : txs.qs*sizeof(int) + txs.dm*sizeof(half2) + txs.sc*sizeof(int);
- const size_t nbs_y = mmq_x * (sizeof(block_q8_1_mmq));
- return nbs_ids + nbs_x + GGML_PAD(nbs_y, nwarps*warp_size*sizeof(int));
+static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
+ const size_t nbs_ids = config.J*sizeof(int);
+ const size_t nbs_x = ggml_cuda_mmq_get_nbytes_shared_x(config, cc);
+ const size_t nbs_y = config.J * (sizeof(block_q8_1_mmq));
+ return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int));
}
-template <ggml_type type, int mmq_x>
+template <ggml_type type, int J, bool fallback>
static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
const int id = ggml_cuda_get_device();
const int cc = ggml_cuda_info().devices[id].cc;
const int nsm = ggml_cuda_info().devices[id].nsm;
const int warp_size = ggml_cuda_info().devices[id].warp_size;
- const int nwarps = mmq_get_nwarps_host(cc, warp_size);
- const int mmq_y = get_mmq_y_host(cc);
- const dim3 block_dims(warp_size, nwarps, 1);
+ const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc);
+ GGML_ASSERT(config.nthreads % warp_size == 0);
+ const int nwarps = config.nthreads / warp_size;
+ const int nbytes_shared = mmq_get_nbytes_shared(config, cc);
- const int nbytes_shared = mmq_get_nbytes_shared<type>(mmq_x, mmq_y, cc, warp_size, nwarps);
+ const dim3 block_dims(warp_size, nwarps, 1);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, mmq_x, false>), nbytes_shared);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, mmq_x, true>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true>), nbytes_shared);
- const int nty = (args.nrows_x + mmq_y - 1) / mmq_y;
- const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x;
+ const int nty = (args.nrows_x + config.I - 1) / config.I;
+ const int ntx = (args.ncols_max + config.J - 1) / config.J;
const int ntzw = args.nchannels_y * args.nsamples_y;
const dim3 block_nums_xy_tiling(nty, ntx, ntzw);
const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio);
const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio);
- if (!args.use_stream_k) {
- if (args.nrows_x % mmq_y == 0) {
- constexpr bool need_check = false;
- mul_mat_q<type, mmq_x, need_check><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr,
- blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
- channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
- sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
- ntx_fd);
- } else {
- constexpr bool need_check = true;
- mul_mat_q<type, mmq_x, need_check><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr,
- blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
- channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
- sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
- ntx_fd);
- }
+ if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) {
+ mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
+ (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
+ blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
+ channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
+ sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
+ ntx_fd);
return;
}
ggml_cuda_pool & pool = ctx.pool(id);
ggml_cuda_pool_alloc<float> tmp_fixup(pool);
if (fixup_needed) {
- tmp_fixup.alloc(block_nums_stream_k.x * mmq_x*mmq_y);
+ tmp_fixup.alloc(block_nums_stream_k.x * config.J*config.I);
}
- const dim3 block_nums_fixup(block_nums_stream_k.x, mmq_y/warp_size, 1);
+ const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
- if (args.nrows_x % mmq_y == 0) {
- constexpr bool need_check = false;
- mul_mat_q<type, mmq_x, need_check><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr,
- blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
- channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
- sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
- ntx_fd);
-
- if (!fixup_needed) {
- return;
- }
-
- CUDA_CHECK(cudaGetLastError());
- mul_mat_q_stream_k_fixup<type, mmq_x, need_check><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
- (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
- args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
- ntx_fd);
- } else {
- constexpr bool need_check = true;
- mul_mat_q<type, mmq_x, need_check><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr,
- blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
- channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
- sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
- ntx_fd);
-
- if (!fixup_needed) {
- return;
- }
+ mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
+ (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
+ blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
+ channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
+ sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
+ ntx_fd);
- CUDA_CHECK(cudaGetLastError());
- mul_mat_q_stream_k_fixup<type, mmq_x, need_check><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
- (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
- args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
- ntx_fd);
+ if (!fixup_needed) {
+ return;
}
-}
-template <ggml_type type>
-void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
- const int id = ggml_cuda_get_device();
- const int cc = ggml_cuda_info().devices[id].cc;
- const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
- const int warp_size = ggml_cuda_info().devices[id].warp_size;
- const int nwarps = mmq_get_nwarps_host(cc, warp_size);
+ CUDA_CHECK(cudaGetLastError());
+ mul_mat_q_stream_k_fixup<type, J, fallback><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
+ (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
+ args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
+ ntx_fd);
+}
- const int mmq_x_max = get_mmq_x_max_host(cc);
- const int mmq_y = get_mmq_y_host(cc);
+template <ggml_type type, bool fallback>
+void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+ const int id = ggml_cuda_get_device();
+ const int cc = ggml_cuda_info().devices[id].cc;
+ const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
- int mmq_x_best = 0;
- int ntiles_x_best = INT_MAX;
+ int J_best = 0;
+ int ntiles_J_best = INT_MAX;
- for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) {
- const int granularity = mmq_get_granularity_host(mmq_x, cc);
+ for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
+ const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc);
+ if (config.type == GGML_TYPE_COUNT) {
+ continue;
+ }
- if (mmq_x % granularity != 0 || mmq_get_nbytes_shared<type>(mmq_x, mmq_y, cc, warp_size, nwarps) > smpbo) {
+ if (mmq_get_nbytes_shared(config, cc) > smpbo) {
continue;
}
- const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x;
+ const int ntiles_x = (args.ncols_max + config.J - 1) / config.J;
- if (ntiles_x < ntiles_x_best) {
- mmq_x_best = mmq_x;
- ntiles_x_best = ntiles_x;
+ if (ntiles_x < ntiles_J_best) {
+ J_best = J;
+ ntiles_J_best = ntiles_x;
}
}
- switch (mmq_x_best) {
+ switch (J_best) {
case 8:
- launch_mul_mat_q<type, 8>(ctx, args, stream);
+ launch_mul_mat_q<type, 8, fallback>(ctx, args, stream);
break;
case 16:
- launch_mul_mat_q<type, 16>(ctx, args, stream);
+ launch_mul_mat_q<type, 16, fallback>(ctx, args, stream);
break;
case 24:
- launch_mul_mat_q<type, 24>(ctx, args, stream);
+ launch_mul_mat_q<type, 24, fallback>(ctx, args, stream);
break;
case 32:
- launch_mul_mat_q<type, 32>(ctx, args, stream);
+ launch_mul_mat_q<type, 32, fallback>(ctx, args, stream);
break;
case 40:
- launch_mul_mat_q<type, 40>(ctx, args, stream);
+ launch_mul_mat_q<type, 40, fallback>(ctx, args, stream);
break;
case 48:
- launch_mul_mat_q<type, 48>(ctx, args, stream);
+ launch_mul_mat_q<type, 48, fallback>(ctx, args, stream);
break;
case 56:
- launch_mul_mat_q<type, 56>(ctx, args, stream);
+ launch_mul_mat_q<type, 56, fallback>(ctx, args, stream);
break;
case 64:
- launch_mul_mat_q<type, 64>(ctx, args, stream);
+ launch_mul_mat_q<type, 64, fallback>(ctx, args, stream);
break;
case 72:
- launch_mul_mat_q<type, 72>(ctx, args, stream);
+ launch_mul_mat_q<type, 72, fallback>(ctx, args, stream);
break;
case 80:
- launch_mul_mat_q<type, 80>(ctx, args, stream);
+ launch_mul_mat_q<type, 80, fallback>(ctx, args, stream);
break;
case 88:
- launch_mul_mat_q<type, 88>(ctx, args, stream);
+ launch_mul_mat_q<type, 88, fallback>(ctx, args, stream);
break;
case 96:
- launch_mul_mat_q<type, 96>(ctx, args, stream);
+ launch_mul_mat_q<type, 96, fallback>(ctx, args, stream);
break;
case 104:
- launch_mul_mat_q<type, 104>(ctx, args, stream);
+ launch_mul_mat_q<type, 104, fallback>(ctx, args, stream);
break;
case 112:
- launch_mul_mat_q<type, 112>(ctx, args, stream);
+ launch_mul_mat_q<type, 112, fallback>(ctx, args, stream);
break;
case 120:
- launch_mul_mat_q<type, 120>(ctx, args, stream);
+ launch_mul_mat_q<type, 120, fallback>(ctx, args, stream);
break;
case 128:
- launch_mul_mat_q<type, 128>(ctx, args, stream);
+ launch_mul_mat_q<type, 128, fallback>(ctx, args, stream);
break;
default:
- fprintf(stderr, "mmq_x_best=%d\n", mmq_x_best);
+ fprintf(stderr, "J_best=%d\n", J_best);
GGML_ABORT("fatal error");
break;
}
}
+template <ggml_type type>
+void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+ if (args.nrows_x % 128 == 0) {
+ constexpr bool fallback = false;
+ mul_mat_q_switch_J<type, fallback>(ctx, args, stream);
+ } else {
+ constexpr bool fallback = true;
+ mul_mat_q_switch_J<type, fallback>(ctx, args, stream);
+ }
+}
+
#define DECL_MMQ_CASE(type) \
template void mul_mat_q_case<type>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \
+extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
+extern DECL_MMQ_CASE(GGML_TYPE_Q2_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_1);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_1);
extern DECL_MMQ_CASE(GGML_TYPE_Q8_0);
-extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
-extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
+// -----------------------------------------
extern DECL_MMQ_CASE(GGML_TYPE_Q2_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q3_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_K);
extern DECL_MMQ_CASE(GGML_TYPE_Q6_K);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS);
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XS);
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS);
extern DECL_MMQ_CASE(GGML_TYPE_IQ3_S);
-extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL);
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
+extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
// -------------------------------------------------------------------------------------------------------------------------
void ggml_cuda_mul_mat_q(
ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst);
-void ggml_cuda_op_mul_mat_q(
- ggml_backend_cuda_context & ctx,
- const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i,
- const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
- const int64_t src1_padded_row_size, cudaStream_t stream);
-
bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);
-
static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) {
switch (type) {
case GGML_TYPE_Q1_0: return vec_dot_q1_0_q8_1;
+ case GGML_TYPE_Q2_0: return vec_dot_q2_0_q8_1;
case GGML_TYPE_Q4_0: return vec_dot_q4_0_q8_1;
case GGML_TYPE_Q4_1: return vec_dot_q4_1_q8_1;
case GGML_TYPE_Q5_0: return vec_dot_q5_0_q8_1;
static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) {
switch (type) {
case GGML_TYPE_Q1_0: return VDR_Q1_0_Q8_1_MMVQ;
+ case GGML_TYPE_Q2_0: return VDR_Q2_0_Q8_1_MMVQ;
case GGML_TYPE_Q4_0: return VDR_Q4_0_Q8_1_MMVQ;
case GGML_TYPE_Q4_1: return VDR_Q4_1_Q8_1_MMVQ;
case GGML_TYPE_Q5_0: return VDR_Q5_0_Q8_1_MMVQ;
nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
break;
+ case GGML_TYPE_Q2_0:
+ mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q2_0>
+ (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
+ nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
+ nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream);
+ break;
case GGML_TYPE_Q4_0:
mul_mat_vec_q_switch_ncols_dst<GGML_TYPE_Q4_0>
(vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst,
#include "quantize.cuh"
#include <cstdint>
+#if defined(BLACKWELL_MMA_AVAILABLE)
+// this maps to 256-bit loads in PTX on supported devices,
+// and otherwise falls back to 2 128-bit loads
+struct __builtin_align__(32) float8 {
+ float x; float y; float z; float w;
+ float p; float q; float r; float s;
+};
+#endif
+
+#if CUDART_VERSION >= 12080
+static __device__ __forceinline__ float nvfp4_native_scale_error(
+ const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) {
+ const float scale_dequant = 2.0f * scale;
+ float err = 0.0f;
+
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; k += 4) {
+ const float v0 = vals[k + 0] * inv_col_scale;
+ const float v1 = vals[k + 1] * inv_col_scale;
+ const float v2 = vals[k + 2] * inv_col_scale;
+ const float v3 = vals[k + 3] * inv_col_scale;
+
+ const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale));
+ const __nv_fp4x4_storage_t q_storage = q.__x;
+ const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage);
+ const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U);
+
+ const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1);
+ const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1);
+ const __half2 h2_lo = static_cast<__half2>(hraw2_lo);
+ const __half2 h2_hi = static_cast<__half2>(hraw2_hi);
+ const float2 dq_lo = __half22float2(h2_lo);
+ const float2 dq_hi = __half22float2(h2_hi);
+
+ const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant;
+ const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant;
+ const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant;
+ const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant;
+
+ err = fmaf(err0, err0, err);
+ err = fmaf(err1, err1, err);
+ err = fmaf(err2, err2, err);
+ err = fmaf(err3, err3, err);
+ }
+
+ return err;
+}
+#endif // CUDART_VERSION >= 12080
+
__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
static __global__ void quantize_q8_1(
const float * x_ptr, void * vy_ptr,
return static_cast<uint8_t>(biased);
}
-
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <bool scatter, bool use_aligned_float8>
static __global__ void quantize_mmq_nvfp4(
- const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+ const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
- const int64_t ne0, const int64_t ne1, const int64_t ne2) {
+ const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) {
#if defined(BLACKWELL_MMA_AVAILABLE)
- const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
- if (i0_base >= ne0) {
- return;
- }
+ const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
- const int64_t i1 = blockIdx.x;
- const int64_t i2 = blockIdx.z % ne2;
- const int64_t i3 = blockIdx.z / ne2;
- const int64_t i01 = ids ? ids[i1] : i1;
- const int64_t k_block = i0_base / QK_K;
- const int64_t blocks_per_col = (ne0 + QK_K - 1) / QK_K;
- if (k_block >= blocks_per_col) {
- return;
+ int64_t base_idx;
+ if constexpr (scatter) {
+ base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
+ } else {
+ const int64_t i2 = blockIdx.y % ne2;
+ const int64_t i3 = blockIdx.y / ne2;
+ const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+ base_idx = i3 * s03 + i2 * s02 + i01 * s01;
+ }
+ const float * __restrict__ x_row = x + base_idx;
+
+ float amax = 0.0f;
+ if constexpr (use_aligned_float8) {
+ for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) {
+ const float * x_base = x_row + i0;
+ const float8 v = reinterpret_cast<const float8 *>(x_base)[0];
+ amax = fmaxf(amax, fabsf(v.x));
+ amax = fmaxf(amax, fabsf(v.y));
+ amax = fmaxf(amax, fabsf(v.z));
+ amax = fmaxf(amax, fabsf(v.w));
+ amax = fmaxf(amax, fabsf(v.p));
+ amax = fmaxf(amax, fabsf(v.q));
+ amax = fmaxf(amax, fabsf(v.r));
+ amax = fmaxf(amax, fabsf(v.s));
+ }
+ } else {
+ for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) {
+ amax = fmaxf(amax, fabsf(x_row[i0]));
+ }
}
- const int64_t ib = blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x;
- block_fp4_mmq * y = (block_fp4_mmq *) vy;
- block_fp4_mmq * yb = y + ib;
+ amax = warp_reduce_max<WARP_SIZE>(amax);
- const int sub = (i0_base % QK_K) / QK_NVFP4_SUB;
+ __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE];
+ const int lane = threadIdx.x % WARP_SIZE;
+ const int warp = threadIdx.x / WARP_SIZE;
- float vals_raw[QK_NVFP4_SUB];
- float amax_raw = 0.0f;
- const int64_t base_idx = i3 * s03 + i2 * s02 + i01 * s01;
+ if (lane == 0) {
+ warp_amax[warp] = amax;
+ }
+ __syncthreads();
+
+ if (warp == 0) {
+ amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f;
+ amax = warp_reduce_max<WARP_SIZE>(amax);
+ if (lane == 0) {
+ warp_amax[0] = amax / (6.0f * 448.0f);
+ if constexpr (scatter) {
#pragma unroll
- for (int k = 0; k < QK_NVFP4_SUB; k++) {
- const int64_t i00 = i0_base + k;
- if (i00 < ne00) {
- const float v = x[base_idx + i00];
- vals_raw[k] = v;
- amax_raw = fmaxf(amax_raw, fabsf(v));
- } else {
- vals_raw[k] = 0.0f;
+ for (int slot = 0; slot < n_expert_used; ++slot) {
+ const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+ scale[i] = warp_amax[0];
+ }
+ } else {
+ scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0];
+ }
}
}
+ __syncthreads();
- static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2};
- const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f);
+ block_fp4_mmq * y = (block_fp4_mmq *) vy;
+ const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB;
+
+ for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) {
+ const int64_t i0_base = isb * QK_NVFP4_SUB;
+ const int64_t k_block = i0_base / QK_FP4_MMQ;
+ const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
+
+ const float row_scale = warp_amax[0];
+ const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f;
+
+ float vals[QK_NVFP4_SUB];
+ if constexpr (use_aligned_float8) {
+ const float * x_base = x_row + i0_base;
+ const float8 v0 = i0_base + 7 < ne00 ? reinterpret_cast<const float8 *>(x_base)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+ const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast<const float8 *>(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+ vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w;
+ vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s;
+ vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w;
+ vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s;
+ } else {
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+ const int64_t i00 = i0_base + k;
+ vals[k] = i00 < ne00 ? x_row[i00] : 0.0f;
+ }
+ }
- float best_err = FLT_MAX;
- uint8_t fp8_code = 0;
- float subblock_scale = 0.0f;
+ uint32_t q0 = 0;
+ uint32_t q1 = 0;
-#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell.
- for (int i = 0; i < 5; i++) {
- const int test_code = first_fp8_code + test_offsets[i];
- if (test_code < 0 || test_code > 0x7e) {
- continue;
+ float amax_sub = 0.0f;
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+ amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale));
}
- const uint8_t code = (uint8_t) test_code;
- const float test_scale = ggml_cuda_ue4m3_to_fp32(code);
- const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
- float cur_err = 0.0f;
+
+ static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 };
+ const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f);
+
+ uint8_t fp8_code = (uint8_t) first_fp8_code;
+ float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code);
+ float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+ float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale);
+#else
+ float best_err = 0.0f;
#pragma unroll
for (int k = 0; k < QK_NVFP4_SUB; ++k) {
- const float v = vals_raw[k];
- const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
- const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale;
- cur_err = fmaf(err_diff, err_diff, cur_err);
+ const float v = vals[k] * inv_col_scale;
+ const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err);
+ const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale;
+ best_err = fmaf(err_diff, err_diff, best_err);
}
+#endif // CUDART_VERSION >= 12080
+
+#pragma unroll
+ for (int i = 1; i < 5; ++i) {
+ const int test_code = first_fp8_code + test_offsets[i];
+ if (test_code < 0 || test_code > 0x7e) {
+ continue;
+ }
+
+ const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code);
+ const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+ const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale);
+#else
+ float cur_err = 0.0f;
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+ const float v = vals[k] * inv_col_scale;
+ const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
+ const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale;
+ cur_err = fmaf(err_diff, err_diff, cur_err);
+ }
+#endif // CUDART_VERSION >= 12080
- if (cur_err < best_err) {
- best_err = cur_err;
- fp8_code = test_code;
- subblock_scale = test_scale;
+ if (cur_err < best_err) {
+ best_err = cur_err;
+ fp8_code = (uint8_t) test_code;
+ subblock_scale = test_scale;
+ }
}
- }
+#if CUDART_VERSION >= 12080
+ const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+ const float s = inv_col_scale * inv_scale;
+
+ __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8] * s, vals[1] * s, vals[9] * s));
+ __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s));
+ __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s));
+ __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s));
+
+ const char2 q0_lo_c = *reinterpret_cast<char2 *>(&q0_lo);
+ const char2 q0_hi_c = *reinterpret_cast<char2 *>(&q0_hi);
+ const char2 q1_lo_c = *reinterpret_cast<char2 *>(&q1_lo);
+ const char2 q1_hi_c = *reinterpret_cast<char2 *>(&q1_hi);
+
+ q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) << 8) |
+ (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24);
+ q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) << 8) |
+ (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24);
+#else
+ const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#pragma unroll
+ for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
+ q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k);
+ q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4);
+ q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k);
+ q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4);
+ }
+#endif // CUDART_VERSION >= 12080
- const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
- uint32_t q0 = 0;
- uint32_t q1 = 0;
-#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1
- for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
- q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 0], inv_scale) << (8 * k);
- q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 8], inv_scale) << (8 * k + 4);
- q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 4], inv_scale) << (8 * k);
- q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
+ if constexpr (scatter) {
+#pragma unroll
+ for (int slot = 0; slot < n_expert_used; ++slot) {
+ const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+ block_fp4_mmq * yb = y + (k_block * ne1 + i);
+ uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+ yqs[2 * sub + 0] = q0;
+ yqs[2 * sub + 1] = q1;
+ reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+ }
+ } else {
+ block_fp4_mmq * yb = y + (blockIdx.y * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
+ uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+ yqs[2 * sub + 0] = q0;
+ yqs[2 * sub + 1] = q1;
+ reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+ }
}
-
- uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
- yqs[2 * sub + 0] = q0;
- yqs[2 * sub + 1] = q1;
- reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
#else
+ GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used);
NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
#endif // defined(BLACKWELL_MMA_AVAILABLE)
// quantize values in the format mxfp4 is stored which is interleaved nibbles
// i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <bool scatter>
static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
const int32_t * __restrict__ ids,
void * __restrict__ vy,
const int64_t s03,
const int64_t ne0,
const int ne1,
- const int ne2) {
+ const int ne2,
+ const int n_expert_used) {
constexpr int vals_per_scale = 32;
constexpr int vals_per_warp = 2 * vals_per_scale; // Each warp processes 2 blocks of 32 = 64 values
return;
}
- const int64_t i1 = blockIdx.x;
- const int64_t i2 = blockIdx.z % ne2;
- const int64_t i3 = blockIdx.z / ne2;
-
- ggml_cuda_pdl_sync();
- const int64_t i01 = ids ? ids[i1] : i1;
- const int64_t i02 = i2;
- const int64_t i03 = i3;
-
- block_fp4_mmq * y = (block_fp4_mmq *) vy;
-
- const int64_t block_fp4_mmq_size = 8 * QK_MXFP4; // 256 values
- const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
- const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x;
+ const int64_t block_fp4_mmq_size = QK_FP4_MMQ;
+ const int64_t k_block = warp_start_offset / block_fp4_mmq_size;
const int64_t quad_idx_in_block = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
const int group_id = lane_id_32 / 4;
const int lane_in_group = lane_id_32 % 4;
const int base = group_id * 2;
- char2 * yqs2 = (char2 *) y[ib].qs;
- const int64_t base_pos = i03 * s03 + i02 * s02 + i01 * s01;
+ ggml_cuda_pdl_sync();
+ int64_t base_pos;
+ if constexpr (scatter) {
+ base_pos = (int64_t) blockIdx.x * s02; // one physical row per token
+ } else {
+ const int64_t i2 = blockIdx.z % ne2;
+ const int64_t i3 = blockIdx.z / ne2;
+ const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+ base_pos = i3 * s03 + i2 * s02 + i01 * s01;
+ }
uint8_t scales[2];
+ char2 packed[2];
#pragma unroll
for (int b = 0; b < 2; ++b) {
const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE);
const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE);
- if (lane_in_group == 0) {
- __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
-
- yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = *(char2 *) &fp4_packed;
- }
+ __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
+ packed[b] = *(char2 *) &fp4_packed;
#else
// Fallback: manual FP4 conversion using LUT
const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s);
const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE);
const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE);
- if (lane_in_group == 0) {
- char2 q;
- q.x = (q_hi_0 << 4) | q_lo_0;
- q.y = (q_hi_1 << 4) | q_lo_1;
- yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = q;
- }
+ char2 q;
+ q.x = (q_hi_0 << 4) | q_lo_0;
+ q.y = (q_hi_1 << 4) | q_lo_1;
+ packed[b] = q;
#endif // CUDART_VERSION >= 12080
}
- if (lane_id_32 == 0) {
- // Store 2 scales packed into 1 uint32
- y[ib].d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+ block_fp4_mmq * y = (block_fp4_mmq *) vy;
+ if constexpr (scatter) {
+#pragma unroll
+ for (int slot = 0; slot < n_expert_used; ++slot) {
+ const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+ block_fp4_mmq * yb = y + (k_block * ne1 + i);
+ char2 * yqs2 = (char2 *) yb->qs;
+ if (lane_in_group == 0) {
+ yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
+ yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
+ }
+ if (lane_id_32 == 0) {
+ yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+ }
+ }
+ } else {
+ const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
+ block_fp4_mmq * yb = y + (ib0 + k_block * ne1 + blockIdx.x);
+ char2 * yqs2 = (char2 *) yb->qs;
+ if (lane_in_group == 0) {
+ yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
+ yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
+ }
+ if (lane_id_32 == 0) {
+ yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+ }
}
+ GGML_UNUSED(n_expert_used);
}
-template <mmq_q8_1_ds_layout ds_layout>
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <mmq_q8_1_ds_layout ds_layout, bool scatter>
static __global__ void quantize_mmq_q8_1(
const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
- const int64_t ne0, const int ne1, const int ne2) {
+ const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) {
constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
return;
}
- const int64_t i1 = blockIdx.x;
- const int64_t i2 = blockIdx.z % ne2;
- const int64_t i3 = blockIdx.z / ne2;
-
const int64_t i00 = i0;
ggml_cuda_pdl_sync();
- const int64_t i01 = ids ? ids[i1] : i1;
- const int64_t i02 = i2;
- const int64_t i03 = i3;
- const float4 * x4 = (const float4 *) x;
+ int64_t base_idx;
+ if constexpr (scatter) {
+ base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
+ } else {
+ const int64_t i2 = blockIdx.z % ne2;
+ const int64_t i3 = blockIdx.z / ne2;
+ const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+ base_idx = i3*s03 + i2*s02 + i01*s01;
+ }
+ const float4 * x4 = (const float4 *) x;
block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
- const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
- const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel
- const int64_t iqs = i0 % (4*QK8_1); // quant index in block
+ const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel
+ const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block
// Load 4 floats per thread and calculate max. abs. value between them:
- const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+ const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float amax = fabsf(xi.x);
amax = fmaxf(amax, fabsf(xi.y));
amax = fmaxf(amax, fabsf(xi.z));
q.y = roundf(xi.y*d_inv);
q.z = roundf(xi.z*d_inv);
q.w = roundf(xi.w*d_inv);
+ const float d = 1.0f / d_inv;
- // Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
- char4 * yqs4 = (char4 *) y[ib].qs;
- yqs4[iqs/4] = q;
-
- if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
- if (iqs % 16 != 0 || iqs >= 96) {
- return;
+ // write the block once (normal) or to each of the token's compact rows (scatter)
+ const int nwrite = scatter ? n_expert_used : 1;
+#pragma unroll
+ for (int slot = 0; slot < nwrite; ++slot) {
+ int64_t ib;
+ if constexpr (scatter) {
+ const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+ ib = k_block*ne1 + i;
+ } else {
+ const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
+ ib = ib0 + k_block*ne1 + blockIdx.x;
}
- y[ib].d2s6[2 + iqs/16] = sum;
-
- if (iqs % 64 != 0) {
- return;
+ // Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
+ char4 * yqs4 = (char4 *) y[ib].qs;
+ yqs4[iqs/4] = q;
+
+ if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
+ if (iqs % 16 == 0 && iqs < 96) {
+ y[ib].d2s6[2 + iqs/16] = sum;
+ if (iqs % 64 == 0) {
+ y[ib].d2s6[iqs/64] = d;
+ }
+ }
+ } else if (iqs % 32 == 0) {
+ if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
+ y[ib].ds4[iqs/32] = make_half2(d, sum);
+ } else {
+ y[ib].d4[iqs/32] = d;
+ }
}
-
- const float d = 1.0f / d_inv;
-
- y[ib].d2s6[iqs/64] = d;
-
- return;
- }
-
- if (iqs % 32 != 0) {
- return;
- }
-
- const float d = 1.0f / d_inv;
-
- if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
- y[ib].ds4[iqs/32] = make_half2(d, sum);
- } else {
- y[ib].d4[iqs/32] = d;
}
+ GGML_UNUSED(n_expert_used);
}
void quantize_row_q8_1_cuda(
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
GGML_ASSERT(ne00 % 4 == 0);
- GGML_ASSERT(ne0 % (4*QK8_1) == 0);
+ GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
// ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid:
const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
switch (mmq_get_q8_1_ds_layout(type_src0)) {
case MMQ_Q8_1_DS_LAYOUT_D4:
- quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>
- <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+ quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false>
+ <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
break;
case MMQ_Q8_1_DS_LAYOUT_DS4:
- quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>
- <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+ quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false>
+ <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
break;
case MMQ_Q8_1_DS_LAYOUT_D2S6:
- quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6>
- <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+ quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false>
+ <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
break;
default:
GGML_ABORT("fatal error");
}
}
+// scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
+void quantize_scatter_mmq_q8_1_cuda(
+ const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+ const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+ const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+ GGML_ASSERT(ne00 % 4 == 0);
+ GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
+
+ const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+ const dim3 num_blocks(n_tokens, block_num_y, 1);
+ const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+ switch (mmq_get_q8_1_ds_layout(type_src0)) {
+ case MMQ_Q8_1_DS_LAYOUT_D4:
+ quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+ break;
+ case MMQ_Q8_1_DS_LAYOUT_DS4:
+ quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+ break;
+ case MMQ_Q8_1_DS_LAYOUT_D2S6:
+ quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+ break;
+ default:
+ GGML_ABORT("fatal error");
+ break;
+ }
+}
+
+// scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row)
+void quantize_scatter_mmq_fp4_cuda(
+ const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
+ const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+ const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+ GGML_ASSERT(ne0 > 0);
+ if (type_src0 == GGML_TYPE_NVFP4) {
+ GGML_ASSERT(scale);
+ GGML_ASSERT(ne00 % QK_NVFP4 == 0);
+ const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+ const dim3 num_blocks(n_tokens, 1, 1);
+ if (use_aligned_float8) {
+ quantize_mmq_nvfp4<true, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+ } else {
+ quantize_mmq_nvfp4<true, false><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+ }
+ } else {
+ GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
+ constexpr int nwarps = 8;
+ constexpr int vals_per_block = nwarps * 2 * QK_MXFP4;
+ const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
+ const dim3 block_size(WARP_SIZE, nwarps, 1);
+ const dim3 num_blocks(n_tokens, block_num_y, 1);
+ quantize_mmq_mxfp4<true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+ }
+}
+
void quantize_mmq_fp4_cuda(
- const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+ const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
GGML_ASSERT(ne0 > 0);
if (type_src0 == GGML_TYPE_NVFP4) {
+ GGML_ASSERT(scale);
GGML_ASSERT(ne00 % QK_NVFP4 == 0);
- constexpr int nvfp4_block_size = 128;
- const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
- const dim3 block_size(nvfp4_block_size, 1, 1);
- const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
- quantize_mmq_nvfp4<<<num_blocks, block_size, 0, stream>>>(
- x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+ const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+ const dim3 num_blocks(ne1, ne2 * ne3, 1);
+ if (use_aligned_float8) {
+ quantize_mmq_nvfp4<false, true><<<num_blocks, block_size, 0, stream>>>(
+ x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+ } else {
+ quantize_mmq_nvfp4<false, false><<<num_blocks, block_size, 0, stream>>>(
+ x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+ }
} else {
GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
const dim3 block_size(WARP_SIZE, nwarps, 1);
- quantize_mmq_mxfp4<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+ quantize_mmq_mxfp4<false><<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
}
}
void quantize_mmq_fp4_cuda(const float * x,
const int32_t * ids,
void * vy,
+ float * scale,
ggml_type type_src0,
+ bool use_aligned_float8,
int64_t ne00,
int64_t s01,
int64_t s02,
int64_t ne2,
int64_t ne3,
cudaStream_t stream);
+
+// quantize each token once and scatter the block to its compact rows (via the inverse map)
+void quantize_scatter_mmq_fp4_cuda(const float * x,
+ const int32_t * ids_src1_inv,
+ void * vy,
+ float * scale,
+ ggml_type type_src0,
+ bool use_aligned_float8,
+ int64_t ne00,
+ int64_t stride_token,
+ int64_t ne0,
+ int64_t n_tokens,
+ int64_t nrows_dst,
+ int n_expert_used,
+ cudaStream_t stream);
+
+void quantize_scatter_mmq_q8_1_cuda(const float * x,
+ const int32_t * ids_src1_inv,
+ void * vy,
+ ggml_type type_src0,
+ int64_t ne00,
+ int64_t stride_token,
+ int64_t ne0,
+ int64_t n_tokens,
+ int64_t nrows_dst,
+ int n_expert_used,
+ cudaStream_t stream);
#include "ssm-scan.cuh"
+
+// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path.
+// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences).
+#define SSM_SSD_MIN_TOKENS 128
+
+// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items.
+#define SSM_SSD_DT_BLOCK 256
+#define SSM_SSD_DT_MAX_ITEMS 32
+
+// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity.
+#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS)
+
+// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk.
+#define SSM_SSD_CHUNK_SIZE 256
+
// We would like to keep pragma unroll for cases where L_template is not 0,
// so we suppress the clang transformation warning.
#ifdef __clang__
}
}
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+// ============================================================================
+// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS)
+//
+// Instead of a sequential scan, SSD reformulates the output as:
+// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init
+// where L is a causal decay mask derived from A and dt.
+//
+// This converts the O(T*N) sequential scan into parallel matmuls.
+// ============================================================================
+// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan.
+// Grid: (n_head, n_seqs)
+template <int BLOCK_SIZE, int MAX_ITEMS>
+__global__ void ssm_ssd_prepare_dt_kernel(
+ const float * __restrict__ dt_raw,
+ float * __restrict__ dt_sp_out,
+ float * __restrict__ cs_out,
+ const int n_head, const int n_tok,
+ const int dt_stride_tok, // elements between tokens in dt
+ const int dt_stride_seq) { // elements between sequences in dt
+
+ const int h = blockIdx.x;
+ const int s = blockIdx.y;
+
+ const float * dt_seq = dt_raw + s * dt_stride_seq;
+
+ float * dt_sp_seq = dt_sp_out + s * n_tok * n_head;
+ float * cs_seq = cs_out + s * n_tok * n_head;
+
+ const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+ // Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x).
+ // Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads
+ // (stride n_head between threads vs. items_per_thread*n_head in blocked layout).
+ float local_vals[MAX_ITEMS];
+ for (int i = 0; i < items_per_thread; i++) {
+ const int t = i * BLOCK_SIZE + threadIdx.x;
+ if (t < n_tok) {
+ float val = dt_seq[h + t * dt_stride_tok];
+ float sp = (val <= 20.0f) ? log1pf(expf(val)) : val;
+ local_vals[i] = sp;
+ dt_sp_seq[t * n_head + h] = sp;
+ } else {
+ local_vals[i] = 0.0f;
+ }
+ }
+
+ // Phase 2+3: per-step inclusive scan to build cs[] in token order.
+ // With interleaved distribution the per-thread total scan would not give token-order
+ // prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total.
+#ifdef USE_CUB
+ using BlockScan = cub::BlockScan<float, BLOCK_SIZE>;
+ __shared__ typename BlockScan::TempStorage scan_temp;
+ __shared__ float step_total;
+
+ float running = 0.0f;
+ for (int i = 0; i < items_per_thread; i++) {
+ float inclusive;
+ BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive);
+ const int t = i * BLOCK_SIZE + threadIdx.x;
+ if (t < n_tok) {
+ cs_seq[t * n_head + h] = running + inclusive;
+ }
+ if (threadIdx.x == BLOCK_SIZE - 1) {
+ step_total = inclusive;
+ }
+ __syncthreads();
+ running += step_total;
+ }
+#else
+ // Fallback: sequential prefix scan in shared memory, one slab at a time.
+ __shared__ float sdata[BLOCK_SIZE];
+ float running = 0.0f;
+ for (int i = 0; i < items_per_thread; i++) {
+ const int t = i * BLOCK_SIZE + threadIdx.x;
+ sdata[threadIdx.x] = local_vals[i];
+ __syncthreads();
+ if (threadIdx.x == 0) {
+ for (int j = 1; j < BLOCK_SIZE; j++) {
+ sdata[j] += sdata[j - 1];
+ }
+ }
+ __syncthreads();
+ if (t < n_tok) {
+ cs_seq[t * n_head + h] = running + sdata[threadIdx.x];
+ }
+ running += sdata[BLOCK_SIZE - 1];
+ __syncthreads();
+ }
+#endif
+}
+
+// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled.
+// T_matmul controls precision for X_dt, B_weighted (float or half).
+// C_scaled is always float (pairs with float s_cur in step 3c).
+// Computation is always FP32; only the final store converts to T_matmul.
+// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch).
+// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs)
+template <int BLOCK_SIZE, typename T_matmul>
+__global__ void ssm_ssd_pre_matmul_kernel(
+ const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums
+ const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt)
+ const float * __restrict__ A, // {1, n_head}
+ const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs}
+ const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs}
+ const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs}
+ T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest
+ T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end
+ float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float)
+ const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs}
+ half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs}
+ const int chunk_len, const int head_dim, const int n_head, const int n_group,
+ const int d_state, const int A_stride,
+ const int x_stride_tok, const int x_stride_seq,
+ const int B_stride_tok, const int B_stride_seq,
+ const int C_stride_tok, const int C_stride_seq,
+ const int chunk_offset,
+ const int n_tok_total) {
+
+ const int h = blockIdx.y;
+ const int s = blockIdx.z;
+ const int g = h / (n_head / n_group);
+
+ const float A_h = A[h * A_stride];
+ const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
+
+ const int cs_seq_off = s * n_tok_total * n_head;
+ const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f;
+ const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base;
+
+ // Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes.
+ const int n_xdt = chunk_len * head_dim;
+ if (idx < n_xdt) {
+ const int d = idx % head_dim;
+ const int t = idx / head_dim;
+
+ const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim];
+ const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h];
+
+ X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val);
+ }
+
+ // Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len)
+ // and the same cs_t load, so merging halves the cs[] global memory traffic.
+ const int n_bw = d_state * chunk_len;
+ if (idx < n_bw) {
+ const int n = idx % d_state;
+ const int t = idx / d_state;
+
+ const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base;
+
+ const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n];
+ B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t)));
+
+ const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n];
+ C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t);
+ }
+
+ // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask.
+ const int n_M = chunk_len * chunk_len;
+ if (idx < n_M) {
+ const int t_out = idx % chunk_len;
+ const int t_in = idx / chunk_len;
+
+ half val;
+ if (t_in <= t_out) {
+ const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base;
+ const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base;
+ const float decay = __expf(A_h * (cs_out - cs_in));
+ const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group
+ + (int64_t)g * chunk_len * chunk_len;
+ const float cb_val = CB_g[t_out + t_in * chunk_len];
+ val = __float2half(decay * cb_val);
+ } else {
+ val = __float2half(0.0f);
+ }
+
+ M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val;
+ }
+}
+
+// Scale running state in-place: s_cur *= decay_total(chunk).
+// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay.
+// Eliminates the s_old buffer and D2D memcpy vs the old approach of:
+// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old
+// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs)
+template <int BLOCK_SIZE>
+__global__ void ssm_ssd_scale_state_kernel(
+ float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs}
+ const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums
+ const float * __restrict__ A, // {1, n_head}
+ const int d_state, const int head_dim, const int n_head,
+ const int chunk_offset, const int chunk_len,
+ const int n_tok_total, const int A_stride) {
+
+ const int h = blockIdx.y;
+ const int s = blockIdx.z;
+ const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
+ const int state_per_head = d_state * head_dim;
+ if (idx >= state_per_head) return;
+
+ const float A_h = A[h * A_stride];
+ const int cs_seq_off = s * n_tok_total * n_head;
+ const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f;
+ const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base;
+ const float decay_total = __expf(A_h * cs_last);
+
+ const int off = s * state_per_head * n_head + h * state_per_head + idx;
+ s_cur[off] *= decay_total;
+}
+
+// Copy initial state from src0[ids[s]] into s_cur for each sequence.
+// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs)
+template <int BLOCK_SIZE>
+__global__ void ssm_ssd_init_state_kernel(
+ const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs}
+ const int32_t * __restrict__ ids, // {n_seqs}
+ float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs}
+ const int state_size, // d_state * head_dim * n_head
+ const int64_t s0_stride_seq) { // elements between state rows
+ const int s = blockIdx.y;
+ const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
+ if (idx >= state_size) return;
+
+ const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq;
+ s_cur[s * state_size + idx] = s_src[idx];
+}
+
+// SSD (State Space Duality) dispatch for Mamba-2 prefill.
+// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt.
+// All strides are in elements (floats), not bytes.
+static void ssm_scan_ssd_f32_cuda(
+ ggml_backend_cuda_context & ctx,
+ const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d,
+ const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d,
+ const int64_t s0_stride_seq, // state (src0) stride between seqs
+ const int x_stride_tok, const int x_stride_seq, // x (src1) strides
+ const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides
+ const int A_stride, // A (src3) stride between heads
+ const int B_stride_tok, const int B_stride_seq, // B (src4) strides
+ const int C_stride_tok, const int C_stride_seq, // C (src5) strides
+ const int64_t s_off, const int64_t d_state, const int64_t head_dim,
+ const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) {
+
+ cudaStream_t stream = ctx.stream();
+ const int64_t d_inner = head_dim * n_head;
+
+ const int64_t chunk_size = SSM_SSD_CHUNK_SIZE;
+ const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size;
+
+ const int64_t state_per_head = d_state * head_dim;
+
+ using matmul_t = half;
+ static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F;
+
+ ggml_cuda_pool_alloc<float> dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq);
+ ggml_cuda_pool_alloc<float> cs_buf(ctx.pool(), n_tok * n_head * n_seq);
+ ggml_cuda_pool_alloc<float> CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq);
+ ggml_cuda_pool_alloc<matmul_t> X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq);
+ ggml_cuda_pool_alloc<matmul_t> B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq);
+ ggml_cuda_pool_alloc<float> C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq);
+ float * dt_sp = dt_sp_buf.get();
+ float * cs = cs_buf.get();
+ float * CB = CB_buf.get();
+ matmul_t * X_dt = X_dt_buf.get();
+ matmul_t * B_weighted = B_w_buf.get();
+ float * C_scaled = C_s_buf.get();
+ float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst
+
+ // Step 1: softplus(dt) and parallel prefix sum over full sequence
+ {
+ dim3 grid(n_head, n_seq);
+ ssm_ssd_prepare_dt_kernel<SSM_SSD_DT_BLOCK, SSM_SSD_DT_MAX_ITEMS><<<grid, SSM_SSD_DT_BLOCK, 0, stream>>>(
+ src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq);
+ CUDA_CHECK(cudaGetLastError());
+ }
+
+ // Step 2: initialize running state from src0[ids[s]]
+ {
+ constexpr int BLOCK = 256;
+ const int64_t state_size = d_state * head_dim * n_head;
+ dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq);
+ ssm_ssd_init_state_kernel<BLOCK><<<grid, BLOCK, 0, stream>>>(
+ src0_d, src6_d, s_cur, state_size, s0_stride_seq);
+ CUDA_CHECK(cudaGetLastError());
+ }
+
+ // Step 3: chunked SSD loop
+ // Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state
+ cublasHandle_t handle = ctx.cublas_handle();
+ CUBLAS_CHECK(cublasSetStream(handle, stream));
+ const float alpha_one = 1.0f;
+ const float beta_zero = 0.0f;
+ const float beta_one = 1.0f;
+ const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B
+ const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B
+
+ // Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size)
+ const int64_t n_M_max = chunk_size * chunk_size;
+ ggml_cuda_pool_alloc<half> M_buf(ctx.pool(), n_M_max * n_head * n_seq);
+ half * M_mat = M_buf.get();
+
+ for (int64_t k = 0; k < n_chunks; k++) {
+ const int64_t chunk_offset = k * chunk_size;
+ const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset);
+
+ // 3a: CB = C^T @ B per group
+ for (int64_t s = 0; s < n_seq; s++) {
+ const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok;
+ const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok;
+ float * CB_s = CB + s * chunk_len * chunk_len * n_group;
+
+ if (n_group == 1) {
+ CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,
+ chunk_len, chunk_len, d_state,
+ &alpha_one, C_s, lda_C_src, B_s, ldb_B_src,
+ &beta_zero, CB_s, (int)chunk_len));
+ } else {
+ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
+ chunk_len, chunk_len, d_state,
+ &alpha_one,
+ C_s, CUDA_R_32F, lda_C_src, d_state,
+ B_s, CUDA_R_32F, ldb_B_src, d_state,
+ &beta_zero,
+ CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len),
+ n_group,
+ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
+ }
+ }
+
+ // 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix
+ const int64_t n_M = chunk_len * chunk_len;
+ {
+ constexpr int BLOCK = 256;
+ const int64_t n_xdt = chunk_len * head_dim;
+ const int64_t n_bw = d_state * chunk_len;
+ int64_t max_work = n_xdt;
+ if (n_bw > max_work) max_work = n_bw;
+ if (n_M > max_work) max_work = n_M;
+ dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq);
+ ssm_ssd_pre_matmul_kernel<BLOCK, matmul_t><<<grid, BLOCK, 0, stream>>>(
+ cs, dt_sp, src3_d, src1_d, src4_d, src5_d,
+ X_dt, B_weighted, C_scaled,
+ CB, M_mat,
+ chunk_len, head_dim, n_head, n_group, d_state, A_stride,
+ x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq,
+ chunk_offset, n_tok);
+ CUDA_CHECK(cudaGetLastError());
+ }
+
+ // 3c: dst = S_cur^T @ C_scaled (state contribution)
+ {
+ const int64_t stride_S = state_per_head;
+ const int64_t stride_Cs = d_state * chunk_len;
+
+ for (int64_t s = 0; s < n_seq; s++) {
+ float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner;
+
+ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
+ head_dim, chunk_len, d_state,
+ &alpha_one,
+ s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S,
+ C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs,
+ &beta_zero,
+ dst_chunk, CUDA_R_32F, d_inner, head_dim,
+ n_head,
+ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
+ }
+ }
+
+ // 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result)
+ // M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T.
+ {
+ const int64_t stride_M = n_M;
+ const int64_t stride_X_h = (int64_t)chunk_len * head_dim;
+
+ for (int64_t s = 0; s < n_seq; s++) {
+ float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner;
+ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T,
+ head_dim, chunk_len, chunk_len,
+ &alpha_one,
+ X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h,
+ M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M,
+ &beta_one,
+ dst_chunk, CUDA_R_32F, d_inner, head_dim,
+ n_head,
+ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
+ }
+ }
+
+ // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update)
+ {
+ // Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it
+ constexpr int BLOCK = 256;
+ dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq);
+ ssm_ssd_scale_state_kernel<BLOCK><<<grid, BLOCK, 0, stream>>>(
+ s_cur, cs, src3_d,
+ d_state, head_dim, n_head,
+ chunk_offset, chunk_len, n_tok, A_stride);
+ CUDA_CHECK(cudaGetLastError());
+
+ // cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled)
+ const int64_t stride_Bw = d_state * chunk_len;
+ const int64_t stride_X = chunk_len * head_dim;
+ const int64_t stride_S = state_per_head;
+
+ for (int64_t s = 0; s < n_seq; s++) {
+ // X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd}
+ CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T,
+ d_state, head_dim, chunk_len,
+ &alpha_one,
+ B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw,
+ X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X,
+ &beta_one,
+ s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S,
+ n_head,
+ CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
+ }
+ }
+ }
+}
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const struct ggml_tensor * src0 = dst->src[0]; // s
const struct ggml_tensor * src1 = dst->src[1]; // x
GGML_ASSERT(src6->type == GGML_TYPE_I32);
GGML_ASSERT(dst->type == GGML_TYPE_F32);
+ // Byte strides are narrowed to int for both scan and SSD paths.
+ GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX);
+ GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX);
+ GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX);
+ GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX);
+ GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX);
+ GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX);
+ GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX);
+ GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX);
+ GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX);
+ GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX);
+ GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX);
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ // Mamba-2 with scalar A per head: use SSD matmul path for long sequences.
+ // Requires NVIDIA Turing+ otherwise fallback to scan.
+ const bool is_mamba2 = (src3->nb[1] == sizeof(float));
+ const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+ const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS
+ && n_t <= SSM_SSD_MAX_TOKENS
+ && GGML_CUDA_CC_IS_NVIDIA(cc)
+ && cc >= GGML_CUDA_CC_TURING
+ && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment
+
+ if (use_ssd) {
+ // ssm_ssd_init_state_kernel uses flat linear indexing within each sequence,
+ // so src0 must be fully contiguous across all inner dimensions.
+ // The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1].
+ GGML_ASSERT(src0->nb[1] == nc * sizeof(float));
+ GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float));
+
+ ssm_scan_ssd_f32_cuda(ctx,
+ src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
+ (int64_t)(src0->nb[3] / sizeof(float)),
+ (int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)),
+ (int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)),
+ (int)(src3->nb[1] / sizeof(float)),
+ (int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)),
+ (int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)),
+ s_off, nc, nr, nh, ng, n_t, n_s);
+ return;
+ }
+#endif
ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
TYPES_MMQ = [
"GGML_TYPE_Q1_0",
+ "GGML_TYPE_Q2_0",
"GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0",
"GGML_TYPE_Q2_K", "GGML_TYPE_Q3_K", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", "GGML_TYPE_Q6_K",
"GGML_TYPE_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S",
--- /dev/null
+// This file has been autogenerated by generate_cu_files.py, do not edit manually.
+
+#include "../mmq.cuh"
+
+DECL_MMQ_CASE(GGML_TYPE_Q2_0);
const int ncols_pad = next_power_of_2(ncols);
const size_t shared_mem = ncols_pad * sizeof(int);
const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;
+ const bool use_bitonic = shared_mem <= max_shared_mem && ncols <= 1024;
+ const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows);
- ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
+ ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * chunk_nrows);
int * tmp_dst = temp_dst_alloc.get();
- if (shared_mem > max_shared_mem || ncols > 1024) {
- argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
- } else {
- argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
+ for (int64_t i = 0; i < nrows; i += chunk_nrows) {
+ int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i);
+
+ if (use_bitonic) {
+ argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream);
+ } else {
+ argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream);
+ }
+ CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), iter_nrows,
+ cudaMemcpyDeviceToDevice, stream));
+
+ src0_d += ncols * iter_nrows;
+ dst_d += k * iter_nrows;
}
- CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
- cudaMemcpyDeviceToDevice, stream));
#else // GGML_CUDA_USE_CUB
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
int * tmp_dst = temp_dst_alloc.get();
// Kernel config struct - passed by value to CUDA kernel
struct topk_moe_config {
bool use_sigmoid;
+ bool use_sqrt_softplus;
bool with_norm;
bool delayed_softmax;
};
}
}
+template <int experts_per_thread, bool use_limit>
+__device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ const int idx = lane + i * WARP_SIZE;
+ const bool active = !use_limit || (idx < limit);
+ vals[i] = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
+ }
+}
+
/*
This kernel does the following:
1. optionally softmax over the logits per token [n_experts, n_tokens]
if (!config.delayed_softmax) {
if (config.use_sigmoid) {
sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
+ } else if (config.use_sqrt_softplus) {
+ sqrt_softplus_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
} else {
softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
}
}
topk_moe_config config;
- config.use_sigmoid = args.sigmoid;
- config.with_norm = with_norm;
- config.delayed_softmax = args.delayed_softmax;
+ config.use_sigmoid = args.sigmoid;
+ config.use_sqrt_softplus = args.sqrt_softplus;
+ config.with_norm = with_norm;
+ config.delayed_softmax = args.delayed_softmax;
if (bias) {
launch_topk_moe_cuda<true>(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val,
} else if (gating_op->op == GGML_OP_UNARY) {
ggml_unary_op op = ggml_get_unary_op(gating_op);
- if (op != GGML_UNARY_OP_SIGMOID) {
+ if (op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SOFTPLUS) {
return false;
}
}
struct ggml_cuda_topk_moe_args {
bool sigmoid{};
+ bool sqrt_softplus{};
bool softmax{};
bool delayed_softmax{};
bool prob_bias{};
#define VDR_Q1_0_Q8_1_MMVQ 1 // Process one 32-element chunk at a time for parallelism
#define VDR_Q1_0_Q8_1_MMQ 4 // Q1_0 has 128 bits (4 ints) per block
+#define VDR_Q2_0_Q8_1_MMVQ 1 // Process one 32-element chunk at a time for parallelism
+#define VDR_Q2_0_Q8_1_MMQ 2 // Q2_0 group 64: 128 bits (4 ints) per block, 2 32-element chunks
+
#define VDR_Q4_0_Q8_1_MMVQ 2
#define VDR_Q4_0_Q8_1_MMQ 4
// Q8_1: 32 elements per block with individual scales
// iqs selects which of the 4 chunks of 32 elements to process (0-3)
- const float d1 = bq1_0->d;
+ const float d1 = bq1_0->d;
+ const int16_t * qs = (const int16_t *) bq1_0->qs + iqs * 2;
+
+ // Process only the chunk specified by iqs
+ const block_q8_1 * bq8_1_chunk = bq8_1 + iqs;
+
+ int sumi = 0;
+#pragma unroll
+ for (int j = 0; j < 2; ++j) {
+ const int q = qs[j];
+
+ const int u0 = get_int_b4(bq8_1_chunk->qs, j*4+0);
+ const int u1 = get_int_b4(bq8_1_chunk->qs, j*4+1);
+ const int u2 = get_int_b4(bq8_1_chunk->qs, j*4+2);
+ const int u3 = get_int_b4(bq8_1_chunk->qs, j*4+3);
+
+ // unpack crumbs into nibble indices
+ const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8, 9, 12, 13]
+ const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15]
+ // unpack nibbles into byte values
+ const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >> 0);
+ const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >> 0);
+ const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16);
+ const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16);
+ // unshuffle values
+ const int v0 = __byte_perm(s0, s1, 0x5410);
+ const int v1 = __byte_perm(s0, s1, 0x7632);
+ const int v2 = __byte_perm(s2, s3, 0x5410);
+ const int v3 = __byte_perm(s2, s3, 0x7632);
+
+ sumi = ggml_cuda_dp4a(v0, u0, sumi);
+ sumi = ggml_cuda_dp4a(v1, u1, sumi);
+ sumi = ggml_cuda_dp4a(v2, u2, sumi);
+ sumi = ggml_cuda_dp4a(v3, u3, sumi);
+ }
+
+ // Apply Q1_0's single scale and this chunk's Q8_1 scale
+ const float d8 = __low2float(bq8_1_chunk->ds);
+ return d1 * d8 * sumi;
+}
+
+static __device__ __forceinline__ float vec_dot_q2_0_q8_1(
+ const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {
+
+ const block_q2_0 * bq2_0 = (const block_q2_0 *) vbq + kbx;
+
+ // Q2_0 (group 64): 64 elements with ONE scale, 2 bits per element (4 elements per byte)
+ // Q8_1: 32 elements per block with individual scales
+ // iqs selects which of the 2 chunks of 32 elements to process (0-1)
+
+ const float d2 = bq2_0->d;
// Process only the chunk specified by iqs
const block_q8_1 * bq8_1_chunk = bq8_1 + iqs;
- // Load 32 bits (4 bytes) for this chunk from Q1_0
- const int offset = iqs * 4;
- const int v = bq1_0->qs[offset + 0] | (bq1_0->qs[offset + 1] << 8) |
- (bq1_0->qs[offset + 2] << 16) | (bq1_0->qs[offset + 3] << 24);
+ // Load 64 bits (8 bytes) for this chunk from Q2_0: bytes [8*iqs, 8*iqs+8)
+ const int offset = iqs * 8;
+ const int v0 = bq2_0->qs[offset + 0] | (bq2_0->qs[offset + 1] << 8) |
+ (bq2_0->qs[offset + 2] << 16) | (bq2_0->qs[offset + 3] << 24);
+ const int v1 = bq2_0->qs[offset + 4] | (bq2_0->qs[offset + 5] << 8) |
+ (bq2_0->qs[offset + 6] << 16) | (bq2_0->qs[offset + 7] << 24);
- // Unpack 32 bits into 32 signed values (-1 or +1)
+ // Unpack 32 2-bit codes into 8 int32s, each holding 4 signed int8 symbols in {-1,0,1,2}.
+ // Stored code c in {0,1,2,3} -> symbol s = c - 1.
int vi_bytes[8];
#pragma unroll
- for (int j = 0; j < 8; ++j) {
- const int shift = j * 4;
- const int bits4 = (v >> shift) & 0x0F;
- const int b0 = (bits4 & 0x01) ? 1 : -1;
- const int b1 = (bits4 & 0x02) ? 1 : -1;
- const int b2 = (bits4 & 0x04) ? 1 : -1;
- const int b3 = (bits4 & 0x08) ? 1 : -1;
- vi_bytes[j] = (b0 & 0xFF) | ((b1 & 0xFF) << 8) | ((b2 & 0xFF) << 16) | ((b3 & 0xFF) << 24);
+ for (int j = 0; j < 4; ++j) {
+ const int shift = j * 8;
+ const int codes = (v0 >> shift) & 0xFF;
+ const int c0 = ((codes >> 0) & 0x3) - 1;
+ const int c1 = ((codes >> 2) & 0x3) - 1;
+ const int c2 = ((codes >> 4) & 0x3) - 1;
+ const int c3 = ((codes >> 6) & 0x3) - 1;
+ vi_bytes[j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
+ }
+#pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ const int shift = j * 8;
+ const int codes = (v1 >> shift) & 0xFF;
+ const int c0 = ((codes >> 0) & 0x3) - 1;
+ const int c1 = ((codes >> 2) & 0x3) - 1;
+ const int c2 = ((codes >> 4) & 0x3) - 1;
+ const int c3 = ((codes >> 6) & 0x3) - 1;
+ vi_bytes[4 + j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
}
// Compute dot product for this 32-element chunk
sumi = ggml_cuda_dp4a(vi_bytes[j], u, sumi);
}
- // Apply Q1_0's single scale and this chunk's Q8_1 scale
+ // Apply Q2_0's single scale and this chunk's Q8_1 scale
const float d8 = __low2float(bq8_1_chunk->ds);
- return d1 * d8 * sumi;
+ return d2 * d8 * sumi;
}
static __device__ __forceinline__ float vec_dot_q4_0_q8_1(
#include <hip/hip_fp16.h>
#include <hip/hip_bf16.h>
-#if defined(GGML_HIP_ROCWMMA_FATTN)
-#include <rocwmma/rocwmma-version.hpp>
-#endif // defined(GGML_HIP_ROCWMMA_FATTN)
-
#ifdef GGML_USE_NCCL
#include <rccl/rccl.h>
#endif // GGML_USE_NCCL
--- /dev/null
+
+message(STATUS "Using ET backend")
+
+# Configure ET platform path
+if (DEFINED ENV{ET_PLATFORM})
+ set(ET_PLATFORM_PATH $ENV{ET_PLATFORM})
+else()
+ set(ET_PLATFORM_PATH "/opt/et")
+endif()
+
+# Use sysemu for ET backend if compiled with `-DGGML_ET_SYSEMU=ON`
+if (GGML_ET_SYSEMU)
+ message(STATUS "Using ET backend with sysemu instead of hardware")
+else()
+ message(STATUS "Using ET backend with hardware device")
+endif()
+
+# Add ET platform CMake modules and config files to search paths
+list(APPEND CMAKE_PREFIX_PATH ${ET_PLATFORM_PATH}/lib/cmake)
+list(APPEND CMAKE_MODULE_PATH ${ET_PLATFORM_PATH}/lib/cmake)
+include(aifoundry-utils/ProjectFunctions)
+
+message(STATUS "Using ET Platform at ${ET_PLATFORM_PATH}")
+
+find_package(runtime REQUIRED)
+
+# Kernel list
+set(KERNELS
+ el_map_f32
+ flash_attn_ext_f32
+ glu_f32
+ scale_f32
+ mul_mat_f32
+ mul_mat_f32_matrix_engine
+ mul_mat_id_f32
+ mul_mat_id_Q4_0
+ mul_mat_id_Q8_0
+ mul_mat_Q8_0
+ mul_mat_Q4_0
+ mul_mat_Q4_0_matrix_engine
+ mul_mat_f16
+ mul_mat_f16_matrix_engine
+ rope_f32
+ unary_f32
+ sqr_f32
+ clamp_f32
+ sum_rows_f32
+ mean_f32
+ cumsum_f32
+ norm_f32
+ l2_norm_f32
+ group_norm_f32
+ rms_norm_f32
+ rms_norm_mul_f32
+ softmax_f32
+ im2col
+ get_rows_f32
+ concat_f32
+ repeat_f32
+ rwkv_wkv6_f32
+ rwkv_wkv7_f32
+ gated_delta_net_f32
+ cont_f32
+ cont_f16
+ cpy_f32_f16
+ flash_attn_ext_f16_me
+ set_rows_f32
+ set_f32
+ fill_f32
+ pad_f32
+ diag_f32
+ tri_f32
+ solve_tri_f32
+ ssm_conv_f32
+ ssm_scan_f32
+ conv_2d_f32_me
+ memops
+ uberkernel
+)
+
+# Kernels that we support dispatch form Uberkernel
+set(UBERKERNEL_SUPPORTED_KERNELS
+ el_map_f32
+ # unary_f32
+ # cpy_f32_f16
+ # cont_f32
+ # get_rows_f32
+ concat_f32
+ cont_f16
+ cumsum_f32
+ diag_f32
+ fill_f32
+ flash_attn_ext_f16_me
+ flash_attn_ext_f32
+ gated_delta_net_f32
+ glu_f32
+ group_norm_f32
+ im2col
+ l2_norm_f32
+ mul_mat_f16
+ mul_mat_f16_matrix_engine
+ mul_mat_f32
+ mul_mat_f32_matrix_engine
+ mul_mat_id_f32
+ mul_mat_Q4_0
+ mul_mat_Q8_0
+ norm_f32
+ pad_f32
+ repeat_f32
+ rms_norm_f32
+ rms_norm_mul_f32
+ rope_f32
+ rwkv_wkv6_f32
+ rwkv_wkv7_f32
+ scale_f32
+ set_f32
+ set_rows_f32
+ softmax_f32
+ solve_tri_f32
+ sqr_f32
+ # ssm_conv_f32
+ ssm_scan_f32
+ sum_rows_f32
+ tri_f32
+)
+
+set(UBERKERNEL_MAP_HPP ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-uberkernel-kernel-map.h)
+set(UBERKERNEL_MAP_CPP ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-uberkernel-kernel-map.cpp)
+
+set(UBERKERNEL_KERNELS_SORTED ${UBERKERNEL_SUPPORTED_KERNELS})
+list(SORT UBERKERNEL_KERNELS_SORTED)
+
+set(UBERKERNEL_ENUM_ENTRIES "")
+set(UBERKERNEL_MAP_ENTRIES "")
+set(_uk_idx 1)
+foreach(KERNEL ${UBERKERNEL_KERNELS_SORTED})
+ string(TOUPPER ${KERNEL} _uk_upper)
+ string(APPEND UBERKERNEL_ENUM_ENTRIES
+ " GGML_ET_UBERKERNEL_KERNEL_${_uk_upper} = ${_uk_idx},\n")
+ string(APPEND UBERKERNEL_MAP_ENTRIES
+ " {\"${KERNEL}\", GGML_ET_UBERKERNEL_KERNEL_${_uk_upper}},\n")
+ math(EXPR _uk_idx "${_uk_idx} + 1")
+endforeach()
+
+configure_file(
+ ${CMAKE_CURRENT_SOURCE_DIR}/cmake/ggml-et-uberkernel-kernel-map.h.in
+ ${UBERKERNEL_MAP_HPP}
+ @ONLY)
+configure_file(
+ ${CMAKE_CURRENT_SOURCE_DIR}/cmake/ggml-et-uberkernel-kernel-map.cpp.in
+ ${UBERKERNEL_MAP_CPP}
+ @ONLY)
+
+add_custom_target(et-uberkernel-map
+ DEPENDS ${UBERKERNEL_MAP_HPP} ${UBERKERNEL_MAP_CPP}
+)
+
+# Build ET kernels (cross-compiled in subdirectory scope)
+add_subdirectory(et-kernels)
+
+# Embed kernels into C++ source
+set(EMBED_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_one_kernel.cmake)
+set(EMBED_HPP ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-kernels-embed.hpp)
+set(EMBED_CPP ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-kernels-embed.cpp)
+set(EMBED_DIR ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/embed)
+file(MAKE_DIRECTORY ${EMBED_DIR})
+
+set(EMBED_KERNEL_SOURCES)
+set(EMBED_EXTERNS "")
+set(EMBED_MAP_ENTRIES "")
+
+foreach(KERNEL ${KERNELS})
+ set(ELF_PATH ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/${KERNEL}.elf)
+ set(OUT_CPP ${EMBED_DIR}/${KERNEL}.cpp)
+
+ add_custom_command(
+ OUTPUT ${OUT_CPP}
+ COMMAND ${CMAKE_COMMAND}
+ -DELF_FILE=${ELF_PATH}
+ -DOUT_FILE=${OUT_CPP}
+ -DVAR_NAME=${KERNEL}
+ -P ${EMBED_SCRIPT}
+ DEPENDS ${KERNEL}.elf ${EMBED_SCRIPT}
+ COMMENT "Embedding ${KERNEL}.elf"
+ VERBATIM
+ )
+ list(APPEND EMBED_KERNEL_SOURCES ${OUT_CPP})
+
+ string(APPEND EMBED_EXTERNS
+ "extern unsigned char ${KERNEL}_data[];\n"
+ "extern const uint64_t ${KERNEL}_len;\n")
+ string(APPEND EMBED_MAP_ENTRIES
+ " {\"${KERNEL}\", {${KERNEL}_data, ${KERNEL}_len}},\n")
+endforeach()
+
+configure_file(
+ ${CMAKE_CURRENT_SOURCE_DIR}/cmake/ggml-et-kernels-embed.hpp.in
+ ${EMBED_HPP}
+ @ONLY)
+configure_file(
+ ${CMAKE_CURRENT_SOURCE_DIR}/cmake/ggml-et-kernels-embed.cpp.in
+ ${EMBED_CPP}
+ @ONLY)
+
+add_custom_target(et-kernels-embed ALL
+ DEPENDS ${EMBED_KERNEL_SOURCES} ${EMBED_HPP} ${EMBED_CPP} et-uberkernel-map
+)
+
+ggml_add_backend_library(ggml-et
+ ggml-et.cpp
+ ggml-et-kernels.cpp
+ ggml-et-memops.cpp
+ ggml-et-ops.cpp
+ ggml-et-cpu-compare.cpp
+ )
+
+# Mark generated files as such
+set_source_files_properties(
+ ${EMBED_CPP}
+ ${EMBED_HPP}
+ ${EMBED_KERNEL_SOURCES}
+ ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-uberkernel-kernel-map.cpp
+ ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-uberkernel-kernel-map.h
+ PROPERTIES GENERATED TRUE
+)
+
+# Add embedded kernel sources
+target_sources(ggml-et PRIVATE
+ ${EMBED_CPP}
+ ${EMBED_HPP}
+ ${EMBED_KERNEL_SOURCES}
+ ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-uberkernel-kernel-map.cpp
+ ${CMAKE_CURRENT_BINARY_DIR}/et-kernels/ggml-et-uberkernel-kernel-map.h
+)
+
+# Include directory for embedded headers
+target_include_directories(ggml-et PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/et-kernels)
+
+target_link_libraries(ggml-et PRIVATE runtime::etrt_static deviceLayer::deviceLayer)
+target_compile_definitions(ggml-et PRIVATE GGML_ET_UBERKERNEL_HOST_LOOKUP)
+if (GGML_ET_SYSEMU)
+ target_compile_definitions(ggml-et PRIVATE GGML_ET_SYSEMU=1)
+endif()
+
+# Ensure kernels are built and embedded before the backend library
+add_dependencies(ggml-et et-kernels-embed et-uberkernel-map)
--- /dev/null
+# Inputs (via -D):
+# ELF_FILE - path to source .elf
+# OUT_FILE - path to output .cpp
+# VAR_NAME - C symbol base name (kernel name)
+
+file(READ "${ELF_FILE}" HEX HEX)
+string(LENGTH "${HEX}" HEX_LEN)
+math(EXPR SIZE "${HEX_LEN} / 2")
+string(REGEX REPLACE "(..)" "0x\\1," BYTES "${HEX}")
+
+file(WRITE "${OUT_FILE}"
+"// Auto-generated by embed_one_kernel.cmake. Do not edit.\n"
+"#include <cstdint>\n"
+"unsigned char ${VAR_NAME}_data[${SIZE}] = { ${BYTES} };\n"
+"extern const uint64_t ${VAR_NAME}_len = ${SIZE};\n")
--- /dev/null
+// Auto-generated kernel embeddings. Do not edit.
+#include "ggml-et-kernels-embed.hpp"
+
+const std::unordered_map<std::string, std::pair<const unsigned char*, uint64_t>> ggml_et_embedded_kernels = {
+@EMBED_MAP_ENTRIES@
+};
--- /dev/null
+// Auto-generated kernel embeddings. Do not edit.
+#pragma once
+
+#include <cstdint>
+#include <unordered_map>
+#include <string>
+#include <utility>
+
+@EMBED_EXTERNS@
+
+// Kernel name -> (data, length) lookup map
+extern const std::unordered_map<std::string, std::pair<const unsigned char*, uint64_t>> ggml_et_embedded_kernels;
--- /dev/null
+// Auto-generated uberkernel kernel-id mapping. Do not edit.
+#include "ggml-et-uberkernel-kernel-map.h"
+
+#ifdef GGML_ET_UBERKERNEL_HOST_LOOKUP
+#include <string>
+#include <unordered_map>
+
+uint16_t ggml_et_uberkernel_kernel_id_from_name(const char * kernel_name) {
+ if (kernel_name == nullptr) {
+ return GGML_ET_UBERKERNEL_KERNEL_INVALID;
+ }
+ static const std::unordered_map<std::string, uint16_t> kernel_id_map = {
+@UBERKERNEL_MAP_ENTRIES@
+ };
+ auto it = kernel_id_map.find(std::string(kernel_name));
+ return it == kernel_id_map.end() ? GGML_ET_UBERKERNEL_KERNEL_INVALID : it->second;
+}
+#endif
--- /dev/null
+// Auto-generated uberkernel kernel-id mapping. Do not edit.
+#pragma once
+
+#include <stdint.h>
+
+enum ggml_et_uberkernel_kernel_id {
+ GGML_ET_UBERKERNEL_KERNEL_INVALID = 0,
+@UBERKERNEL_ENUM_ENTRIES@
+};
+
+#ifdef GGML_ET_UBERKERNEL_HOST_LOOKUP
+uint16_t ggml_et_uberkernel_kernel_id_from_name(const char * kernel_name);
+#endif
--- /dev/null
+# ggml-et: Device kernels (cross-compiled within the main build)
+#
+# The RISC-V toolchain is set up in-scope so these targets use the
+# cross-compiler while the rest of the build uses the host compiler.
+# This keeps kernels in compile_commands.json for full IDE support.
+
+# --- RISC-V toolchain setup (scoped to this directory) ---
+set(TOOLCHAIN_DIR ${ET_PLATFORM_PATH})
+include(${ET_PLATFORM_PATH}/lib/cmake/riscv64-ec-toolchain.cmake)
+set(CMAKE_ADDR2LINE "${TOOLCHAIN_DIR}/bin/riscv64-unknown-elf-addr2line")
+set(CMAKE_LINKER_TYPE LLD)
+
+# Ensure kernels are built in this directory even if a global output directory is set
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+
+message(STATUS "ET kernels using RISC-V toolchain at: ${TOOLCHAIN_DIR}")
+
+# DeviceUtils provides the add_riscv_executable macro
+list(APPEND CMAKE_MODULE_PATH "${ET_PLATFORM_PATH}/lib/cmake/cmake-modules")
+list(APPEND CMAKE_PREFIX_PATH "${ET_PLATFORM_PATH}/lib/cmake")
+include(DeviceUtils)
+
+find_package(et-common-libs REQUIRED)
+find_package(esperantoTrace REQUIRED)
+
+# --- Kernel configuration ---
+if(NOT DEFINED ADDRESS)
+ set(ADDRESS "0x8005801000")
+ message(STATUS "ADDRESS not specified, using default: ${ADDRESS}")
+endif()
+
+set(LINKER_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/src/linker.ld)
+set(CHECK_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/scripts/check_unimplemented_instructions.sh)
+
+# Track address changes to trigger relinking
+set(ADDRESS_FILE ${CMAKE_CURRENT_BINARY_DIR}/et_address.txt)
+file(CONFIGURE OUTPUT ${ADDRESS_FILE} CONTENT "${ADDRESS}" @ONLY)
+
+# KERNELS defined in upper CMakeLists.txt
+foreach(KERNEL ${KERNELS})
+ add_riscv_executable(${KERNEL})
+ target_sources(${KERNEL}.elf PRIVATE
+ src/${KERNEL}.c
+ src/crt.S
+ )
+ target_include_directories(${KERNEL}.elf PRIVATE
+ ${CMAKE_CURRENT_SOURCE_DIR}/src
+ ${CMAKE_CURRENT_SOURCE_DIR}/..
+ ${CMAKE_CURRENT_BINARY_DIR}
+ ${CMAKE_SOURCE_DIR}/ggml/include
+ ${CMAKE_SOURCE_DIR}/ggml/src
+ )
+ target_link_libraries(${KERNEL}.elf PRIVATE et-common-libs::cm-umode)
+ # C-only flags — must not apply to .S files
+ target_compile_options(${KERNEL}.elf PRIVATE
+ $<$<COMPILE_LANGUAGE:C>:-fno-zero-initialized-in-bss>
+ $<$<COMPILE_LANGUAGE:C>:-ffreestanding>
+ $<$<COMPILE_LANGUAGE:C>:-std=gnu99>
+ $<$<COMPILE_LANGUAGE:C>:-ffat-lto-objects>
+ $<$<COMPILE_LANGUAGE:C>:-mcmodel=medany>
+ $<$<COMPILE_LANGUAGE:C>:-mabi=lp64f>
+ $<$<COMPILE_LANGUAGE:C>:-march=rv64imf>
+ $<$<COMPILE_LANGUAGE:C>:-ffunction-sections>
+ $<$<COMPILE_LANGUAGE:C>:-fdata-sections>
+ $<$<COMPILE_LANGUAGE:C>:-O3>
+ $<$<COMPILE_LANGUAGE:C>:-g0>
+ $<$<COMPILE_LANGUAGE:C>:-nostdlib>
+ $<$<COMPILE_LANGUAGE:C>:-ffreestanding>
+ )
+ target_link_options(${KERNEL}.elf PRIVATE
+ -Wl,--defsym=BASE_ADDRESS=${ADDRESS}
+ -Wl,--entry=_start
+ )
+ # Append to LINK_DEPENDS (macro already sets it for the linker script)
+ set_property(TARGET ${KERNEL}.elf APPEND PROPERTY
+ LINK_DEPENDS "${ADDRESS_FILE}"
+ )
+
+ # Post-build: strip and check (fails build if check script fails)
+ add_custom_command(TARGET ${KERNEL}.elf POST_BUILD
+ COMMAND ${CMAKE_STRIP} --strip-debug $<TARGET_FILE:${KERNEL}.elf>
+ COMMAND ${CHECK_SCRIPT}
+ ${CMAKE_OBJDUMP} ${CMAKE_ADDR2LINE} $<TARGET_FILE:${KERNEL}.elf>
+ DEPENDS ${CHECK_SCRIPT}
+ VERBATIM
+ )
+endforeach()
+
+add_dependencies(uberkernel.elf et-uberkernel-map)
+
+# Each supported kernel is compiled in its own translation unit with
+# -Dentry_point=<kernel>_entry
+# so symbols and macros don't leak between kernels. The dispatcher
+# (uberkernel.c) calls the renamed entries via extern declarations.
+#
+# HACK: we need to supresse _me kernels from setting up SCP themselves
+set(_UBER_ME_KERNELS mul_mat_f16_matrix_engine mul_mat_f32_matrix_engine flash_attn_ext_f16_me)
+
+foreach(UK_KERNEL ${UBERKERNEL_SUPPORTED_KERNELS})
+ set(_obj uber_${UK_KERNEL})
+ add_library(${_obj} OBJECT src/${UK_KERNEL}.c)
+ target_compile_definitions(${_obj} PRIVATE "entry_point=${UK_KERNEL}_entry" ET_UBERKERNEL)
+ target_include_directories(${_obj} PRIVATE
+ ${CMAKE_CURRENT_SOURCE_DIR}/src
+ ${CMAKE_CURRENT_SOURCE_DIR}/..
+ ${CMAKE_CURRENT_BINARY_DIR}
+ ${CMAKE_SOURCE_DIR}/ggml/include
+ ${CMAKE_SOURCE_DIR}/ggml/src
+ )
+ target_link_libraries(${_obj} PRIVATE et-common-libs::cm-umode)
+ target_compile_options(${_obj} PRIVATE
+ $<$<COMPILE_LANGUAGE:C>:-fno-zero-initialized-in-bss>
+ $<$<COMPILE_LANGUAGE:C>:-ffreestanding>
+ $<$<COMPILE_LANGUAGE:C>:-std=gnu99>
+ $<$<COMPILE_LANGUAGE:C>:-ffat-lto-objects>
+ $<$<COMPILE_LANGUAGE:C>:-mcmodel=medany>
+ $<$<COMPILE_LANGUAGE:C>:-mabi=lp64f>
+ $<$<COMPILE_LANGUAGE:C>:-march=rv64imf>
+ $<$<COMPILE_LANGUAGE:C>:-ffunction-sections>
+ $<$<COMPILE_LANGUAGE:C>:-fdata-sections>
+ $<$<COMPILE_LANGUAGE:C>:-O3>
+ $<$<COMPILE_LANGUAGE:C>:-g0>
+ $<$<COMPILE_LANGUAGE:C>:-nostdlib>
+ )
+ # ME kernels: suppress setup_cache_scp() (called once by the dispatcher)
+ if(UK_KERNEL IN_LIST _UBER_ME_KERNELS)
+ target_compile_definitions(${_obj} PRIVATE UBERKERNEL_SUPPRESS_SCP_SETUP)
+ endif()
+ target_sources(uberkernel.elf PRIVATE $<TARGET_OBJECTS:${_obj}>)
+endforeach()
+
+# Print summary
+message(STATUS "GGML ET Kernels configured:")
+foreach(KERNEL ${KERNELS})
+ message(STATUS " - ${KERNEL}")
+endforeach()
+message(STATUS "Base address: ${ADDRESS}")
--- /dev/null
+#!/bin/bash
+
+OBJDUMP=$1
+ADDR2LINE=$2
+TARGET_DEBUG=$3
+TARGET_ASM=${TARGET_DEBUG}.S
+BAD_INST_FILE=${TARGET_DEBUG}-BAD-INST.log
+
+# grep expression to find unimplemented instructions
+UNIMPLEMENTED_EXPR="fdiv.s\\|fsqrt.s\\|fcvt.l.s\\|fcvt.lu.s\\|fcvt.s.l\\|fcvt.s.lu\\|fdiv.pi\\|fdivu.pi\\|fremu.pi\\|frem.pi\\|fdiv.ps\\|fsqrt.ps\\|frsq.ps\\|fsin.ps"
+
+# dump assembly into .S file
+${OBJDUMP} -lwdSC ${TARGET_DEBUG} > ${TARGET_ASM}
+
+# check with grep for unimplemented instructions
+# Note: The exit status is 0 if selected lines are found, and 1 if not found.
+grep ${UNIMPLEMENTED_EXPR} ${TARGET_ASM} > /dev/null
+ret=$?
+
+if [ ${ret} -eq 0 ]
+then
+ # unimplemented instructions are found
+ echo -e "BUILD ERROR: Executable file ${TARGET_DEBUG} contains unimplemented instructions. Please review the lines of code listed in ${BAD_INST_FILE}"
+ echo -e "\t For further details, please read paragraph 3.4 of the ETSoC-1 Programmer's Reference Manual (PRM)"
+
+ # addr2line
+ grep ${UNIMPLEMENTED_EXPR} ${TARGET_ASM} | cut -d: -f 1 | ${ADDR2LINE} -i -e ${TARGET_DEBUG} > ${BAD_INST_FILE}
+ grep ${UNIMPLEMENTED_EXPR} ${TARGET_ASM} >> ${BAD_INST_FILE}
+ echo "------------------------------------------------------------"
+ cat ${BAD_INST_FILE}
+ echo "------------------------------------------------------------"
+ exit 1
+
+else
+ rm -f ${BAD_INST_FILE}
+fi
--- /dev/null
+#!/usr/bin/env bash
+set -euo pipefail
+
+LOG="llama_bench_$(date +%Y%m%d_%H%M%S).log"
+
+{
+ echo "===== START ====="
+ date
+ hostname
+ uname -a
+ echo "Command:"
+ echo "./build/bin/llama-bench -m ../../models/Llama-3.2-1B-Instruct-Q8_0.gguf -fa 0 -p 32,64,128,256,512 -n 32,64,128,256,512"
+ echo "================="
+
+ ./build/bin/llama-bench \
+ -m ../../models/Llama-3.2-1B-Instruct-Q8_0.gguf \
+ -fa 0 \
+ -p 32,64,128,256,512 \
+ -n 32,64,128,256,512
+
+ echo "===== END ====="
+ date
+} 2>&1 | tee "$LOG"
--- /dev/null
+//******************************************************************************
+// ET Vectorized Block Operations Library
+// Provides optimized block-level operations using ET hardware vector instructions
+//******************************************************************************
+
+#ifndef BLOCK_OPS_H
+# define BLOCK_OPS_H
+
+# include "math_fp.h"
+# include "quants.h"
+
+# include <stdint.h>
+
+//******************************************************************************
+// Block Dot Product Operations
+//******************************************************************************
+inline void __attribute__((always_inline)) excl_mode(uint64_t val) {
+ __asm__ __volatile__("csrw 0x7d3, %[csr_enc]\n" : : [csr_enc] "r"(val) : "x31");
+}
+
+static inline float compute_block_dot_product_q4_0(const block_q4_0 * a_block, const float * b_col_start) {
+ // Set mask register to enable all 8 vector elements
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+
+ // Use f10 as accumulator, init to 0
+ __asm__ volatile("fbci.ps f10, 0" ::: "f10");
+
+ static const int32_t gather_pattern[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+ __asm__ volatile("flw.ps f31, %[gather]\n" : : [gather] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+
+ // Process 32 elements in 2 chunks of 16 elements (8 bytes) each
+ for (int chunk = 0; chunk < 2; chunk++) {
+ int offset_a = chunk * 8;
+ int offset_b_low = chunk * 8; // Activations for lower nibbles
+ int offset_b_high = chunk * 8 + 16; // Activations for upper nibbles (16 elements later)
+
+ __asm__ volatile(
+ "fgb.ps f11, f31(%[a_ptr])\n" // Gather 8 bytes (16 packed q4_0 weights)
+
+ // 1. Extract & Multiply Lower Nibbles
+ "fandi.pi f12, f11, 15\n" // Mask lower 4 bits (x & 0xF)
+ "faddi.pi f12, f12, -8\n" // GGML offset to signed: (x & 0xF) - 8
+ "fcvt.ps.pw f12, f12, rne\n" // Convert INT32 to FP32
+ "flw.ps f13, 0(%[b_low])\n" // Load 8 B values (floats)
+ "fmadd.ps f10, f12, f13, f10, rne\n" // acc += A_low * B_low
+
+ // 2. Extract & Multiply Upper Nibbles
+ "fsrli.pi f14, f11, 4\n" // Shift upper 4 bits down
+ "fandi.pi f14, f14, 15\n" // Mask new lower 4 bits
+ "faddi.pi f14, f14, -8\n" // GGML offset to signed
+ "fcvt.ps.pw f14, f14, rne\n" // Convert INT32 to FP32
+ "flw.ps f15, 0(%[b_high])\n" // Load next 8 B values (floats)
+ "fmadd.ps f10, f14, f15, f10, rne\n" // acc += A_high * B_high
+ :
+ : [a_ptr] "r"(&a_block->qs[offset_a]), [b_low] "r"(&b_col_start[offset_b_low]),
+ [b_high] "r"(&b_col_start[offset_b_high])
+ // Note: f10 is explicitly NOT listed in the clobbers here to ensure the compiler
+ // preserves the running sum across C loop iterations safely.
+ : "f11", "f12", "f13", "f14", "f15");
+ }
+
+ // Horizontal sum: reduce f10 into a single scalar
+ float final_sum;
+ __asm__ __volatile__(
+ // Pairwise sum within each 128-bit half
+ "fswizz.ps f1, f10, 0xB1 \n\t" // Swaps: e0<->e1 and e2<->e3
+ "fadd.ps f2, f10, f1, rne \n\t"
+ // Complete the sum for each 128-bit half
+ "fswizz.ps f3, f2, 0x4E \n\t" // Swaps: e0,e1 <-> e2,e3
+ "fadd.ps f4, f2, f3, rne \n\t"
+ // Sum across the two 128b halfs
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(final_sum)::"t0", "f1", "f2", "f3", "f4", "f5", "f10");
+
+ // Restore original mask
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ const float scale = fp16_to_fp32(a_block->d);
+ return final_sum * scale;
+}
+
+// Compute dot product between dequantized q8_0 block and f32 column vector
+// Vectorized: processes 8 elements at a time using ET vector instructions
+// Block size: 32 int8 values (QK8_0)
+static inline float compute_block_dot_product_q8_0(const block_q8_0 * a_block, const float * b_col_start) {
+ // Set mask register to enable all 8 vector elements
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10"); // Use f10 as accumulator, init to 0
+
+ static const int32_t gather_pattern[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+
+ __asm__ volatile("flw.ps f31, %[gather]\n" : : [gather] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+
+ // Process 32 elements in 4 chunks of 8 elements each
+ for (int chunk = 0; chunk < 4; chunk++) {
+ int offset = chunk << 3; // chunk * 8
+
+ __asm__ volatile(
+ "flw.ps f12, %[b_vec]\n" // Load 8 B values (floats)
+ "fgb.ps f11, f31(%[a_ptr])\n" // Gather 8 int8 bytes from A using pattern
+ "fcvt.ps.pw f11, f11\n" // Convert int8 vector to float vector
+ "fmadd.ps f10, f11, f12, f10\n" // acc += a_vec * b_vec (8-wide)
+ :
+ : [a_ptr] "r"(&a_block->qs[offset]), [b_vec] "m"(*(const float (*)[8]) & b_col_start[offset]),
+ [scale] "m"(a_block->d)
+ : "f10", "f11", "f12");
+ }
+
+ // Horizontal sum: reduce f10 into a single scalar
+ float final_sum;
+ __asm__ __volatile__(
+ // Pairwise sum within each 128-bit half
+ "fswizz.ps f1, f10, 0xB1 \n\t" // Swaps: e0<->e1 and e2<->e3
+ "fadd.ps f2, f10, f1, rne \n\t"
+ // Complete the sum for each 128-bit half
+ "fswizz.ps f3, f2, 0x4E \n\t" // Swaps: e0,e1 <-> e2,e3
+ "fadd.ps f4, f2, f3, rne \n\t"
+ // Sum across the two 128b halfs
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(final_sum)::"t0", "f10", "f2", "f3", "f4", "f5");
+
+ // Restore original mask
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ const float scale = fp16_to_fp32(a_block->d);
+ return final_sum * scale;
+}
+
+//******************************************************************************
+// Split-phase Q8_0 dot product API
+//
+// q8_dot_begin(st) — save mask, set mask 0xFF
+// q8_dot_reset() — zero vector accumulator f20
+// q8_dot_tile(q, b, n) — accumulate n Q8_0 blocks into f20
+// q8_dot_reduce() — horizontal sum of f20, return scalar float
+// q8_dot_teardown(st) — restore original mask
+//
+// Register contract:
+// f20 — row accumulator (persistent across tiles, reset per row)
+// f31 — gather pattern (reloaded per q8_dot_tile call)
+// f10-f12 — scratch within tile
+// f15 — scale broadcast within tile
+// f1-f5, t0 — scratch within reduce
+//******************************************************************************
+
+static inline void __attribute__((always_inline)) q8_dot_reset(void) {
+ __asm__ volatile("fbci.pi f20, 0" ::: "f20");
+}
+
+// Accumulate n_blocks Q8_0 blocks into f20.
+// Uses fg32b.ps (fast gather with scalar pattern) for aligned chunks,
+// falls back to fgb.ps for chunks crossing a 32-byte boundary.
+static inline void __attribute__((always_inline)) q8_dot_tile(const block_q8_0 * q_row,
+ const float * b_col,
+ int64_t n_blocks) {
+ const int32_t gather_pattern[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+ const uint64_t gather_0_to_7 = 0x398a418820ULL;
+
+ __asm__ volatile("flw.ps f31, %[g]\n" : : [g] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+
+ for (int64_t kb = 0; kb < n_blocks; kb++) {
+ const block_q8_0 * blk = q_row + kb;
+ const float * b_ptr = b_col + (kb << 5);
+ const uintptr_t qs_addr = (uintptr_t) blk->qs;
+ const uintptr_t qs_aligned = qs_addr & ~(uintptr_t) 31;
+ const uintptr_t qs_low = qs_addr & 31;
+ const int fast_chunks = (int) ((32 - qs_low) >> 3);
+
+ if (fast_chunks >= 3) {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "flw.ps f12, %[bv0]\n"
+ "fg32b.ps f11, %[gi](%[ap0])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv1]\n"
+ "fg32b.ps f11, %[gi](%[ap1])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv2]\n"
+ "fg32b.ps f11, %[gi](%[ap2])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv3]\n"
+ "fgb.ps f11, f31(%[ap3])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ :
+ : [gi] "r"(gather_0_to_7), [ap0] "r"(qs_addr), [ap1] "r"(qs_aligned | ((qs_addr + 8) & 31)),
+ [ap2] "r"(qs_aligned | ((qs_addr + 16) & 31)), [ap3] "r"(&blk->qs[24]),
+ [bv0] "m"(*(const float (*)[8]) & b_ptr[0]), [bv1] "m"(*(const float (*)[8]) & b_ptr[8]),
+ [bv2] "m"(*(const float (*)[8]) & b_ptr[16]), [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12");
+ } else if (fast_chunks == 2) {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "flw.ps f12, %[bv0]\n"
+ "fg32b.ps f11, %[gi](%[ap0])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv1]\n"
+ "fg32b.ps f11, %[gi](%[ap1])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv2]\n"
+ "fgb.ps f11, f31(%[ap2])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv3]\n"
+ "fgb.ps f11, f31(%[ap3])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ :
+ : [gi] "r"(gather_0_to_7), [ap0] "r"(qs_addr), [ap1] "r"(qs_aligned | ((qs_addr + 8) & 31)),
+ [ap2] "r"(&blk->qs[16]), [ap3] "r"(&blk->qs[24]), [bv0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [bv1] "m"(*(const float (*)[8]) & b_ptr[8]), [bv2] "m"(*(const float (*)[8]) & b_ptr[16]),
+ [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12");
+ } else if (fast_chunks == 1) {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "flw.ps f12, %[bv0]\n"
+ "fg32b.ps f11, %[gi](%[ap0])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv1]\n"
+ "fgb.ps f11, f31(%[ap1])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv2]\n"
+ "fgb.ps f11, f31(%[ap2])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv3]\n"
+ "fgb.ps f11, f31(%[ap3])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ :
+ : [gi] "r"(gather_0_to_7), [ap0] "r"(qs_addr), [ap1] "r"(&blk->qs[8]), [ap2] "r"(&blk->qs[16]),
+ [ap3] "r"(&blk->qs[24]), [bv0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [bv1] "m"(*(const float (*)[8]) & b_ptr[8]), [bv2] "m"(*(const float (*)[8]) & b_ptr[16]),
+ [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12");
+ } else {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "flw.ps f12, %[bv0]\n"
+ "fgb.ps f11, f31(%[ap0])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv1]\n"
+ "fgb.ps f11, f31(%[ap1])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv2]\n"
+ "fgb.ps f11, f31(%[ap2])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "flw.ps f12, %[bv3]\n"
+ "fgb.ps f11, f31(%[ap3])\n"
+ "fcvt.ps.pw f11, f11\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ :
+ : [ap0] "r"(&blk->qs[0]), [ap1] "r"(&blk->qs[8]), [ap2] "r"(&blk->qs[16]), [ap3] "r"(&blk->qs[24]),
+ [bv0] "m"(*(const float (*)[8]) & b_ptr[0]), [bv1] "m"(*(const float (*)[8]) & b_ptr[8]),
+ [bv2] "m"(*(const float (*)[8]) & b_ptr[16]), [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12");
+ }
+
+ // f20 += f10 * broadcast(scale) — hardware fp16→fp32 via FCVT.PS.F16
+ uint32_t scale_raw = (uint32_t) blk->d;
+ __asm__ volatile(
+ "fbcx.ps f15, %[sb]\n"
+ "fcvt.ps.f16 f15, f15\n"
+ "fmadd.ps f20, f10, f15, f20\n"
+ :
+ : [sb] "r"(scale_raw)
+ : "f15", "f20");
+ }
+}
+
+// Horizontal sum of 8-element vector accumulator f20.
+static inline float __attribute__((always_inline)) q8_dot_reduce(void) {
+ float result;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f20, 0xB1 \n\t"
+ "fadd.ps f2, f20, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result)::"t0", "f1", "f2", "f3", "f4", "f5");
+ return result;
+}
+
+// Full-row dot product (convenience wrapper)
+static inline float compute_row_dot_q8_0(const block_q8_0 * q_row, const float * b_col, int64_t K_blocks) {
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+ q8_dot_reset();
+ q8_dot_tile(q_row, b_col, K_blocks);
+ float result = q8_dot_reduce();
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return result;
+}
+
+//******************************************************************************
+// Hoisted Q8_0 dot API
+//
+// q8_dot_begin/end save/restore the vector mask once around a long sequence of
+// dot products, so the per-row mask shuffles are hoisted out of the inner
+// loops. q8_dot_compute does a full-row dot (no mask handling). The _x2
+// variant computes two rows together while reusing each loaded B chunk —
+// only safe when both row pointers share the same 32-byte alignment phase
+// (i.e. the Q8 row stride is a multiple of 32).
+//******************************************************************************
+
+typedef struct {
+ unsigned long saved_mask;
+} q8_dot_state;
+
+static inline void q8_dot_begin(q8_dot_state * state) {
+ __asm__ volatile("mova.x.m %0" : "=r"(state->saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+}
+
+static inline void q8_dot_end(const q8_dot_state * state) {
+ __asm__ volatile("mova.m.x %0" ::"r"(state->saved_mask));
+}
+
+// Equivalent to q8_dot_reset+tile+reduce, without touching the mask register.
+// Caller is responsible for q8_dot_begin/end around the surrounding loop.
+static inline float q8_dot_compute(const block_q8_0 * q_row, const float * b_col, int64_t K_blocks) {
+ q8_dot_reset();
+ q8_dot_tile(q_row, b_col, K_blocks);
+ return q8_dot_reduce();
+}
+
+// Compute two row dots together while reusing the same loaded B chunks.
+//
+// Safe when every row starts at the same 32-byte offset, i.e. the Q8 row stride
+// is a multiple of 32. In that case the gather/alignment pattern is the same
+// for both rows at a given `kb`, so one set of B vector loads feeds both row
+// accumulators.
+static inline void q8_dot_compute_x2_aligned(const block_q8_0 * q_row0,
+ const block_q8_0 * q_row1,
+ const float * b_col,
+ int64_t K_blocks,
+ float * out0,
+ float * out1) {
+ const int32_t gather_pattern[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+ const uint64_t gather_0_to_7 = 0x398a418820ULL;
+ __asm__ volatile("flw.ps f31, %[g]\n" : : [g] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+ __asm__ volatile(
+ "fbci.pi f20, 0\n"
+ "fbci.pi f21, 0\n" ::
+ : "f20", "f21");
+
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const block_q8_0 * blk0 = q_row0 + kb;
+ const block_q8_0 * blk1 = q_row1 + kb;
+ const float * b_ptr = b_col + (kb << 5);
+
+ const uintptr_t qs_addr0 = (uintptr_t) blk0->qs;
+ const uintptr_t qs_addr1 = (uintptr_t) blk1->qs;
+ const uintptr_t qs_aligned0 = qs_addr0 & ~(uintptr_t) 31;
+ const uintptr_t qs_aligned1 = qs_addr1 & ~(uintptr_t) 31;
+ const int fast_chunks = (int) ((32 - (qs_addr0 & 31)) >> 3);
+
+ if (fast_chunks >= 3) {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "fbci.pi f11, 0\n"
+
+ "flw.ps f12, %[bv0]\n"
+ "fg32b.ps f16, %[gi](%[r0ap0])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f12, f10\n"
+ "fg32b.ps f17, %[gi](%[r1ap0])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f12, f11\n"
+
+ "flw.ps f13, %[bv1]\n"
+ "fg32b.ps f16, %[gi](%[r0ap1])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f13, f10\n"
+ "fg32b.ps f17, %[gi](%[r1ap1])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f13, f11\n"
+
+ "flw.ps f14, %[bv2]\n"
+ "fg32b.ps f16, %[gi](%[r0ap2])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f14, f10\n"
+ "fg32b.ps f17, %[gi](%[r1ap2])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f14, f11\n"
+
+ "flw.ps f15, %[bv3]\n"
+ "fgb.ps f16, f31(%[r0ap3])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f15, f10\n"
+ "fgb.ps f17, f31(%[r1ap3])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f15, f11\n"
+ :
+ : [gi] "r"(gather_0_to_7), [r0ap0] "r"(qs_addr0), [r0ap1] "r"(qs_aligned0 | ((qs_addr0 + 8) & 31)),
+ [r0ap2] "r"(qs_aligned0 | ((qs_addr0 + 16) & 31)), [r0ap3] "r"(&blk0->qs[24]), [r1ap0] "r"(qs_addr1),
+ [r1ap1] "r"(qs_aligned1 | ((qs_addr1 + 8) & 31)), [r1ap2] "r"(qs_aligned1 | ((qs_addr1 + 16) & 31)),
+ [r1ap3] "r"(&blk1->qs[24]), [bv0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [bv1] "m"(*(const float (*)[8]) & b_ptr[8]), [bv2] "m"(*(const float (*)[8]) & b_ptr[16]),
+ [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17");
+ } else if (fast_chunks == 2) {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "fbci.pi f11, 0\n"
+
+ "flw.ps f12, %[bv0]\n"
+ "fg32b.ps f16, %[gi](%[r0ap0])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f12, f10\n"
+ "fg32b.ps f17, %[gi](%[r1ap0])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f12, f11\n"
+
+ "flw.ps f13, %[bv1]\n"
+ "fg32b.ps f16, %[gi](%[r0ap1])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f13, f10\n"
+ "fg32b.ps f17, %[gi](%[r1ap1])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f13, f11\n"
+
+ "flw.ps f14, %[bv2]\n"
+ "fgb.ps f16, f31(%[r0ap2])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f14, f10\n"
+ "fgb.ps f17, f31(%[r1ap2])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f14, f11\n"
+
+ "flw.ps f15, %[bv3]\n"
+ "fgb.ps f16, f31(%[r0ap3])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f15, f10\n"
+ "fgb.ps f17, f31(%[r1ap3])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f15, f11\n"
+ :
+ : [gi] "r"(gather_0_to_7), [r0ap0] "r"(qs_addr0), [r0ap1] "r"(qs_aligned0 | ((qs_addr0 + 8) & 31)),
+ [r0ap2] "r"(&blk0->qs[16]), [r0ap3] "r"(&blk0->qs[24]), [r1ap0] "r"(qs_addr1),
+ [r1ap1] "r"(qs_aligned1 | ((qs_addr1 + 8) & 31)), [r1ap2] "r"(&blk1->qs[16]),
+ [r1ap3] "r"(&blk1->qs[24]), [bv0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [bv1] "m"(*(const float (*)[8]) & b_ptr[8]), [bv2] "m"(*(const float (*)[8]) & b_ptr[16]),
+ [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17");
+ } else if (fast_chunks == 1) {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "fbci.pi f11, 0\n"
+
+ "flw.ps f12, %[bv0]\n"
+ "fg32b.ps f16, %[gi](%[r0ap0])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f12, f10\n"
+ "fg32b.ps f17, %[gi](%[r1ap0])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f12, f11\n"
+
+ "flw.ps f13, %[bv1]\n"
+ "fgb.ps f16, f31(%[r0ap1])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f13, f10\n"
+ "fgb.ps f17, f31(%[r1ap1])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f13, f11\n"
+
+ "flw.ps f14, %[bv2]\n"
+ "fgb.ps f16, f31(%[r0ap2])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f14, f10\n"
+ "fgb.ps f17, f31(%[r1ap2])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f14, f11\n"
+
+ "flw.ps f15, %[bv3]\n"
+ "fgb.ps f16, f31(%[r0ap3])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f15, f10\n"
+ "fgb.ps f17, f31(%[r1ap3])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f15, f11\n"
+ :
+ : [gi] "r"(gather_0_to_7), [r0ap0] "r"(qs_addr0), [r0ap1] "r"(&blk0->qs[8]), [r0ap2] "r"(&blk0->qs[16]),
+ [r0ap3] "r"(&blk0->qs[24]), [r1ap0] "r"(qs_addr1), [r1ap1] "r"(&blk1->qs[8]),
+ [r1ap2] "r"(&blk1->qs[16]), [r1ap3] "r"(&blk1->qs[24]), [bv0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [bv1] "m"(*(const float (*)[8]) & b_ptr[8]), [bv2] "m"(*(const float (*)[8]) & b_ptr[16]),
+ [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17");
+ } else {
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "fbci.pi f11, 0\n"
+
+ "flw.ps f12, %[bv0]\n"
+ "fgb.ps f16, f31(%[r0ap0])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f12, f10\n"
+ "fgb.ps f17, f31(%[r1ap0])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f12, f11\n"
+
+ "flw.ps f13, %[bv1]\n"
+ "fgb.ps f16, f31(%[r0ap1])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f13, f10\n"
+ "fgb.ps f17, f31(%[r1ap1])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f13, f11\n"
+
+ "flw.ps f14, %[bv2]\n"
+ "fgb.ps f16, f31(%[r0ap2])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f14, f10\n"
+ "fgb.ps f17, f31(%[r1ap2])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f14, f11\n"
+
+ "flw.ps f15, %[bv3]\n"
+ "fgb.ps f16, f31(%[r0ap3])\n"
+ "fcvt.ps.pw f16, f16\n"
+ "fmadd.ps f10, f16, f15, f10\n"
+ "fgb.ps f17, f31(%[r1ap3])\n"
+ "fcvt.ps.pw f17, f17\n"
+ "fmadd.ps f11, f17, f15, f11\n"
+ :
+ : [r0ap0] "r"(&blk0->qs[0]), [r0ap1] "r"(&blk0->qs[8]), [r0ap2] "r"(&blk0->qs[16]),
+ [r0ap3] "r"(&blk0->qs[24]), [r1ap0] "r"(&blk1->qs[0]), [r1ap1] "r"(&blk1->qs[8]),
+ [r1ap2] "r"(&blk1->qs[16]), [r1ap3] "r"(&blk1->qs[24]), [bv0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [bv1] "m"(*(const float (*)[8]) & b_ptr[8]), [bv2] "m"(*(const float (*)[8]) & b_ptr[16]),
+ [bv3] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17");
+ }
+
+ const uint32_t scale_raw0 = (uint32_t) blk0->d;
+ const uint32_t scale_raw1 = (uint32_t) blk1->d;
+ __asm__ volatile(
+ "fbcx.ps f24, %[s0]\n"
+ "fcvt.ps.f16 f24, f24\n"
+ "fmadd.ps f20, f10, f24, f20\n"
+ "fbcx.ps f25, %[s1]\n"
+ "fcvt.ps.f16 f25, f25\n"
+ "fmadd.ps f21, f11, f25, f21\n"
+ :
+ : [s0] "r"(scale_raw0), [s1] "r"(scale_raw1)
+ : "f20", "f21", "f24", "f25");
+ }
+
+ float result0;
+ float result1;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f20, 0xB1 \n\t"
+ "fadd.ps f2, f20, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result0)::"t0", "f1", "f2", "f3", "f4", "f5");
+ __asm__ __volatile__(
+ "fswizz.ps f1, f21, 0xB1 \n\t"
+ "fadd.ps f2, f21, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result1)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ *out0 = result0;
+ *out1 = result1;
+}
+
+// Compute dot product between f16 block and f32 column vector (NAIVE VERSION)
+// Scalar implementation for debugging - no vectorization
+// Block size: 32 f16 values (64 bytes = 1 cache line)
+static inline float compute_block_dot_product_f16_naive(const uint16_t * a_block, const float * b_col_start) {
+ float acc_vec[8] __attribute__((aligned(32))) = { 0.0f };
+ // Byte offsets for 16-bit (half-word) elements
+ static const int32_t gather_pattern[8] = { 0, 2, 4, 6, 8, 10, 12, 14 };
+ unsigned long temp_mask;
+
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Load the pattern once into f31 for the duration of all 4 chunks
+ __asm__ volatile("flw.ps f31, %[gather]\n" : : [gather] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+
+ for (int chunk = 0; chunk < 4; chunk++) {
+ // Correct pointers:
+ // a_block elements are 2 bytes, b_col elements are 4 bytes
+ const uint16_t * a_ptr = &a_block[chunk << 3]; // chunk * 8
+ const float * b_ptr = &b_col_start[chunk << 3]; // chunk * 8
+
+ __asm__ volatile(
+ "flw.ps f10, %[acc]\n"
+ "fgh.ps f11, f31(%[a_p])\n" // Uses {0,2,4,6,8,10,12,14} byte offsets
+ "fcvt.ps.f16 f11, f11\n"
+ "flw.ps f12, (%[b_p])\n" // Standard vector load (32-bit floats)
+ "fmadd.ps f10, f11, f12, f10\n"
+ "fsw.ps f10, %[result]\n"
+
+ : [result] "=m"(*(float (*)[8]) acc_vec)
+ : [acc] "m"(*(const float (*)[8]) acc_vec), [a_p] "r"(a_ptr), [b_p] "r"(b_ptr)
+ : "f10", "f11", "f12");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ return acc_vec[0] + acc_vec[1] + acc_vec[2] + acc_vec[3] + acc_vec[4] + acc_vec[5] + acc_vec[6] + acc_vec[7];
+}
+
+// Compute dot product between f16 block and f32 column vector
+// SCALAR implementation for partial blocks
+// Block size: up to 32 f16 values (can handle partial blocks for misaligned K)
+static inline float compute_block_dot_product_f16_partial(const uint16_t * a_block,
+ const float * b_col_start,
+ int elements) {
+ // This matches compute_block_dot_product_f16_naive behavior
+ float sum = 0.0f;
+
+ for (int i = 0; i < elements; i++) {
+ float a_val = fp16_to_fp32(a_block[i]);
+ float b_val = b_col_start[i];
+ sum += a_val * b_val;
+ }
+
+ return sum;
+}
+
+// Compute dot product between f16 block and f16 column vector
+// Scalar implementation for generic non-matrix-engine fallback paths.
+static inline float compute_block_dot_product_f16_f16_partial(const uint16_t * a_block,
+ const uint16_t * b_col_start,
+ int elements) {
+ float sum = 0.0f;
+
+ for (int i = 0; i < elements; i++) {
+ sum += fp16_to_fp32(a_block[i]) * fp16_to_fp32(b_col_start[i]);
+ }
+
+ return sum;
+}
+
+// Compute dot product between f16 block and f32 column vector
+// Vectorized: processes 8 elements at a time using ET vector instructions
+// Block size: 32 f16 values (64 bytes = 1 cache line)
+static inline float compute_block_dot_product_f16(const uint16_t * a_block, const float * b_col_start) {
+ return compute_block_dot_product_f16_partial(a_block, b_col_start, QK_F16);
+}
+
+// Compute dot product between f32 block and f32 column vector
+// Vectorized: processes 8 elements at a time using ET vector instructions
+// Block size: up to 16 f32 values (can handle partial blocks for misaligned K)
+static inline float compute_block_dot_product_f32_partial(const float * a_block,
+ const float * b_col_start,
+ int elements) {
+ float acc_vec[8] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; // Accumulator vector
+
+ // Calculate how many full 8-element chunks we can process
+ int vec_end = (elements / 8) * 8;
+
+ if (vec_end > 0) {
+ // Set mask register to enable all 8 vector elements
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+
+ // Process full 8-element chunks
+ for (int i = 0; i < vec_end; i += 8) {
+ // Vectorized f32 multiply-accumulate
+ __asm__ volatile(
+ "flw.ps f10, %[acc]\n" // Load current accumulator (8 floats)
+ "flw.ps f11, %[a_vec]\n" // Load 8 A values (f32)
+ "flw.ps f12, %[b_vec]\n" // Load 8 B values (f32)
+ "fmadd.ps f10, f11, f12, f10\n" // acc += a_vec * b_vec (8-wide)
+ "fsw.ps f10, %[result]\n" // Store back to accumulator
+
+ : [result] "=m"(*(float (*)[8]) acc_vec)
+ : [acc] "m"(*(const float (*)[8]) acc_vec), [a_vec] "m"(*(const float (*)[8])(a_block + i)),
+ [b_vec] "m"(*(const float (*)[8])(b_col_start + i))
+ : "f10", "f11", "f12");
+ }
+
+ // Restore original mask
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+ }
+
+ // Horizontal sum: reduce 8 accumulator elements to single scalar
+ float final_sum = 0.0f;
+ for (int i = 0; i < 8; i++) {
+ final_sum += acc_vec[i];
+ }
+
+ // Handle remaining elements (< 8) with scalar operations
+ for (int i = vec_end; i < elements; i++) {
+ final_sum += a_block[i] * b_col_start[i];
+ }
+
+ return final_sum;
+}
+
+// Compute dot product between f32 block and f16 column vector
+// Scalar implementation for generic non-matrix-engine fallback paths.
+static inline float compute_block_dot_product_f32_f16_partial(const float * a_block,
+ const uint16_t * b_col_start,
+ int elements) {
+ float sum = 0.0f;
+
+ for (int i = 0; i < elements; i++) {
+ sum += a_block[i] * fp16_to_fp32(b_col_start[i]);
+ }
+
+ return sum;
+}
+
+// Compute dot product between f32 block and f32 column vector
+// Vectorized: processes 8 elements at a time using ET vector instructions
+// Block size: 16 f32 values (64 bytes = 1 cache line)
+static inline float compute_block_dot_product_f32(const float * a_block, const float * b_col_start) {
+ return compute_block_dot_product_f32_partial(a_block, b_col_start, QK_F32);
+
+ // float acc_vec[8];
+ // unsigned long old_mask;
+ // __asm__ volatile(
+ // // Save current mask
+ // "mova.x.m %[old_mask]\n"
+ // // Enable all 8 lanes
+ // "mov.m.x m0, x0, 0xFF\n"
+
+ // "flw.ps f11, %[a]\n"
+ // "flw.ps f12, %[b]\n"
+ // "fmadd.ps f10, f11, f12, f10\n"
+ // "fsw.ps f10, %[out]\n"
+ // "mova.m.x %[old_mask]\n"
+
+ // : [out] "=m" (*(float(*)[8])acc_vec),
+ // [old_mask] "=r"(old_mask)
+ // : [a] "m" (*(const float(*)[8])a_block),
+ // [b] "m" (*(const float(*)[8])b_col_start)
+ // : "f10", "f11", "f12"
+ // );
+
+ // // Horizontal reduction
+ // return acc_vec[0] + acc_vec[1] + acc_vec[2] + acc_vec[3] +
+ // acc_vec[4] + acc_vec[5] + acc_vec[6] + acc_vec[7];
+}
+
+#endif // BLOCK_OPS_H
+
+static inline void __attribute__((always_inline)) q4_dot_reset(void) {
+ __asm__ volatile("fbci.pi f20, 0" ::: "f20");
+}
+
+static inline void __attribute__((always_inline)) q4_dot_tile(const block_q4_0 * q_row,
+ const float * b_col,
+ int64_t n_blocks) {
+ const int32_t gather_pattern[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+ __asm__ volatile("flw.ps f31, %[g]\n" : : [g] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+
+ for (int64_t kb = 0; kb < n_blocks; kb++) {
+ const block_q4_0 * blk = q_row + kb;
+ const float * b_ptr = b_col + (kb << 5);
+
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+
+ "fgb.ps f11, f31(%[a_ptr0])\n"
+ "fandi.pi f12, f11, 15\n"
+ "faddi.pi f12, f12, -8\n"
+ "fcvt.ps.pw f12, f12, rne\n"
+ "flw.ps f13, %[b_low0]\n"
+ "fmadd.ps f10, f12, f13, f10, rne\n"
+
+ "fsrli.pi f14, f11, 4\n"
+ "fandi.pi f14, f14, 15\n"
+ "faddi.pi f14, f14, -8\n"
+ "fcvt.ps.pw f14, f14, rne\n"
+ "flw.ps f15, %[b_high0]\n"
+ "fmadd.ps f10, f14, f15, f10, rne\n"
+
+ "fgb.ps f11, f31(%[a_ptr1])\n"
+ "fandi.pi f12, f11, 15\n"
+ "faddi.pi f12, f12, -8\n"
+ "fcvt.ps.pw f12, f12, rne\n"
+ "flw.ps f13, %[b_low1]\n"
+ "fmadd.ps f10, f12, f13, f10, rne\n"
+
+ "fsrli.pi f14, f11, 4\n"
+ "fandi.pi f14, f14, 15\n"
+ "faddi.pi f14, f14, -8\n"
+ "fcvt.ps.pw f14, f14, rne\n"
+ "flw.ps f15, %[b_high1]\n"
+ "fmadd.ps f10, f14, f15, f10, rne\n"
+ :
+ : [a_ptr0] "r"(&blk->qs[0]), [b_low0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [b_high0] "m"(*(const float (*)[8]) & b_ptr[16]), [a_ptr1] "r"(&blk->qs[8]),
+ [b_low1] "m"(*(const float (*)[8]) & b_ptr[8]), [b_high1] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12", "f13", "f14", "f15");
+
+ uint32_t scale_raw = (uint32_t) blk->d;
+ __asm__ volatile(
+ "fbcx.ps f15, %[sb]\n"
+ "fcvt.ps.f16 f15, f15\n"
+ "fmadd.ps f20, f10, f15, f20\n"
+ :
+ : [sb] "r"(scale_raw)
+ : "f15", "f20");
+ }
+}
+
+static inline float __attribute__((always_inline)) q4_dot_reduce(void) {
+ float result;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f20, 0xB1 \n\t"
+ "fadd.ps f2, f20, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result)::"t0", "f1", "f2", "f3", "f4", "f5");
+ return result;
+}
+
+static inline float compute_row_dot_q4_0(const block_q4_0 * q_row, const float * b_col, int64_t K_blocks) {
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+ q4_dot_reset();
+ q4_dot_tile(q_row, b_col, K_blocks);
+ float result = q4_dot_reduce();
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return result;
+}
+
+typedef struct {
+ unsigned long saved_mask;
+} q4_dot_state;
+
+static inline void q4_dot_begin(q4_dot_state * state) {
+ __asm__ volatile("mova.x.m %0" : "=r"(state->saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+}
+
+static inline void q4_dot_end(const q4_dot_state * state) {
+ __asm__ volatile("mova.m.x %0" ::"r"(state->saved_mask));
+}
+
+static inline float q4_dot_compute(const block_q4_0 * q_row, const float * b_col, int64_t K_blocks) {
+ q4_dot_reset();
+ q4_dot_tile(q_row, b_col, K_blocks);
+ return q4_dot_reduce();
+}
+
+static inline void q4_dot_compute_x2_aligned(const block_q4_0 * q_row0,
+ const block_q4_0 * q_row1,
+ const float * b_col,
+ int64_t K_blocks,
+ float * out0,
+ float * out1) {
+ const int32_t gather_pattern[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+ __asm__ volatile("flw.ps f31, %[g]\n" : : [g] "m"(*(const int32_t (*)[8]) gather_pattern) : "f31");
+ __asm__ volatile(
+ "fbci.pi f20, 0\n"
+ "fbci.pi f21, 0\n" ::
+ : "f20", "f21");
+
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const block_q4_0 * blk0 = q_row0 + kb;
+ const block_q4_0 * blk1 = q_row1 + kb;
+ const float * b_ptr = b_col + (kb << 5);
+
+ __asm__ volatile(
+ "fbci.pi f10, 0\n"
+ "fbci.pi f16, 0\n"
+
+ "flw.ps f13, %[b_low0]\n"
+ "flw.ps f15, %[b_high0]\n"
+
+ "fgb.ps f11, f31(%[a_ptr0_0])\n"
+ "fgb.ps f17, f31(%[a_ptr1_0])\n"
+
+ "fandi.pi f12, f11, 15\n"
+ "faddi.pi f12, f12, -8\n"
+ "fcvt.ps.pw f12, f12, rne\n"
+ "fmadd.ps f10, f12, f13, f10, rne\n"
+
+ "fandi.pi f18, f17, 15\n"
+ "faddi.pi f18, f18, -8\n"
+ "fcvt.ps.pw f18, f18, rne\n"
+ "fmadd.ps f16, f18, f13, f16, rne\n"
+
+ "fsrli.pi f14, f11, 4\n"
+ "fandi.pi f14, f14, 15\n"
+ "faddi.pi f14, f14, -8\n"
+ "fcvt.ps.pw f14, f14, rne\n"
+ "fmadd.ps f10, f14, f15, f10, rne\n"
+
+ "fsrli.pi f19, f17, 4\n"
+ "fandi.pi f19, f19, 15\n"
+ "faddi.pi f19, f19, -8\n"
+ "fcvt.ps.pw f19, f19, rne\n"
+ "fmadd.ps f16, f19, f15, f16, rne\n"
+
+ "flw.ps f13, %[b_low1]\n"
+ "flw.ps f15, %[b_high1]\n"
+
+ "fgb.ps f11, f31(%[a_ptr0_1])\n"
+ "fgb.ps f17, f31(%[a_ptr1_1])\n"
+
+ "fandi.pi f12, f11, 15\n"
+ "faddi.pi f12, f12, -8\n"
+ "fcvt.ps.pw f12, f12, rne\n"
+ "fmadd.ps f10, f12, f13, f10, rne\n"
+
+ "fandi.pi f18, f17, 15\n"
+ "faddi.pi f18, f18, -8\n"
+ "fcvt.ps.pw f18, f18, rne\n"
+ "fmadd.ps f16, f18, f13, f16, rne\n"
+
+ "fsrli.pi f14, f11, 4\n"
+ "fandi.pi f14, f14, 15\n"
+ "faddi.pi f14, f14, -8\n"
+ "fcvt.ps.pw f14, f14, rne\n"
+ "fmadd.ps f10, f14, f15, f10, rne\n"
+
+ "fsrli.pi f19, f17, 4\n"
+ "fandi.pi f19, f19, 15\n"
+ "faddi.pi f19, f19, -8\n"
+ "fcvt.ps.pw f19, f19, rne\n"
+ "fmadd.ps f16, f19, f15, f16, rne\n"
+ :
+ : [a_ptr0_0] "r"(&blk0->qs[0]), [a_ptr0_1] "r"(&blk0->qs[8]), [a_ptr1_0] "r"(&blk1->qs[0]),
+ [a_ptr1_1] "r"(&blk1->qs[8]), [b_low0] "m"(*(const float (*)[8]) & b_ptr[0]),
+ [b_high0] "m"(*(const float (*)[8]) & b_ptr[16]), [b_low1] "m"(*(const float (*)[8]) & b_ptr[8]),
+ [b_high1] "m"(*(const float (*)[8]) & b_ptr[24])
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17", "f18", "f19");
+
+ const uint32_t scale_raw0 = (uint32_t) blk0->d;
+ const uint32_t scale_raw1 = (uint32_t) blk1->d;
+ __asm__ volatile(
+ "fbcx.ps f24, %[s0]\n"
+ "fcvt.ps.f16 f24, f24\n"
+ "fmadd.ps f20, f10, f24, f20\n"
+ "fbcx.ps f25, %[s1]\n"
+ "fcvt.ps.f16 f25, f25\n"
+ "fmadd.ps f21, f16, f25, f21\n"
+ :
+ : [s0] "r"(scale_raw0), [s1] "r"(scale_raw1)
+ : "f20", "f21", "f24", "f25");
+ }
+
+ float result0, result1;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f20, 0xB1 \n\t"
+ "fadd.ps f2, f20, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result0)::"t0", "f1", "f2", "f3", "f4", "f5");
+ __asm__ __volatile__(
+ "fswizz.ps f1, f21, 0xB1 \n\t"
+ "fadd.ps f2, f21, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result1)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ *out0 = result0;
+ *out1 = result1;
+}
--- /dev/null
+//******************************************************************************
+// CLAMP F32 Kernel
+// Element-wise: dst[i] = min(max(src0[i], min_val), max_val)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_clamp_params {
+ struct ggml_tensor src0; // F32 input (contiguous)
+ struct ggml_tensor dst; // F32 output (contiguous; may alias src0.data)
+ float min_val;
+ float max_val;
+};
+
+// Vectorized fmax/fmin clamp with scalar tail. n may be any non-negative int.
+static inline void clamp_block_f32(float * dst, const float * src, float min_val, float max_val, int32_t n) {
+ int32_t i = 0;
+ const int32_t vec_end = (n / 8) * 8;
+
+ if (vec_end > 0) {
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ for (; i < vec_end; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[s]\n"
+ "fbc.ps f11, %[mn]\n"
+ "fbc.ps f12, %[mx]\n"
+ "fmax.ps f13, f10, f11\n"
+ "fmin.ps f13, f13, f12\n"
+ "fsw.ps f13, %[d]\n"
+ : [d] "=m"(*(float (*)[8]) & dst[i])
+ : [s] "m"(*(const float (*)[8]) & src[i]), [mn] "m"(min_val), [mx] "m"(max_val)
+ : "f10", "f11", "f12", "f13");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+ }
+
+ for (; i < n; i++) {
+ float v = src[i];
+ if (v < min_val) {
+ v = min_val;
+ }
+ if (v > max_val) {
+ v = max_val;
+ }
+ dst[i] = v;
+ }
+}
+
+int entry_point(struct ggml_et_clamp_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t total_elements = src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3];
+ if (total_elements <= 0) {
+ return 0;
+ }
+
+ const float min_val = params->min_val;
+ const float max_val = params->max_val;
+
+ // Distribute by cache lines (16 F32 elements). Each thread owns disjoint
+ // cache lines, so a partial trailing line is written by exactly one
+ // thread — safe under non-coherent caches.
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = (int64_t) thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ clamp_block_f32(dst_data + es, src0_data + es, min_val, max_val, (int32_t) (ee - es));
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Concat F32 Kernel
+// Concatenates two F32 tensors along a specified dimension.
+// All copies are aligned to cacheline boundaries (64 bytes = 16 floats).
+//
+// For dim >= 1, entire rows are copied from src0 or src1 into dst.
+// For dim == 0, use:
+// - a fast vector path when both source row segments are cacheline-aligned
+// - a scalar stride-aware path otherwise
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+#include <string.h>
+
+struct ggml_et_concat_params {
+ struct ggml_tensor src0; // F32 input tensor 0
+ struct ggml_tensor src1; // F32 input tensor 1
+ struct ggml_tensor dst; // F32 output tensor
+ int32_t dim; // Concatenation dimension
+};
+
+// Copy n floats from src to dst using 8-wide vector loads/stores.
+// n must be a multiple of 16 (cacheline-aligned).
+static inline void copy_row_aligned(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[src_vec]\n"
+ "fsw.ps f11, %[dst_vec]\n"
+ : [dst_vec] "=m"(*(float (*)[8]) & dst[i])
+ : [src_vec] "m"(*(const float (*)[8]) & src[i])
+ : "f11");
+ }
+}
+
+int entry_point(struct ggml_et_concat_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * dst = ¶ms->dst;
+ int32_t dim = params->dim;
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * src1_data = (float *) src1->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne00 = src0->ne[0], ne01 = src0->ne[1], ne02 = src0->ne[2], ne03 = src0->ne[3];
+ const int64_t ne10 = src1->ne[0], ne11 = src1->ne[1], ne12 = src1->ne[2], ne13 = src1->ne[3];
+ const int64_t ne0 = dst->ne[0], ne1 = dst->ne[1], ne2 = dst->ne[2], ne3 = dst->ne[3];
+
+ // src strides in bytes
+ const size_t nb00 = src0->nb[0], nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ const size_t nb10 = src1->nb[0], nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
+ // dst strides in bytes
+ const size_t dnb1 = dst->nb[1], dnb2 = dst->nb[2], dnb3 = dst->nb[3];
+
+ // Total rows across all higher dimensions
+ const int64_t total_rows = ne1 * ne2 * ne3;
+
+ // Generic slow path for dim==0 when either source segment is not suitable for
+ // aligned vector copies. Threading is done by cacheline-aligned row groups,
+ // so writers do not share destination cache lines.
+ if (dim == 0 && (ne00 % 16 != 0 || ne10 % 16 != 0 || nb00 != sizeof(float) || nb10 != sizeof(float))) {
+ const int64_t rows_per_group = et_rows_per_cacheline_group(ne0, sizeof(float));
+ const int64_t total_groups = (total_rows + rows_per_group - 1) / rows_per_group;
+
+ for (int64_t grp = thread_id; grp < total_groups; grp += num_threads) {
+ const int64_t row_start = grp * rows_per_group;
+ int64_t row_end = row_start + rows_per_group;
+ if (row_end > total_rows) {
+ row_end = total_rows;
+ }
+
+ for (int64_t row = row_start; row < row_end; row++) {
+ int64_t i1 = row % ne1;
+ int64_t i2 = (row / ne1) % ne2;
+ int64_t i3 = row / (ne1 * ne2);
+
+ float * dst_row = (float *) ((char *) dst_data + i1 * dnb1 + i2 * dnb2 + i3 * dnb3);
+
+ const char * s0_base = (const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03;
+ for (int64_t i0 = 0; i0 < ne00; i0++) {
+ dst_row[i0] = *(const float *) (s0_base + i0 * nb00);
+ }
+
+ const char * s1_base = (const char *) src1_data + i1 * nb11 + i2 * nb12 + i3 * nb13;
+ for (int64_t i0 = 0; i0 < ne10; i0++) {
+ dst_row[ne00 + i0] = *(const float *) (s1_base + i0 * nb10);
+ }
+ }
+ }
+ return 0;
+ }
+
+ // Standard path: ne0 % 16 == 0, aligned rows
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ // Decompose linear row index into (i1, i2, i3)
+ int64_t i1 = row % ne1;
+ int64_t i2 = (row / ne1) % ne2;
+ int64_t i3 = row / (ne1 * ne2);
+
+ float * dst_row = (float *) ((char *) dst_data + i1 * dnb1 + i2 * dnb2 + i3 * dnb3);
+
+ if (dim == 0) {
+ // Concat along innermost dimension: [src0_row | src1_row]
+ // Both ne00 and ne10 are multiples of 16 (cacheline-aligned)
+ const float * s0_row = (const float *) ((const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+ const float * s1_row = (const float *) ((const char *) src1_data + i1 * nb11 + i2 * nb12 + i3 * nb13);
+
+ copy_row_aligned(dst_row, s0_row, (int32_t) ne00);
+ copy_row_aligned(dst_row + ne00, s1_row, (int32_t) ne10);
+
+ } else if (dim == 1) {
+ // Concat along dim 1: first ne01 rows from src0, rest from src1
+ if (i1 < ne01) {
+ const float * s0_row = (const float *) ((const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+ copy_row_aligned(dst_row, s0_row, (int32_t) ne0);
+ } else {
+ const float * s1_row =
+ (const float *) ((const char *) src1_data + (i1 - ne01) * nb11 + i2 * nb12 + i3 * nb13);
+ copy_row_aligned(dst_row, s1_row, (int32_t) ne0);
+ }
+
+ } else if (dim == 2) {
+ // Concat along dim 2: first ne02 slices from src0, rest from src1
+ if (i2 < ne02) {
+ const float * s0_row = (const float *) ((const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+ copy_row_aligned(dst_row, s0_row, (int32_t) ne0);
+ } else {
+ const float * s1_row =
+ (const float *) ((const char *) src1_data + i1 * nb11 + (i2 - ne02) * nb12 + i3 * nb13);
+ copy_row_aligned(dst_row, s1_row, (int32_t) ne0);
+ }
+
+ } else {
+ // dim == 3: first ne03 batches from src0, rest from src1
+ if (i3 < ne03) {
+ const float * s0_row = (const float *) ((const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+ copy_row_aligned(dst_row, s0_row, (int32_t) ne0);
+ } else {
+ const float * s1_row =
+ (const float *) ((const char *) src1_data + i1 * nb11 + i2 * nb12 + (i3 - ne03) * nb13);
+ copy_row_aligned(dst_row, s1_row, (int32_t) ne0);
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal CONT F16 Kernel
+// Converts non-contiguous F16 tensors to contiguous memory layout
+//
+// Note: F16 is represented as uint16_t (IEEE 754 binary16 format)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+struct ggml_et_cont_params {
+ struct ggml_tensor src0; // F16 input tensor (non-contiguous)
+ struct ggml_tensor dst; // F16 output tensor (contiguous)
+};
+
+int entry_point(struct ggml_et_cont_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = 2048; //get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0; // Non-contiguous input
+ struct ggml_tensor * dst = ¶ms->dst; // Contiguous output
+
+ if (src0->type != GGML_TYPE_F16 || dst->type != GGML_TYPE_F16) {
+ return -1; // Unsupported type combination
+ }
+
+ uint16_t * src0_data = (uint16_t *) src0->data;
+ uint16_t * dst_data = (uint16_t *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ const int64_t src_elements = src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3];
+ const int64_t dst_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+ if (src_elements != dst_elements) {
+ return -1; // Element count mismatch
+ }
+
+ // Source tensor dimensions and strides
+ const int64_t ne00 = src0->ne[0];
+ const int64_t ne01 = src0->ne[1];
+ const int64_t ne02 = src0->ne[2];
+ const int64_t ne03 = src0->ne[3];
+
+ const int64_t nb00 = src0->nb[0];
+ const int64_t nb01 = src0->nb[1];
+ const int64_t nb02 = src0->nb[2];
+ const int64_t nb03 = src0->nb[3];
+
+ // Parallelize by rows (dimension 1)
+ const int64_t total_rows = ne01;
+ const int64_t rows_per_thread = (total_rows + num_threads - 1) / num_threads;
+ const int64_t start_row = thread_id * rows_per_thread;
+ const int64_t end_row = (start_row + rows_per_thread < total_rows) ? (start_row + rows_per_thread) : total_rows;
+
+ if (start_row >= total_rows) {
+ return 0;
+ }
+
+ // Iterate over source tensor dimensions
+ for (int64_t i03 = 0; i03 < ne03; i03++) {
+ for (int64_t i02 = 0; i02 < ne02; i02++) {
+ // Calculate base linear index for this (i03, i02) slice in destination
+ const int64_t dst_linear_base = i03 * ne02 * ne01 * ne00 + i02 * ne01 * ne00;
+
+ // Process this thread's assigned rows
+ for (int64_t i01 = start_row; i01 < end_row; i01++) {
+ // Linear index for start of this row in destination
+ const int64_t dst_linear_row_base = dst_linear_base + i01 * ne00;
+
+ // Inner loop over dimension 0
+ for (int64_t i00 = 0; i00 < ne00; i00++) {
+ // Source offset using non-contiguous strides
+ const int64_t src_offset_bytes = i00 * nb00 + i01 * nb01 + i02 * nb02 + i03 * nb03;
+ const uint16_t * src_ptr = (const uint16_t *) ((const char *) src0_data + src_offset_bytes);
+
+ // Destination linear index (contiguous layout)
+ const int64_t dst_linear_idx = dst_linear_row_base + i00;
+
+ // Use atomic store for thread safety
+ atomic_store_f16((volatile uint16_t *) &dst_data[dst_linear_idx], *src_ptr);
+ }
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal CONT F32 Kernel
+// Converts non-contiguous tensors to contiguous memory layout
+//
+// Fast path: src contiguous: flat vectorized copy by cache lines
+// Aligned path: nb00==4 and ne00 % 16 == 0: distribute rows, no coherency issue
+// Unaligned: nb00==4 and ne00 not aligned: distribute by cache lines,
+// reverse-compute src coords, handle partial rows at boundaries
+// Fallback: nb00 != 4: scalar per-element
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdbool.h>
+#include <stdint.h>
+
+struct ggml_et_cont_params {
+ struct ggml_tensor src0; // F32 input tensor (non-contiguous)
+ struct ggml_tensor dst; // F32 output tensor (contiguous)
+};
+
+// Vectorized copy with scalar tail
+static inline void vec_copy_f32(float * dst, const float * src, int32_t n) {
+ int32_t i = 0;
+ const int32_t vec_end = (n / 8) * 8;
+ for (; i < vec_end; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[s]\n"
+ "fsw.ps f10, %[d]\n"
+ : [d] "=m"(*(float (*)[8]) & dst[i])
+ : [s] "m"(*(const float (*)[8]) & src[i])
+ : "f10");
+ }
+ for (; i < n; i++) {
+ dst[i] = src[i];
+ }
+}
+
+// Scalar copy
+static inline void scalar_copy_f32(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ dst[i] = src[i];
+ }
+}
+
+// static inline size_t tensor_bytes(const struct ggml_tensor *t) {
+// return (size_t)t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+// }
+
+int entry_point(struct ggml_et_cont_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne00 = src0->ne[0];
+ const int64_t ne01 = src0->ne[1];
+ const int64_t ne02 = src0->ne[2];
+ const int64_t ne03 = src0->ne[3];
+
+ const int64_t nb00 = src0->nb[0];
+ const int64_t nb01 = src0->nb[1];
+ const int64_t nb02 = src0->nb[2];
+ const int64_t nb03 = src0->nb[3];
+
+ const int64_t total_elements = ne00 * ne01 * ne02 * ne03;
+
+ if (total_elements == 0) {
+ return 0;
+ }
+
+ const bool src_contiguous = ggml_tensor_is_contiguous(src0, 4);
+
+ //==========================================================================
+ // Fast path: src is contiguous: flat vectorized copy by cache lines
+ //==========================================================================
+ if (src_contiguous) {
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ vec_copy_f32(dst_data + es, src0_data + es, (int32_t) (ee - es));
+ return 0;
+ }
+
+ //==========================================================================
+ // Non-contiguous paths: require nb00==4 (dim 0 contiguous in src)
+ //==========================================================================
+ if (nb00 != 4) {
+ // Fully non-contiguous scalar fallback — distribute by cache lines
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ for (int64_t idx = es; idx < ee; idx++) {
+ const int64_t i00 = idx % ne00;
+ const int64_t rem1 = idx / ne00;
+ const int64_t i01 = rem1 % ne01;
+ const int64_t rem2 = rem1 / ne01;
+ const int64_t i02 = rem2 % ne02;
+ const int64_t i03 = rem2 / ne02;
+
+ const float * sp =
+ (const float *) ((const char *) src0_data + i00 * nb00 + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ dst_data[idx] = *sp;
+ }
+ return 0;
+ }
+
+ // nb00 == 4 from here: dim 0 is contiguous in src
+
+ //==========================================================================
+ // Aligned path: ne00 % 16 == 0: rows are cache-line aligned, distribute rows
+ //==========================================================================
+ if (ne00 % 16 == 0) {
+ const int64_t total_rows = ne01 * ne02 * ne03;
+ const int64_t rows_per_thread = (total_rows + num_threads - 1) / num_threads;
+ const int64_t start_row = thread_id * rows_per_thread;
+ const int64_t end_row = (start_row + rows_per_thread < total_rows) ? (start_row + rows_per_thread) : total_rows;
+
+ if (start_row >= total_rows) {
+ return 0;
+ }
+
+ for (int64_t ir = start_row; ir < end_row; ir++) {
+ const int64_t i03 = ir / (ne02 * ne01);
+ const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01;
+ const int64_t i01 = ir - i03 * ne02 * ne01 - i02 * ne01;
+
+ const float * src_row = (const float *) ((const char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ float * dst_row = dst_data + ir * ne00;
+
+ vec_copy_f32(dst_row, src_row, (int32_t) ne00);
+ }
+ return 0;
+ }
+
+ //==========================================================================
+ // Unaligned path: ne00 % 16 != 0, nb00 == 4
+ // Distribute cache-line-aligned chunks of dst, handle partial rows at edges
+ //==========================================================================
+ {
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ int64_t pos = es;
+
+ // Compute starting row coordinates
+ int64_t row_idx = pos / ne00;
+ int64_t col = pos % ne00;
+
+ while (pos < ee) {
+ // Decompose row_idx -> (i01, i02, i03)
+ const int64_t i03 = row_idx / (ne02 * ne01);
+ const int64_t i02 = (row_idx - i03 * ne02 * ne01) / ne01;
+ const int64_t i01 = row_idx - i03 * ne02 * ne01 - i02 * ne01;
+
+ const float * src_row = (const float *) ((const char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03);
+
+ // How many elements left in this row and in our chunk
+ int64_t row_remaining = ne00 - col;
+ int64_t chunk_remaining = ee - pos;
+ int32_t n = (int32_t) (row_remaining < chunk_remaining ? row_remaining : chunk_remaining);
+
+ vec_copy_f32(dst_data + pos, src_row + col, n);
+
+ pos += n;
+ col = 0; // subsequent rows start at column 0
+ row_idx++;
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// 2D F32 convolution on the ET-SoC-1 matrix engine (GGML CONV_2D layout).
+//
+// LAYOUT (matches GGML's standard CONV_2D, cwhn=false; wireable directly):
+// src1 input : ne = [W, H, Cin, N=1] memory: input [n][cin][h][w]
+// src0 filter: ne = [Kw, Kh, Cin, Cout] memory: filter[oc][ic][kh][kw]
+// dst output: ne = [W, H, Cout, N=1] memory: output[n][oc][h][w]
+//
+// CONSTRAINTS (enforced at supports_op):
+// F32 throughout, N == 1, Cin % 16 == 0, Cout % 16 == 0, positive
+// stride/pad, dilation == 1. Tile/L2SCP limits are checked here.
+//
+// MEMORY MODEL:
+// Each active shire uses its own 2 MB local L2 SCP:
+// filter slice | pin buffer 0 | pin buffer 1? | output staging? | scratch
+//
+// The filter slice contains only the output-channel tiles (`mt`) consumed
+// by this shire's tile assignment. That keeps hart-0's inner-loop
+// tensor_loads local to the shire and avoids packing unused filter slabs.
+//
+// THREADING (multi-minion, multi-shire):
+// PHASE 1 (per-shire filter pack): hart-1's pack this shire's filter
+// slice into local L2 SCP. Work is slab-striped across the 32 minions.
+//
+// PHASE 2 (per-shire compute): hart-1's pack the input pin chunks while
+// hart-0's run the matrix engine. Pin double-buffering hides the next
+// chunk pack behind the current chunk's FMA pipeline when Cin does not
+// fit in one local buffer.
+//
+// PERFORMANCE STRATEGIES:
+// 1. Local filter slice: pack only the `mt` values this shire consumes;
+// inner-loop tensor_loads stay shire-local.
+// 2. Pin Cin streaming + chunk double-buffer: pack one
+// chunk while computing the prior one.
+// 3. TenC save/restore: f0..f31 IS the TenC accumulator;
+// spill/refill via L2 SCP scratch lets each hart hold multiple
+// partial accumulators across chunks.
+// 4. OW%16 staging: for partial-tile output, write to a
+// padded L2 SCP region then have one hart scalar-emit to DRAM.
+//
+// WHY THE FILTER PACK EXISTS:
+// GGML's OIHW filter has stride Kh*Kw*4 between consecutive Cin elements
+// (e.g. 36 bytes for 3x3) — usually NOT a multiple of 64, so plain
+// tensor_load cannot gather it directly. The per-slab pack into a
+// Cin-innermost form gives every per-tap slab a flat 64-byte row stride
+// and enables tensor_load.
+//
+// Picking M=Cout, N=W means TenC's natural row stride matches NCHW
+// output's per-channel stride (H*W*4) — the output store is a clean
+// tensor_store with no transpose. The price is that conv_size/conv_ctrl
+// no longer help with W boundaries (mask gates M, not N), so we handle
+// boundaries up-front by zero-padding the input in L2SCP.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+#include "tensor.h"
+
+#include <etsoc/common/utils.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+#define TILE 16 /* matrix engine native tile in M, K, N */
+/* L1 SCP layout: A double-buffered, B single-buffered. Per the SDK doc
+ `dst_start` is a 6-bit field (max 63) but empirical testing shows the
+ physical L1 SCP per minion is 48 lines — writes to lines >= 48 corrupt.
+ So we get 3 × 16-line buffers max: A_0, A_1, B. Pick A as the
+ double-buffered operand (filter-slab loads, the longer of the two). */
+#define LSCP_A_0 0 /* A buffer 0 at L1 SCP lines 0..15 */
+#define LSCP_A_1 16 /* A buffer 1 at L1 SCP lines 16..31 */
+#define LSCP_B 32 /* B (single buffer) at lines 32..47 */
+#define N_MIN_PER_SHIRE 32 /* ET-SoC-1 geometry: 32 minions/shire */
+#define N_SHIRES 32 /* default active shire count */
+#define MAX_TILES_PER_HART 2 /* per-hart TenC slots (save/restore) */
+#define MAX_DBL_BUFS 2 /* chunk pack buffers (double-buffered) */
+
+/* Per-shire L2 SCP local budget. Per-shire SCP is 2 MB; we cap at
+ 1984 KB to leave 64 KB headroom for per-hart TenC scratch (32 minions ×
+ 2 slots × 1 KB), which lives at the tail of the SCP outside the pin
+ sizing budget. Bigger budget here means bigger feasible chunk_KT,
+ which means fewer chunks (each chunk costs 2 SHIRE barriers + ~30
+ TenC save/restore events per hart). */
+#define LOCAL_BUDGET (1984 * 1024)
+
+/* Cap on the per-shire filter region in local L2 SCP. The shire packs the
+ mt values it can consume under the current tile assignment, rather than
+ the whole Cout dimension. Reads in the inner loop are then fully
+ shire-local — no NoC fanout. */
+#define LOCAL_FILTER_CAP (1024 * 1024) /* 1 MB / shire ceiling */
+
+#define SLAB_BYTES ((uint64_t) TILE * TILE * sizeof(float)) /* 1024 */
+#define SLAB_LINES ((SLAB_BYTES + 63) / 64) /* 16 */
+
+/* Upper bound on the number of distinct mt values a single shire may pack.
+ This keeps the mt list stack-resident. Shapes that need more should fall
+ back until the filter-slice bookkeeping is made dynamic. */
+#define MAX_MY_MT (N_MIN_PER_SHIRE * MAX_TILES_PER_HART)
+
+typedef struct {
+ int mt;
+ int mt_idx;
+ int oh;
+ int ow_base;
+} conv_tile_t;
+
+static inline int ceil_div_i32(int x, int y) {
+ return (x + y - 1) / y;
+}
+
+static inline int round_up_tile_i32(int x) {
+ return (x + TILE - 1) & ~(TILE - 1);
+}
+
+static inline int min_i32(int a, int b) {
+ return a < b ? a : b;
+}
+
+static inline uint64_t min_u64(uint64_t a, uint64_t b) {
+ return a < b ? a : b;
+}
+
+/* ===== Vector helpers for hart-1 pack ============================
+ Both assume dst (and src for copy) are 32-byte aligned; n is in floats.
+ The 8-element tail is handled scalar. f30/f31 are scratch — clobbered
+ per-call via the asm clobber list. */
+static inline void vec_zero_aligned(float * dst, int n) {
+ int i = 0;
+ const int n8 = n & ~7;
+ for (; i < n8; i += 8) {
+ __asm__ volatile(
+ "fsub.ps f31, f31, f31\n"
+ "fsw.ps f31, %[d]\n"
+ : [d] "=m"(*(float (*)[8]) & dst[i])
+ :
+ : "f31");
+ }
+ for (; i < n; ++i) {
+ dst[i] = 0.0f;
+ }
+}
+
+static inline void vec_copy_aligned(float * dst, const float * src, int n) {
+ int i = 0;
+ const int n8 = n & ~7;
+ for (; i < n8; i += 8) {
+ __asm__ volatile(
+ "flw.ps f30, %[s]\n"
+ "fsw.ps f30, %[d]\n"
+ : [d] "=m"(*(float (*)[8]) & dst[i])
+ : [s] "m"(*(const float (*)[8]) & src[i])
+ : "f30");
+ }
+ for (; i < n; ++i) {
+ dst[i] = src[i];
+ }
+}
+
+/* ===== TenC save/restore =========================================
+ The TenC accumulator IS the f0..f31 vector register file: row N occupies
+ f(2N) and f(2N+1) (two 8-fp32 vector regs per row). We save by
+ tensor_store-ing TILE rows × 64 bytes, and restore via 32 flw.ps after
+ forcing L1D to refetch from the L2SCP backing (tensor_store bypasses L1D
+ so the backing is always current). See feedback_tenc_save_restore.md. */
+static inline void tenc_restore_from_scratch(uint64_t scr) {
+ FENCE;
+ evict_to_l2((const void *) scr, TILE, 64);
+ WAIT_CACHEOPS;
+ __asm__ volatile(
+ "flw.ps f0, 0(%0)\n"
+ "flw.ps f1, 32(%0)\n"
+ "flw.ps f2, 64(%0)\n"
+ "flw.ps f3, 96(%0)\n"
+ "flw.ps f4, 128(%0)\n"
+ "flw.ps f5, 160(%0)\n"
+ "flw.ps f6, 192(%0)\n"
+ "flw.ps f7, 224(%0)\n"
+ "flw.ps f8, 256(%0)\n"
+ "flw.ps f9, 288(%0)\n"
+ "flw.ps f10, 320(%0)\n"
+ "flw.ps f11, 352(%0)\n"
+ "flw.ps f12, 384(%0)\n"
+ "flw.ps f13, 416(%0)\n"
+ "flw.ps f14, 448(%0)\n"
+ "flw.ps f15, 480(%0)\n"
+ "flw.ps f16, 512(%0)\n"
+ "flw.ps f17, 544(%0)\n"
+ "flw.ps f18, 576(%0)\n"
+ "flw.ps f19, 608(%0)\n"
+ "flw.ps f20, 640(%0)\n"
+ "flw.ps f21, 672(%0)\n"
+ "flw.ps f22, 704(%0)\n"
+ "flw.ps f23, 736(%0)\n"
+ "flw.ps f24, 768(%0)\n"
+ "flw.ps f25, 800(%0)\n"
+ "flw.ps f26, 832(%0)\n"
+ "flw.ps f27, 864(%0)\n"
+ "flw.ps f28, 896(%0)\n"
+ "flw.ps f29, 928(%0)\n"
+ "flw.ps f30, 960(%0)\n"
+ "flw.ps f31, 992(%0)\n"
+ :
+ : "r"(scr)
+ : "f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "f9", "f10", "f11", "f12", "f13", "f14", "f15", "f16",
+ "f17", "f18", "f19", "f20", "f21", "f22", "f23", "f24", "f25", "f26", "f27", "f28", "f29", "f30", "f31",
+ "memory");
+}
+
+/* ===== Pin pack context ==========================================
+ Loop-invariant state hart-1 needs to pack one Cin chunk's worth of
+ pin (Kw shifted, padded copies of input rows) into local L2 SCP. The
+ filter is not touched in this struct; it is packed into the per-shire
+ local slice before the per-chunk loop begins. */
+typedef struct {
+ const float * in_base; /* DRAM input base [Cin][H][W] */
+ int Kw;
+ int chunk_KT; /* number of K_TILES (=16-wide) per chunk */
+ int H, W, Hp, Wp_a;
+ int pad_h, pad_w, s0;
+ int minion; /* this hart's minion id (0..31) */
+ uint64_t pin_copy_floats; /* per-_s pin plane size in floats */
+ uint64_t l2_pad_in_buf[MAX_DBL_BUFS];
+ uint64_t pin_chunk_bytes; /* one chunk pin buffer's total size */
+} pin_ctx_t;
+
+static inline int find_mt_idx(const int * my_mt, int n_my_mt, int mt) {
+ for (int j = 0; j < n_my_mt; ++j) {
+ if (my_mt[j] == mt) {
+ return j;
+ }
+ }
+ return 0;
+}
+
+static inline conv_tile_t decode_tile(int t, int M_TILES, int w_tiles, const int * my_mt, int n_my_mt) {
+ conv_tile_t tile;
+ tile.mt = t % M_TILES;
+ t /= M_TILES;
+ const int wt = t % w_tiles;
+ t /= w_tiles;
+ tile.oh = t;
+ tile.ow_base = wt * TILE;
+ tile.mt_idx = find_mt_idx(my_mt, n_my_mt, tile.mt);
+ return tile;
+}
+
+static inline uint64_t
+filter_slab_addr(uint64_t l2_filter, int Kw, int K_TILES, int n_my_mt, int mt_idx, int kh, int kw, int kt_global) {
+ return l2_filter + (uint64_t) ((((kh * Kw + kw) * n_my_mt + mt_idx) * K_TILES + kt_global)) * SLAB_BYTES;
+}
+
+static inline uint64_t pin_tile_addr(uint64_t l2_pad_in,
+ uint64_t pin_copy_bytes,
+ int ktc,
+ int kw,
+ int Hp,
+ int Wp_a,
+ int oh,
+ int ow_base,
+ int s1,
+ int kh) {
+ const int ir_pad = oh * s1 + kh;
+ return l2_pad_in + (uint64_t) kw * pin_copy_bytes +
+ (((uint64_t) (ktc * TILE) * Hp + ir_pad) * Wp_a + ow_base) * sizeof(float);
+}
+
+static inline char * output_tile_addr(char * out_base,
+ const conv_tile_t * tile,
+ uint64_t out_chan_stride,
+ uint64_t out_row_stride) {
+ return out_base + (size_t) (tile->mt * TILE) * out_chan_stride + (size_t) tile->oh * out_row_stride +
+ (size_t) tile->ow_base * sizeof(float);
+}
+
+static inline void flush_range_to_l2(const void * addr, uint64_t n_bytes) {
+ const uint64_t total_lines = (n_bytes + 63) / 64;
+ const char * fl_addr = (const char *) addr;
+ for (uint64_t done = 0; done < total_lines;) {
+ const uint64_t batch = min_u64(total_lines - done, 16);
+ flush_to_l2((const void *) (fl_addr + done * 64), batch, 64);
+ done += batch;
+ }
+}
+
+static inline void evict_range_past_l2(const void * addr, uint64_t n_bytes) {
+ const uint64_t total_lines = (n_bytes + 63) / 64;
+ const char * fl_addr = (const char *) addr;
+ for (uint64_t done = 0; done < total_lines;) {
+ const uint64_t batch = min_u64(total_lines - done, 16);
+ evict_past_l2((const void *) (fl_addr + done * 64), batch, 64);
+ done += batch;
+ }
+}
+
+/* One matrix-engine tile for one Cin chunk. This is the main optimization
+ surface: A is double-buffered, B is single-buffered due to L1 SCP space. */
+static inline void compute_tile_chunk(uint64_t l2_filter,
+ uint64_t l2_pad_in,
+ uint64_t pin_copy_bytes,
+ int Kh,
+ int Kw,
+ int K_TILES,
+ int chunk_KT,
+ int kt_base,
+ int n_my_mt,
+ int Hp,
+ int Wp_a,
+ int s1,
+ uint64_t a_row_stride,
+ uint64_t b_row_stride,
+ const conv_tile_t * tile,
+ bool first_fma_clears_tenc) {
+ const int n_iters = Kh * Kw * chunk_KT;
+ const uint64_t A_BUFS[2] = { LSCP_A_0, LSCP_A_1 };
+
+ const uint64_t a_addr0 = filter_slab_addr(l2_filter, Kw, K_TILES, n_my_mt, tile->mt_idx, 0, 0, kt_base);
+ tensor_load(false, false, A_BUFS[0], 0, 0, a_addr0, 0, (uint64_t) (TILE - 1), a_row_stride, 0);
+
+ for (int iter = 0; iter < n_iters; ++iter) {
+ const int ktc = iter % chunk_KT;
+ const int rem = iter / chunk_KT;
+ const int kw = rem % Kw;
+ const int kh = rem / Kw;
+
+ const uint64_t b_addr =
+ pin_tile_addr(l2_pad_in, pin_copy_bytes, ktc, kw, Hp, Wp_a, tile->oh, tile->ow_base, s1, kh);
+ tensor_load(false, false, LSCP_B, 0, 0, b_addr, 0, (uint64_t) (TILE - 1), b_row_stride, 1);
+
+ tensor_wait(TENSOR_LOAD_WAIT_0);
+ tensor_wait(TENSOR_LOAD_WAIT_1);
+
+ if (iter + 1 < n_iters) {
+ const int ktc_n = (iter + 1) % chunk_KT;
+ const int rem_n = (iter + 1) / chunk_KT;
+ const int kw_n = rem_n % Kw;
+ const int kh_n = rem_n / Kw;
+ const uint64_t a_addr_n =
+ filter_slab_addr(l2_filter, Kw, K_TILES, n_my_mt, tile->mt_idx, kh_n, kw_n, kt_base + ktc_n);
+ tensor_load(false, false, A_BUFS[(iter + 1) & 1], 0, 0, a_addr_n, 0, (uint64_t) (TILE - 1), a_row_stride,
+ 0);
+ }
+
+ tensor_fma(false, 3, (uint64_t) (TILE - 1), (uint64_t) (TILE - 1), 0, false, false, false, false, LSCP_B,
+ A_BUFS[iter & 1], 0, first_fma_clears_tenc && (iter == 0));
+ tensor_wait(TENSOR_FMA_WAIT);
+ }
+}
+
+/* Pack only the slabs this shire's tiles actually consume, into local
+ L2 SCP. Slab layout in the filter buffer is [Kh][Kw][n_my_mt][K_TILES]
+ of TILE×TILE slabs (Cin-innermost form). Distributed across the 32
+ hart-1's of this shire by `slab % 32 == minion`.
+
+ This deliberately favors local inner-loop reads over global filter fanout.
+ Depending on tile shape, two shires may pack the same mt value; keep that
+ tradeoff visible when experimenting with shared-filter layouts. */
+static void pack_filter_local_mt(const float * flt_base,
+ int Kh,
+ int Kw,
+ int Cin,
+ int K_TILES,
+ const int * my_mt,
+ int n_my_mt,
+ int minion,
+ uint64_t l2_filter_base) {
+ const int n_slabs = Kh * Kw * n_my_mt * K_TILES;
+ const size_t kstep = (size_t) Kh * Kw; /* Cin stride in floats */
+
+ for (int slab = minion; slab < n_slabs; slab += N_MIN_PER_SHIRE) {
+ int t = slab;
+ const int kt = t % K_TILES;
+ t /= K_TILES;
+ const int mt_idx = t % n_my_mt;
+ t /= n_my_mt;
+ const int kw = t % Kw;
+ t /= Kw;
+ const int kh = t;
+ const int mt = my_mt[mt_idx];
+
+ const uint64_t slab_offset = (uint64_t) slab * SLAB_BYTES;
+ float * cell = (float *) (l2_filter_base + slab_offset);
+
+ for (int oc_in = 0; oc_in < TILE; ++oc_in) {
+ const int oc = mt * TILE + oc_in;
+ const float * src = flt_base + (((size_t) oc * Cin + (size_t) kt * TILE) * Kh + kh) * Kw + kw;
+ float * row = cell + (size_t) oc_in * TILE;
+ float scratch[TILE] __attribute__((aligned(32)));
+ for (int ic_in = 0; ic_in < TILE; ++ic_in) {
+ scratch[ic_in] = src[(size_t) ic_in * kstep];
+ }
+ vec_copy_aligned(row, scratch, TILE);
+ }
+ }
+
+ /* Flush this hart's dirty L1D lines for the slabs it wrote. */
+ FENCE;
+ for (int slab = minion; slab < n_slabs; slab += N_MIN_PER_SHIRE) {
+ const uint64_t slab_offset = (uint64_t) slab * SLAB_BYTES;
+ flush_to_l2((const void *) (l2_filter_base + slab_offset), SLAB_LINES, 64);
+ }
+ WAIT_CACHEOPS;
+}
+
+/* Pack one Cin chunk of the input pin (Kw shifted padded copies) into the
+ buf_idx side of local L2SCP. Work distributed across the 32 hart-1's in
+ the shire by `plane % 32 == minion`. The final flush_to_l2 forces L1D
+ write-back so hart-0's tensor_load sees the freshly written bytes. */
+static void pack_pin_chunk(const pin_ctx_t * ctx, int chunk_id, int buf_idx) {
+ const int kt_base = chunk_id * ctx->chunk_KT;
+ const int Kw = ctx->Kw;
+ const int chunk_KT = ctx->chunk_KT;
+ const int H = ctx->H, W = ctx->W, Hp = ctx->Hp, Wp_a = ctx->Wp_a;
+ const int pad_h = ctx->pad_h, pad_w = ctx->pad_w, s0 = ctx->s0;
+ const int minion = ctx->minion;
+
+ /* Pin pack: Kw shifted, padded copies of input rows. Bounds [vlo, vhi)
+ hoisted outside the row loop so the inner loop is three regions
+ (zero-prefix | bulk-copy | zero-suffix) with no per-element predicate. */
+ float * pin0 = (float *) ctx->l2_pad_in_buf[buf_idx];
+ const int chunk_Cin = chunk_KT * TILE;
+ const int n_pin_planes = Kw * chunk_Cin;
+ for (int p = minion; p < n_pin_planes; p += N_MIN_PER_SHIRE) {
+ const int s = p / chunk_Cin;
+ const int icc = p % chunk_Cin;
+ const int ic = kt_base * TILE + icc;
+ float * pin_s = pin0 + (size_t) s * ctx->pin_copy_floats;
+
+ const int offset = s - pad_w;
+ int vlo = 0;
+ while (vlo < Wp_a && (s0 * vlo + offset) < 0) {
+ vlo++;
+ }
+ int vhi = Wp_a;
+ while (vhi > vlo && (s0 * (vhi - 1) + offset) >= W) {
+ vhi--;
+ }
+ const bool aligned = (s0 == 1) && ((vlo & 7) == 0) && (((vlo + offset) & 7) == 0);
+
+ for (int r = 0; r < Hp; ++r) {
+ float * row = pin_s + ((size_t) icc * Hp + r) * Wp_a;
+ const int real_h = r - pad_h;
+ if (real_h < 0 || real_h >= H) {
+ vec_zero_aligned(row, Wp_a);
+ continue;
+ }
+ const float * src_row = ctx->in_base + ((size_t) ic * H + real_h) * W;
+
+ for (int cc = 0; cc < vlo; ++cc) {
+ row[cc] = 0.0f;
+ }
+
+ if (aligned) {
+ vec_copy_aligned(row + vlo, src_row + vlo + offset, vhi - vlo);
+ } else if (s0 == 1) {
+ const float * csrc = src_row + vlo + offset;
+ const int n = vhi - vlo;
+ for (int cc = 0; cc < n; ++cc) {
+ row[vlo + cc] = csrc[cc];
+ }
+ } else {
+ for (int cc = vlo; cc < vhi; ++cc) {
+ row[cc] = src_row[s0 * cc + offset];
+ }
+ }
+
+ for (int cc = vhi; cc < Wp_a; ++cc) {
+ row[cc] = 0.0f;
+ }
+ }
+ }
+
+ /* Flush this buffer's L1D-dirty lines down to L2SCP backing. */
+ FENCE;
+ flush_range_to_l2((const void *) ctx->l2_pad_in_buf[buf_idx], ctx->pin_chunk_bytes);
+ WAIT_CACHEOPS;
+}
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ (void) env;
+
+ const int shire = get_shire_id();
+ const int hart_id = get_hart_id();
+ const int minion = (hart_id >> 1) & 0x1F;
+ const int hart1 = hart_id & 1;
+
+ const struct ggml_tensor * flt = ¶ms->src0; /* [Kw,Kh,Cin,Cout] */
+ const struct ggml_tensor * in = ¶ms->src1; /* [W, H, Cin,N=1 ] */
+ struct ggml_tensor * out = ¶ms->dst; /* [W, H, Cout,N=1] */
+
+ const int Kw = (int) flt->ne[0];
+ const int Kh = (int) flt->ne[1];
+ const int Cin = (int) flt->ne[2];
+ const int Cout = (int) flt->ne[3];
+
+ const int W = (int) in->ne[0];
+ const int H = (int) in->ne[1];
+ const int OW = (int) out->ne[0];
+ const int OH = (int) out->ne[1];
+
+ /* op_params layout (set by ggml_conv_2d):
+ [0]=s0 [1]=s1 [2]=p0 [3]=p1 [4]=d0 [5]=d1 */
+ const int s0 = out->op_params[0];
+ const int s1 = out->op_params[1];
+ const int pad_w = out->op_params[2];
+ const int pad_h = out->op_params[3];
+
+ if (Cin <= 0 || Cout <= 0) {
+ return -1;
+ }
+ if (Cin % TILE != 0 || Cout % TILE != 0) {
+ return -1;
+ }
+ if (W <= 0 || H <= 0) {
+ return -1;
+ }
+ if (s0 <= 0 || s1 <= 0) {
+ return -1;
+ }
+ if (in->ne[2] != Cin || in->ne[3] != 1) {
+ return -1;
+ }
+ if (out->ne[2] != Cout || out->ne[3] != 1) {
+ return -1;
+ }
+ if (!flt->data || !in->data || !out->data) {
+ return -1;
+ }
+
+ const int K_TILES = Cin / TILE;
+ const int M_TILES = Cout / TILE;
+
+ const int Hp = H + 2 * pad_h;
+ const int Wp_a = round_up_tile_i32(OW);
+ const int OW_pad = Wp_a;
+ const bool need_stage = (OW % TILE != 0);
+
+ /* ===================== Tile assignment & active-shire selection =====
+ Computed up front because the per-shire mt set (and thus filter
+ region size) depends on n_active_shires. */
+ const int w_tiles = ceil_div_i32(OW, TILE);
+ const int total_tiles = OH * w_tiles * M_TILES;
+ const int n_active_shires = need_stage ? 1 : min_i32(total_tiles, N_SHIRES);
+
+ /* Inactive shires exit immediately. No global barrier — pack and
+ barriers are now per-shire, so unused shires don't need to vote. */
+ if (shire >= n_active_shires) {
+ return 0;
+ }
+
+ /* ===================== Determine this shire's mt set ================
+ Standard tile assignment: tile t is owned by
+ shire = t % n_active_shires
+ minion = (t / n_active_shires) % N_MIN_PER_SHIRE
+ slot = t / (n_active_shires * N_MIN_PER_SHIRE)
+ So the set of mt's this shire actually consumes is the set of
+ (t % M_TILES) for all t this shire owns. Enumerate all shire-owned
+ tiles, not just the first MAX_TILES_PER_HART slots; the one-chunk
+ path can process more tiles serially. */
+ int my_mt[MAX_MY_MT];
+ int n_my_mt = 0;
+ for (int t = shire; t < total_tiles; t += n_active_shires) {
+ const int mt = t % M_TILES;
+ bool found = false;
+ for (int j = 0; j < n_my_mt; ++j) {
+ if (my_mt[j] == mt) {
+ found = true;
+ break;
+ }
+ }
+ if (!found) {
+ if (n_my_mt >= MAX_MY_MT) {
+ return -1;
+ }
+ my_mt[n_my_mt++] = mt;
+ }
+ }
+ if (n_my_mt == 0) {
+ return 0; /* no tiles for this shire */
+ }
+
+ const uint64_t filter_local_bytes = (uint64_t) Kh * Kw * n_my_mt * K_TILES * SLAB_BYTES;
+ if (filter_local_bytes > LOCAL_FILTER_CAP) {
+ return -1;
+ }
+
+ /* ===================== L2 SCP local layout =========================
+ filter (this shire's mt slice) | pin_buf[0] | pin_buf[1]?
+ | output_stage? | scratch (streaming) */
+ const uint64_t l2_base = (uint64_t) et_shire_l2scp_local(0);
+ const uint64_t l2_filter = l2_base;
+
+ /* Sizing for pin: budget = LOCAL_BUDGET - filter - output_stage. */
+ const int64_t output_stage_bytes_full = need_stage ? (int64_t) Cout * OH * OW_pad * (int64_t) sizeof(float) : 0;
+ const int64_t budget_for_chunks = (int64_t) LOCAL_BUDGET - (int64_t) filter_local_bytes - output_stage_bytes_full;
+ if (budget_for_chunks <= 0) {
+ return -1;
+ }
+ const int64_t per_KT_pin_bytes = (int64_t) Kw * TILE * Hp * Wp_a * (int64_t) sizeof(float);
+
+ int chunk_KT;
+ int n_buffers;
+ if ((int64_t) K_TILES * per_KT_pin_bytes <= budget_for_chunks) {
+ chunk_KT = K_TILES;
+ n_buffers = 1;
+ } else {
+ chunk_KT = K_TILES;
+ while (chunk_KT > 1 && 2 * (int64_t) chunk_KT * per_KT_pin_bytes > budget_for_chunks) {
+ chunk_KT--;
+ }
+ while (chunk_KT > 1 && K_TILES % chunk_KT != 0) {
+ chunk_KT--;
+ }
+ n_buffers = (chunk_KT < K_TILES) ? 2 : 1;
+ if (chunk_KT < 1) {
+ return -1;
+ }
+ }
+ const int n_chunks = K_TILES / chunk_KT;
+
+ /* Streaming keeps partial sums in MAX_TILES_PER_HART scratch slots per
+ hart. The one-chunk path does not need scratch and can stream a longer
+ tile list serially, but multi-chunk shapes must fit this fixed slot
+ count until scratch scheduling is made more general. */
+ const int shire_tile_capacity = shire + MAX_TILES_PER_HART * n_active_shires * N_MIN_PER_SHIRE;
+ if (n_chunks > 1 && shire_tile_capacity < total_tiles) {
+ return -1;
+ }
+
+ const uint64_t pin_copy_floats = (uint64_t) chunk_KT * TILE * Hp * Wp_a;
+ const uint64_t pin_copy_bytes = pin_copy_floats * sizeof(float);
+ const uint64_t pin_chunk_bytes = (uint64_t) Kw * pin_copy_bytes;
+
+ const uint64_t l2_pin_base = l2_filter + filter_local_bytes;
+ const uint64_t l2_pin_buf[MAX_DBL_BUFS] = {
+ l2_pin_base,
+ l2_pin_base + pin_chunk_bytes,
+ };
+
+ const uint64_t l2_output_stage = need_stage ? l2_pin_base + (uint64_t) n_buffers * pin_chunk_bytes : 0;
+
+ const uint64_t scratch_per_hart = (uint64_t) MAX_TILES_PER_HART * (uint64_t) TILE * TILE * sizeof(float);
+ const uint64_t l2_scratch_base = need_stage ? l2_output_stage + (uint64_t) output_stage_bytes_full :
+ l2_pin_base + (uint64_t) n_buffers * pin_chunk_bytes;
+
+ /* ===================== PHASE 1: Filter pack (per-shire mt slice) ====
+ Hart-1's pack only this shire's mt slabs into local L2 SCP. The
+ SHIRE barrier below ensures the filter is in L2 SCP backing before
+ hart-0's first tensor_load. */
+ if (hart1) {
+ pack_filter_local_mt((const float *) flt->data, Kh, Kw, Cin, K_TILES, my_mt, n_my_mt, minion, l2_filter);
+ }
+
+ /* ===================== Hart 1: pin packer (per chunk) ==============
+ Double-buffered prefetch: pack chunk 0 synchronously, then per chunk c
+ signal "buf c ready", pack chunk c+1 into the alternate buffer
+ (overlaps hart-0's compute on c), signal "buf c done". */
+ if (hart1) {
+ const pin_ctx_t ctx = {
+ .in_base = (const float *) in->data,
+ .Kw = Kw,
+ .chunk_KT = chunk_KT,
+ .H = H,
+ .W = W,
+ .Hp = Hp,
+ .Wp_a = Wp_a,
+ .pad_h = pad_h,
+ .pad_w = pad_w,
+ .s0 = s0,
+ .minion = minion,
+ .pin_copy_floats = pin_copy_floats,
+ .l2_pad_in_buf = { l2_pin_buf[0], l2_pin_buf[1] },
+ .pin_chunk_bytes = pin_chunk_bytes,
+ };
+
+ pack_pin_chunk(&ctx, 0, 0); /* prologue */
+
+ for (int c = 0; c < n_chunks; ++c) {
+ et_barrier(ET_BARRIER_SHIRE); /* signal "buf c ready" */
+ if (n_buffers > 1 && c + 1 < n_chunks) {
+ pack_pin_chunk(&ctx, c + 1, (c + 1) & 1);
+ }
+ et_barrier(ET_BARRIER_SHIRE); /* wait "buf c done" */
+ }
+
+ if (need_stage) {
+ et_barrier(ET_BARRIER_SHIRE);
+ }
+ return 0;
+ }
+
+ /* ===================== Hart 0: matrix engine ======================
+ Two execution modes:
+ - n_chunks == 1: full Cin in one shot. Each hart processes a list
+ of tiles serially; TenC resets between tiles via first_pass=true.
+ - n_chunks > 1: streaming. Each hart owns up to MAX_TILES_PER_HART
+ tiles. For each chunk c, restore TenC from scratch[k] (skip on
+ c==0), accumulate this chunk's FMAs, then either save TenC back
+ to scratch[k] (c < last) or tensor_store directly (c == last). */
+ setup_cache_scp();
+ CLEAR_TENSOR_ERROR;
+
+ char * const out_base = need_stage ? (char *) l2_output_stage : (char *) out->data;
+ const int compute_OW = need_stage ? OW_pad : OW;
+ const uint64_t out_chan_stride = (uint64_t) OH * (uint64_t) compute_OW * sizeof(float);
+ const uint64_t out_row_stride = (uint64_t) compute_OW * sizeof(float);
+
+ const uint64_t a_row_stride = (uint64_t) TILE * sizeof(float); /* 64 */
+ const uint64_t b_row_stride = (uint64_t) Hp * (uint64_t) Wp_a * sizeof(float);
+
+ /* Tile assignment: shire-strided so small workloads spread across
+ shires before stacking minions in one shire. */
+ const int t_start = shire + minion * n_active_shires;
+ const int t_stride = n_active_shires * N_MIN_PER_SHIRE;
+
+ if (n_chunks == 1) {
+ et_barrier(ET_BARRIER_SHIRE); /* wait for the (only) pin chunk */
+
+ const uint64_t l2_pad_in = l2_pin_buf[0];
+ for (int t = t_start; t < total_tiles; t += t_stride) {
+ const conv_tile_t tile = decode_tile(t, M_TILES, w_tiles, my_mt, n_my_mt);
+ compute_tile_chunk(l2_filter, l2_pad_in, pin_copy_bytes, Kh, Kw, K_TILES, chunk_KT, 0, n_my_mt, Hp, Wp_a,
+ s1, a_row_stride, b_row_stride, &tile, /*first_fma_clears_tenc=*/true);
+
+ char * dst_addr = output_tile_addr(out_base, &tile, out_chan_stride, out_row_stride);
+ tensor_store(0, 0, 3, (uint64_t) (TILE - 1), (uint64_t) dst_addr, 0, out_chan_stride);
+ tensor_wait(TENSOR_STORE_WAIT);
+ }
+
+ et_barrier(ET_BARRIER_SHIRE); /* matches hart-1's second barrier */
+
+ } else {
+ /* Streaming path: each hart owns up to MAX_TILES_PER_HART tiles. */
+ int my_tiles[MAX_TILES_PER_HART];
+ int n_my_tiles = 0;
+ for (int slot = 0; slot < MAX_TILES_PER_HART; ++slot) {
+ const int t = t_start + slot * t_stride;
+ if (t < total_tiles) {
+ my_tiles[n_my_tiles++] = t;
+ }
+ }
+
+ conv_tile_t tiles[MAX_TILES_PER_HART];
+ for (int k = 0; k < n_my_tiles; ++k) {
+ tiles[k] = decode_tile(my_tiles[k], M_TILES, w_tiles, my_mt, n_my_mt);
+ }
+
+ const uint64_t my_scratch_base = l2_scratch_base + (uint64_t) minion * scratch_per_hart;
+
+ for (int c = 0; c < n_chunks; ++c) {
+ et_barrier(ET_BARRIER_SHIRE); /* pin chunk c packed */
+
+ const int buf = c & 1;
+ const uint64_t l2_pad_in = l2_pin_buf[buf];
+ const int kt_base = c * chunk_KT;
+
+ for (int k = 0; k < n_my_tiles; ++k) {
+ const conv_tile_t * tile = &tiles[k];
+ const uint64_t scr = my_scratch_base + (uint64_t) k * (TILE * TILE * sizeof(float));
+
+ const bool first_pass_chunk = (c == 0);
+ if (!first_pass_chunk) {
+ tenc_restore_from_scratch(scr);
+ }
+
+ compute_tile_chunk(l2_filter, l2_pad_in, pin_copy_bytes, Kh, Kw, K_TILES, chunk_KT, kt_base, n_my_mt,
+ Hp, Wp_a, s1, a_row_stride, b_row_stride, tile, first_pass_chunk);
+
+ if (c == n_chunks - 1) {
+ char * dst_addr = output_tile_addr(out_base, tile, out_chan_stride, out_row_stride);
+ tensor_store(0, 0, 3, (uint64_t) (TILE - 1), (uint64_t) dst_addr, 0, out_chan_stride);
+ } else {
+ tensor_store(0, 0, 3, (uint64_t) (TILE - 1), (uint64_t) scr, 0, 64);
+ }
+ tensor_wait(TENSOR_STORE_WAIT);
+ }
+
+ et_barrier(ET_BARRIER_SHIRE); /* hart-0 done with chunk c */
+ }
+ }
+
+ FENCE;
+
+ /* ----------------------- DRAM emit phase ---------------------------
+ Only relevant when we staged into L2SCP because OW % 16 != 0. */
+ if (need_stage) {
+ et_barrier(ET_BARRIER_SHIRE);
+
+ if (minion == 0) {
+ const float * stage = (const float *) l2_output_stage;
+ float * dram = (float *) out->data;
+ for (int oc = 0; oc < Cout; ++oc) {
+ for (int oh2 = 0; oh2 < OH; ++oh2) {
+ const float * src = stage + ((size_t) oc * OH + oh2) * OW_pad;
+ float * dst = dram + ((size_t) oc * OH + oh2) * OW;
+ for (int ow2 = 0; ow2 < OW; ++ow2) {
+ dst[ow2] = src[ow2];
+ }
+ }
+ }
+ FENCE;
+ const uint64_t total_bytes = (uint64_t) Cout * OH * OW * sizeof(float);
+ evict_range_past_l2((const void *) dram, total_bytes);
+ WAIT_CACHEOPS;
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// CPY F32 -> F16 Kernel
+// Copies F32 source tensor to F16 destination tensor (contiguous output).
+// Source may have arbitrary strides; destination must be contiguous.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdbool.h>
+#include <stdint.h>
+
+struct ggml_et_cont_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+};
+
+int entry_point(struct ggml_et_cont_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env || !params) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F16) {
+ return -1;
+ }
+
+ const char * src_data = (const char *) src0->data;
+ uint16_t * dst_data = (uint16_t *) dst->data;
+
+ if (!src_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne00 = src0->ne[0];
+ const int64_t ne01 = src0->ne[1];
+ const int64_t ne02 = src0->ne[2];
+ const int64_t ne03 = src0->ne[3];
+
+ const int64_t nb00 = src0->nb[0];
+ const int64_t nb01 = src0->nb[1];
+ const int64_t nb02 = src0->nb[2];
+ const int64_t nb03 = src0->nb[3];
+
+ const int64_t total_elements = ne00 * ne01 * ne02 * ne03;
+
+ if (total_elements == 0) {
+ return 0;
+ }
+
+ // Check if src is contiguous F32
+ const bool src_contiguous =
+ (nb00 == 4 && nb01 == ne00 * 4 && nb02 == ne00 * ne01 * 4 && nb03 == ne00 * ne01 * ne02 * 4);
+
+ // Distribute by cache lines (16 F16 elements = 32 bytes = half cache line)
+ // Use 32 elements per chunk to keep output cache-line aligned
+ const int64_t elems_per_cl = 32;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ if (src_contiguous) {
+ // Fast path: src is contiguous F32
+ const float * src_f32 = (const float *) src_data;
+ for (int64_t i = es; i < ee; ++i) {
+ dst_data[i] = fp32_to_fp16(src_f32[i]);
+ }
+ } else {
+ // General path: stride-aware read
+ for (int64_t idx = es; idx < ee; ++idx) {
+ const int64_t i00 = idx % ne00;
+ const int64_t rem1 = idx / ne00;
+ const int64_t i01 = rem1 % ne01;
+ const int64_t rem2 = rem1 / ne01;
+ const int64_t i02 = rem2 % ne02;
+ const int64_t i03 = rem2 / ne02;
+
+ const float val = *(const float *) (src_data + i00 * nb00 + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ dst_data[idx] = fp32_to_fp16(val);
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+.section .text.init, "ax", @progbits
+.global _start
+_start:
+ # initialize global pointer
+.option push
+.option norelax
+ la gp, __global_pointer$
+.option pop
+ # Firmware sets stack pointer before launch
+ # bss not allowed, no init
+ call entry_point
+ li a2, 0 /* KERNEL_RETURN_SUCCESS (0) */
+ mv a1, a0
+ li a0, 8 /* SYSCALL_RETURN_FROM_KERNEL (8) */
+ ecall
--- /dev/null
+//******************************************************************************
+// CUMSUM F32 Kernel
+// Computes an inclusive prefix sum along dim 0 for each row in higher dims.
+// First-pass implementation: scalar and row-contiguous input/output only.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_cumsum_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+};
+
+int entry_point(struct ggml_et_cumsum_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne0 = src0->ne[0];
+ const int64_t ne1 = src0->ne[1];
+ const int64_t ne2 = src0->ne[2];
+ const int64_t ne3 = src0->ne[3];
+
+ const size_t snb0 = src0->nb[0];
+ const size_t snb1 = src0->nb[1];
+ const size_t snb2 = src0->nb[2];
+ const size_t snb3 = src0->nb[3];
+
+ const size_t dnb0 = dst->nb[0];
+ const size_t dnb1 = dst->nb[1];
+ const size_t dnb2 = dst->nb[2];
+ const size_t dnb3 = dst->nb[3];
+
+ if (snb0 != sizeof(float) || dnb0 != sizeof(float)) {
+ return -1;
+ }
+
+ const int64_t total_rows = ne1 * ne2 * ne3;
+ const int64_t rows_per_group = et_rows_per_cacheline_group(ne0, sizeof(float));
+ const int64_t total_groups = (total_rows + rows_per_group - 1) / rows_per_group;
+
+ for (int64_t grp = thread_id; grp < total_groups; grp += num_threads) {
+ const int64_t row_start = grp * rows_per_group;
+ int64_t row_end = row_start + rows_per_group;
+ if (row_end > total_rows) {
+ row_end = total_rows;
+ }
+
+ for (int64_t row = row_start; row < row_end; ++row) {
+ int64_t i1 = row % ne1;
+ int64_t i2 = (row / ne1) % ne2;
+ int64_t i3 = row / (ne1 * ne2);
+
+ const float * src_row = (const float *) ((const char *) src0_data + i1 * snb1 + i2 * snb2 + i3 * snb3);
+ float * dst_row = (float *) ((char *) dst_data + i1 * dnb1 + i2 * dnb2 + i3 * dnb3);
+
+ float acc = 0.0f;
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ acc += src_row[i0];
+ dst_row[i0] = acc;
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Diag F32 Kernel
+// Creates a diagonal matrix from a 1D vector.
+// dst[i][j] = (i == j) ? src0[i] : 0.0f
+//
+// src0: [N, 1, ne2, ne3] (1D vector per batch)
+// dst: [N, N, ne2, ne3] (diagonal matrix per batch)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_diag_params {
+ struct ggml_tensor src0; // F32 input vector
+ struct ggml_tensor dst; // F32 output diagonal matrix
+};
+
+int entry_point(struct ggml_et_diag_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne0 = dst->ne[0]; // N (row width = column count)
+ const int64_t ne1 = dst->ne[1]; // N (number of rows)
+ const int64_t ne2 = dst->ne[2];
+ const int64_t ne3 = dst->ne[3];
+
+ const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+ const size_t nb02 = src0->nb[2], nb03 = src0->nb[3];
+
+ // Total rows across all batches — parallelize over these
+ const int64_t total_rows = ne1 * ne2 * ne3;
+
+ // Prepare zero vector for SIMD zeroing
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ int64_t i1 = row % ne1;
+ int64_t i2 = (row / ne1) % ne2;
+ int64_t i3 = row / (ne1 * ne2);
+
+ float * dst_row = (float *) ((char *) dst_data + i1 * nb1 + i2 * nb2 + i3 * nb3);
+
+ // Zero the entire row with SIMD
+ int64_t i0 = 0;
+ const int64_t vec_end = (ne0 / 8) * 8;
+ for (; i0 < vec_end; i0 += 8) {
+ __asm__ volatile("fsw.ps f10, %[d]\n" : [d] "=m"(*(float (*)[8]) & dst_row[i0])::"f10");
+ }
+ for (; i0 < ne0; i0++) {
+ dst_row[i0] = 0.0f;
+ }
+
+ // Place the diagonal element: dst[i1][i1] = src0[i1]
+ const float * src_ptr = (const float *) ((const char *) src0_data + i2 * nb02 + i3 * nb03);
+ dst_row[i1] = src_ptr[i1];
+ }
+
+ return 0;
+}
--- /dev/null
+// Element-wise operations: dst[i] = src0[i] op src1[i]
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+// Generic m0-gated element-wise block operation.
+// The OP parameter selects the instruction: "fmul.ps", "fadd.ps", "fsub.ps".
+#define DEFINE_BLOCK_OP(name, op_insn) \
+ static inline void name(float * dst_block, const float * src0_block, const float * src1_block, int elements) { \
+ const int32_t vec_end = (elements / 8) * 8; \
+ const int32_t tail = elements - vec_end; \
+ \
+ unsigned long temp_mask; \
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); \
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); \
+ \
+ for (int32_t i = 0; i < vec_end; i += 8) { \
+ __asm__ volatile( \
+ "flw.ps f10, %[s0]\n" \
+ "flw.ps f11, %[s1]\n" op_insn \
+ " f12, f10, f11\n" \
+ "fsw.ps f12, %[d]\n" \
+ : [d] "=m"(*(float (*)[8]) & dst_block[i]) \
+ : [s0] "m"(*(const float (*)[8]) & src0_block[i]), [s1] "m"(*(const float (*)[8]) & src1_block[i]) \
+ : "f10", "f11", "f12"); \
+ } \
+ /* Deal with tail chunks */ \
+ if (tail > 0) { \
+ const unsigned long tail_m0 = (1ul << tail) - 1; \
+ __asm__ volatile( \
+ "mov.m.x m0, %[tm], 0\n" \
+ "flw.ps f10, 0(%[s0])\n" \
+ "flw.ps f11, 0(%[s1])\n" op_insn \
+ " f12, f10, f11\n" \
+ "fsw.ps f12, 0(%[d])\n" \
+ : \
+ : [s0] "r"(&src0_block[vec_end]), [s1] "r"(&src1_block[vec_end]), [d] "r"(&dst_block[vec_end]), \
+ [tm] "r"(tail_m0) \
+ : "f10", "f11", "f12", "memory"); \
+ } \
+ \
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask)); \
+ }
+
+DEFINE_BLOCK_OP(block_mul_cache_aligned, "fmul.ps")
+DEFINE_BLOCK_OP(block_add_cache_aligned, "fadd.ps")
+DEFINE_BLOCK_OP(block_sub_cache_aligned, "fsub.ps")
+
+// Broadcast variants: src1 is a single scalar, broadcast to all 8 lanes.
+#define DEFINE_BLOCK_OP_BROADCAST(name, op_insn) \
+ static inline void name(float * dst_block, const float * src0_block, float scalar, int elements) { \
+ const int32_t vec_end = (elements / 8) * 8; \
+ const int32_t tail = elements - vec_end; \
+ \
+ unsigned long temp_mask; \
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); \
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); \
+ \
+ for (int32_t i = 0; i < vec_end; i += 8) { \
+ __asm__ volatile( \
+ "flw.ps f10, %[s0]\n" \
+ "fbc.ps f11, %[s]\n" op_insn \
+ " f12, f10, f11\n" \
+ "fsw.ps f12, %[d]\n" \
+ : [d] "=m"(*(float (*)[8]) & dst_block[i]) \
+ : [s0] "m"(*(const float (*)[8]) & src0_block[i]), [s] "m"(scalar) \
+ : "f10", "f11", "f12"); \
+ } \
+ \
+ if (tail > 0) { \
+ const unsigned long tail_m0 = (1ul << tail) - 1; \
+ __asm__ volatile( \
+ "mov.m.x m0, %[tm], 0\n" \
+ "flw.ps f10, 0(%[s0])\n" \
+ "fbc.ps f11, 0(%[ps])\n" op_insn \
+ " f12, f10, f11\n" \
+ "fsw.ps f12, 0(%[d])\n" \
+ : \
+ : [s0] "r"(&src0_block[vec_end]), [ps] "r"(&scalar), [d] "r"(&dst_block[vec_end]), [tm] "r"(tail_m0) \
+ : "f10", "f11", "f12", "memory"); \
+ } \
+ \
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask)); \
+ }
+
+DEFINE_BLOCK_OP_BROADCAST(block_mul_broadcast, "fmul.ps")
+DEFINE_BLOCK_OP_BROADCAST(block_add_broadcast, "fadd.ps")
+DEFINE_BLOCK_OP_BROADCAST(block_sub_broadcast, "fsub.ps")
+
+static inline float scalar_el_map(float src0, float src1, enum ggml_op operation) {
+ switch (operation) {
+ case GGML_OP_MUL:
+ return src0 * src1;
+ case GGML_OP_ADD:
+ return src0 + src1;
+ case GGML_OP_SUB:
+ return src0 - src1;
+ default:
+ return 0.0f;
+ }
+}
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * src1_data = (float *) src1->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+#ifdef ET_UBERKERNEL
+ // Consumer-side input eviction. Required because ET caches are
+ // incoherent across minions: if a previous kernel in this UK batch
+ // left stale lines for these addresses in this hart's L1, drop them
+ // so we read fresh from L3/DRAM (where the producer flushed its
+ // results). Standalone launches don't need this -- the host-side
+ // runtime boundary between kernel launches handles it.
+ const size_t src0_bytes = (size_t) src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3] * src0->nb[0];
+ const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * src1->nb[0];
+ evict_region_past_l2(src0_data, src0_bytes);
+ evict_region_past_l2(src1_data, src1_bytes);
+ WAIT_CACHEOPS;
+ FENCE;
+ et_barrier(ET_BARRIER_GLOBAL);
+#endif
+
+ enum ggml_op operation = dst->op;
+
+ if (operation != GGML_OP_MUL && operation != GGML_OP_ADD && operation != GGML_OP_SUB) {
+ return -1; // Unsupported operation
+ }
+
+ const int64_t ne0 = dst->ne[0], ne1 = dst->ne[1], ne2 = dst->ne[2], ne3 = dst->ne[3];
+ const int64_t ne00 = src0->ne[0], ne01 = src0->ne[1], ne02 = src0->ne[2], ne03 = src0->ne[3];
+ const int64_t ne10 = src1->ne[0], ne11 = src1->ne[1], ne12 = src1->ne[2], ne13 = src1->ne[3];
+
+ const size_t nb0 = dst->nb[0], nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+ const size_t nb00 = src0->nb[0], nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ const size_t nb10 = src1->nb[0], nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
+
+ const bool cache_aligned = (dst->ne[0] % 16 == 0);
+
+ // Fast path: no broadcasting, contiguous
+ const bool no_broadcast = (ne10 == ne0 && ne11 == ne1 && ne12 == ne2 && ne13 == ne3);
+ const bool all_contiguous =
+ (nb0 == 4 && nb00 == 4 && nb10 == 4 && nb1 == ne0 * 4 && nb01 == ne0 * 4 && nb11 == ne0 * 4);
+
+ if (no_broadcast && all_contiguous) {
+ const int64_t total_elements = ne0 * ne1 * ne2 * ne3;
+ const int64_t elements_per_cacheline = 16; // 64 bytes / 4 bytes
+ const int64_t total_cachelines = (total_elements + elements_per_cacheline - 1) / elements_per_cacheline;
+
+ const int64_t cl_per_thread = (total_cachelines + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cachelines) {
+ cl_end = total_cachelines;
+ }
+
+ if (cl_start >= total_cachelines) {
+ return 0;
+ }
+
+ const int64_t elem_start = cl_start * elements_per_cacheline;
+ int64_t elem_end = cl_end * elements_per_cacheline;
+ if (elem_end > total_elements) {
+ elem_end = total_elements;
+ }
+ const int32_t count = (int32_t) (elem_end - elem_start);
+
+ switch (operation) {
+ case GGML_OP_MUL:
+ block_mul_cache_aligned(dst_data + elem_start, src0_data + elem_start, src1_data + elem_start, count);
+ break;
+ case GGML_OP_ADD:
+ block_add_cache_aligned(dst_data + elem_start, src0_data + elem_start, src1_data + elem_start, count);
+ break;
+ case GGML_OP_SUB:
+ block_sub_cache_aligned(dst_data + elem_start, src0_data + elem_start, src1_data + elem_start, count);
+ break;
+ default:
+ return 1;
+ }
+#ifdef ET_UBERKERNEL
+ // Producer-side flush: ET caches are incoherent across minions, so
+ // a consumer kernel running on a different minion can't see our
+ // dirty L1 lines via its own evict_region_past_l2. Push our writes
+ // all the way to DRAM so the next batched kernel reads fresh.
+ // Standalone launches don't need this -- the host runtime boundary
+ // between kernel launches handles cache writeback.
+ FENCE;
+ evict_region_past_l2(dst_data + elem_start, (size_t) count * sizeof(float));
+ WAIT_CACHEOPS;
+ FENCE;
+#endif
+ return 0;
+ }
+
+ // Slow path: broadcasting or non-contiguous
+ const int64_t total_rows = ne1 * ne2 * ne3;
+
+ int64_t start_row;
+ int64_t end_row;
+
+ if (cache_aligned) {
+ const int64_t rows_per_thread = (total_rows + num_threads - 1) / num_threads;
+ start_row = thread_id * rows_per_thread;
+ end_row = (start_row + rows_per_thread < total_rows) ? (start_row + rows_per_thread) : total_rows;
+ } else {
+ const int64_t rows_per_group = et_rows_per_cacheline_group(ne0, sizeof(float));
+ const int64_t total_groups = (total_rows + rows_per_group - 1) / rows_per_group;
+
+ if (thread_id >= total_groups) {
+ return 0;
+ }
+
+ const int64_t group_start = thread_id;
+ for (int64_t grp = group_start; grp < total_groups; grp += num_threads) {
+ const int64_t group_row_start = grp * rows_per_group;
+ int64_t group_row_end = group_row_start + rows_per_group;
+ if (group_row_end > total_rows) {
+ group_row_end = total_rows;
+ }
+
+#ifdef ET_UBERKERNEL
+ // First row written by this group (used for producer-side evict).
+ const int64_t first_i03 = group_row_start / (ne2 * ne1);
+ const int64_t first_i02 = (group_row_start - first_i03 * ne2 * ne1) / ne1;
+ const int64_t first_i01 = (group_row_start - first_i03 * ne2 * ne1 - first_i02 * ne1);
+ char * group_dst_base = (char *) dst_data + first_i03 * nb3 + first_i02 * nb2 + first_i01 * nb1;
+#endif
+
+ for (int64_t ir = group_row_start; ir < group_row_end; ir++) {
+ const int64_t i03 = ir / (ne2 * ne1);
+ const int64_t i02 = (ir - i03 * ne2 * ne1) / ne1;
+ const int64_t i01 = (ir - i03 * ne2 * ne1 - i02 * ne1);
+
+ const int64_t i13 = i03 % ne13;
+ const int64_t i12 = i02 % ne12;
+ const int64_t i11 = i01 % ne11;
+
+ float * dst_ptr = (float *) ((char *) dst_data + i03 * nb3 + i02 * nb2 + i01 * nb1);
+ const float * src0_ptr =
+ (const float *) ((const char *) src0_data + i03 * nb03 + i02 * nb02 + i01 * nb01);
+ const float * src1_ptr =
+ (const float *) ((const char *) src1_data + i13 * nb13 + i12 * nb12 + i11 * nb11);
+
+ if (ne10 == 1) {
+ const float scalar = src1_ptr[0];
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ dst_ptr[i0] = scalar_el_map(src0_ptr[i0], scalar, operation);
+ }
+ } else {
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ dst_ptr[i0] = scalar_el_map(src0_ptr[i0], src1_ptr[i0 % ne10], operation);
+ }
+ }
+ }
+
+#ifdef ET_UBERKERNEL
+ // Producer-side flush for this group's rows. Group rows are
+ // contiguous because nb1 = ne0*4 in the cacheline-group layout.
+ // Only needed inside a UK batch; see comment in fast path.
+ const int64_t nrows = group_row_end - group_row_start;
+ if (nrows > 0) {
+ FENCE;
+ evict_region_past_l2(group_dst_base, (size_t) nrows * nb1);
+ WAIT_CACHEOPS;
+ FENCE;
+ }
+#endif
+ }
+
+ return 0;
+ }
+
+ if (start_row >= total_rows) {
+ return 0;
+ }
+
+ for (int64_t ir = start_row; ir < end_row; ir++) {
+ // Convert flat row index to 3D coordinates
+ const int64_t i03 = ir / (ne2 * ne1);
+ const int64_t i02 = (ir - i03 * ne2 * ne1) / ne1;
+ const int64_t i01 = (ir - i03 * ne2 * ne1 - i02 * ne1);
+
+ // Handle broadcasting: src1 coordinates with modulo
+ const int64_t i13 = i03 % ne13;
+ const int64_t i12 = i02 % ne12;
+ const int64_t i11 = i01 % ne11;
+
+ // Calculate base pointers for this row using stride-based addressing
+ float * dst_ptr = (float *) ((char *) dst_data + i03 * nb3 + i02 * nb2 + i01 * nb1);
+ const float * src0_ptr = (const float *) ((const char *) src0_data + i03 * nb03 + i02 * nb02 + i01 * nb01);
+ const float * src1_ptr = (const float *) ((const char *) src1_data + i13 * nb13 + i12 * nb12 + i11 * nb11);
+
+ if (ne10 == 1) {
+ // Broadcast scalar: src1 has ne[0]=1, broadcast across entire row
+ float scalar = src1_ptr[0];
+ switch (operation) {
+ case GGML_OP_MUL:
+ block_mul_broadcast(dst_ptr, src0_ptr, scalar, (int) ne0);
+ break;
+ case GGML_OP_ADD:
+ block_add_broadcast(dst_ptr, src0_ptr, scalar, (int) ne0);
+ break;
+ case GGML_OP_SUB:
+ block_sub_broadcast(dst_ptr, src0_ptr, scalar, (int) ne0);
+ break;
+ default:
+ return 1;
+ }
+ } else {
+ // Broadcasting in dimension 0: src1 repeats across src0
+ const int64_t nr0 = ne0 / ne10;
+
+ for (int64_t r = 0; r < nr0; r++) {
+ const float * src0_block = src0_ptr + r * ne10;
+ float * dst_block = dst_ptr + r * ne10;
+
+ switch (operation) {
+ case GGML_OP_MUL:
+ block_mul_cache_aligned(dst_block, src0_block, src1_ptr, (int) ne10);
+ break;
+ case GGML_OP_ADD:
+ block_add_cache_aligned(dst_block, src0_block, src1_ptr, (int) ne10);
+ break;
+ case GGML_OP_SUB:
+ block_sub_cache_aligned(dst_block, src0_block, src1_ptr, (int) ne10);
+ break;
+ default:
+ return 1;
+ }
+ }
+ }
+ }
+
+#ifdef ET_UBERKERNEL
+ // Producer-side flush for the cache-aligned slow path. Rows
+ // [start_row, end_row) are contiguous in dst because nb1 = ne0 * 4.
+ // Only needed inside a UK batch; see comment in fast path.
+ if (end_row > start_row) {
+ FENCE;
+ evict_region_past_l2((char *) dst_data + start_row * nb1, (size_t) (end_row - start_row) * nb1);
+ WAIT_CACHEOPS;
+ FENCE;
+ }
+#endif
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Fill F32 Kernel
+// Fills entire tensor with a constant scalar value.
+// dst[i] = c for all elements
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_fill_params {
+ struct ggml_tensor dst; // F32 output tensor (contiguous)
+ float c; // Constant value to fill
+};
+
+int entry_point(struct ggml_et_fill_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * dst_data = (float *) dst->data;
+ if (!dst_data) {
+ return -1;
+ }
+
+ const int64_t total_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+
+ if (total_elements == 0) {
+ return 0;
+ }
+
+ // Distribute by cache lines (16 floats = 64 bytes)
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ // Broadcast constant to all SIMD lanes
+ float c = params->c;
+ __asm__ volatile("fbc.ps f10, %[v]\n" : : [v] "m"(c) : "f10");
+
+ // Vector fill (8-wide)
+ int64_t i = es;
+ const int64_t vec_end = es + ((ee - es) / 8) * 8;
+ for (; i < vec_end; i += 8) {
+ __asm__ volatile("fsw.ps f10, %[d]\n" : [d] "=m"(*(float (*)[8]) & dst_data[i])::"f10");
+ }
+ // Scalar tail
+ for (; i < ee; i++) {
+ dst_data[i] = c;
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Flash Attention with TensorFMA16A32 for QK^T
+//
+// Uses the matrix engine for the QK^T dot products (F16×F16→F32),
+// scalar code for online softmax and V accumulation.
+//
+// Hart 0: tensor engine (Q load, K load from SCP, FMA, softmax, V accum)
+// Hart 1: pack K into double-buffered L2 SCP panels, flush for tensor_load
+//
+// Requirements:
+// - Q: F32 (converted to F16 internally)
+// - K, V: F16
+// - dk must be a multiple of 32 (TensorFMA16A32 K-tile)
+// - dv ≤ 512 (accumulator in shire-local L2 SCP)
+//
+// Parallelization: each minion independently processes one (qpos, head, batch)
+// row, round-robin across all minion hart-0s. Hart 1 assists with K packing.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "tensor.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+#include <string.h>
+
+#define NUM_COMPUTE_SHIRES 32
+#define MINIONS_PER_SHIRE 32
+
+// QK^T tiles: 16 KV positions at a time, K in chunks of 32 F16
+#define TILE_KV 16
+#define TILE_K 32
+
+// L1 scratchpad layout: A (Q) in lines 0-15, B (K interleaved) in lines 16-31
+#define A_L1_START 0
+#define B_L1_START 16
+
+// Max head dimensions
+#define FA_DV_MAX 512 // max value head dim (dv)
+#define FA_DK_MAX 512 // max key head dim (dk) - some models use hsk > hsv
+
+typedef uint16_t et_fp16_t;
+
+#define ET_NEG_INF_F (-3.402823466e+38f)
+
+// L2 SCP layout per minion:
+// [0..2047] accumulator (FA_DV_MAX * sizeof(float))
+// [2048..4095] kpanel buffer 0 (32 × 32 × 2 = 2048 bytes)
+// [4096..6143] kpanel buffer 1 (2048 bytes)
+// [6144..6207] stats line - (M_p at +0, S_p at +4), own cache line
+// Double-buffering ensures hart 0 finishes buf[N%2] before hart 1
+// overwrites it at chunk N+2.
+//
+// The stats line reserves a cache-line-aligned slot for split-KV softmax
+// partials (M_p, S_p). With k_splits=1 the slot is currently unused; step 2
+// will populate it and use peer minions' slots during the reduction.
+#define SCP_ACC_OFF 0
+#define SCP_ACC_STRIDE (FA_DV_MAX * sizeof(float)) // 2048
+#define SCP_KPANEL_SIZE (32 * 32 * sizeof(et_fp16_t)) // 2048
+#define SCP_KP0_OFF SCP_ACC_STRIDE // 2048
+#define SCP_KP1_OFF (SCP_KP0_OFF + SCP_KPANEL_SIZE) // 4096
+#define SCP_STATS_OFF (SCP_KP1_OFF + SCP_KPANEL_SIZE) // 6144
+#define SCP_STATS_SIZE 64 // own cache line
+#define SCP_PER_MINION (SCP_STATS_OFF + SCP_STATS_SIZE) // 6208
+
+struct ggml_et_flash_attn_ext_params {
+ struct ggml_tensor src0; // Q (F32)
+ struct ggml_tensor src1; // K (F16)
+ struct ggml_tensor src2; // V (F16)
+ struct ggml_tensor mask; // mask (F16 or F32), zeroed when absent
+ struct ggml_tensor dst; // Output (F32)
+ float scale;
+ int32_t has_mask;
+};
+
+static inline float get_mask_val(const struct ggml_tensor * mask, int64_t iq1, int64_t ik1, int64_t iq2, int64_t iq3) {
+ const char * base = (const char *) mask->data + iq1 * mask->nb[1] + (iq2 % mask->ne[2]) * mask->nb[2] +
+ (iq3 % mask->ne[3]) * mask->nb[3];
+
+ if (mask->type == GGML_TYPE_F32) {
+ return *(const float *) (base + ik1 * mask->nb[0]);
+ }
+ return fp16_to_fp32(*(const uint16_t *) (base + ik1 * mask->nb[0]));
+}
+
+static inline const char * get_mask_row_base(const struct ggml_tensor * mask, int64_t iq1, int64_t iq2, int64_t iq3) {
+ return (const char *) mask->data + iq1 * mask->nb[1] + (iq2 % mask->ne[2]) * mask->nb[2] +
+ (iq3 % mask->ne[3]) * mask->nb[3];
+}
+
+static inline float get_mask_val_from_base(const struct ggml_tensor * mask, const char * base, int64_t ik1) {
+ if (mask->type == GGML_TYPE_F32) {
+ return *(const float *) (base + ik1 * mask->nb[0]);
+ }
+ return fp16_to_fp32(*(const uint16_t *) (base + ik1 * mask->nb[0]));
+}
+
+// Pack K rows for TensorLoadTranspose16 (even/odd deinterleave)
+static inline void __attribute__((always_inline)) pack_k_for_transpose16(et_fp16_t * out,
+ const char * k_base,
+ int64_t kv_start,
+ int64_t dk_start,
+ int64_t kv_count,
+ int64_t nb1_k) {
+ unsigned long old_mask;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ : [ms] "=&r"(old_mask)
+ :
+ :);
+
+ for (int j = 0; j < (int) kv_count; ++j) {
+ const et_fp16_t * k_row = (const et_fp16_t *) (k_base + (kv_start + j) * nb1_k) + dk_start;
+ et_fp16_t * even_row = out + (j * 2) * 32;
+ et_fp16_t * odd_row = out + (j * 2 + 1) * 32;
+ __asm__ volatile(
+ "flw.ps f2, 0(%[src0]) \n\t" // load row[0..15]
+ "flw.ps f3, 0(%[src1]) \n\t" // load row[16..31]
+ "fpackreph.pi f4, f2 \n\t" // even_lo from src0
+ "fpackreph.pi f6, f3 \n\t" // even_lo from src1 (interleaved)
+ "fsrli.pi f5, f2, 16 \n\t" // shift src0 for odd
+ "fsrli.pi f7, f3, 16 \n\t" // shift src1 for odd (interleaved)
+ "fpackreph.pi f5, f5 \n\t" // odd from src0
+ "fpackreph.pi f7, f7 \n\t" // odd from src1
+ "mov.m.x m0, x0, 0x0F \n\t"
+ "fcmovm.ps f4, f4, f6 \n\t" // merge even halves
+ "fcmovm.ps f5, f5, f7 \n\t" // merge odd halves
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fsw.ps f4, 0(%[even]) \n\t"
+ "fsw.ps f5, 0(%[odd]) \n\t"
+ :
+ : [src0] "r"(k_row), [src1] "r"(k_row + 16), [even] "r"(even_row), [odd] "r"(odd_row)
+ : "f2", "f3", "f4", "f5", "f6", "f7", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %[ms] \n\t" : : [ms] "r"(old_mask));
+
+ for (int j = (int) kv_count; j < TILE_KV; ++j) {
+ et_fp16_t * even_row = out + (j * 2) * 32;
+ et_fp16_t * odd_row = out + (j * 2 + 1) * 32;
+ for (int l = 0; l < TILE_K / 2; ++l) {
+ even_row[l] = 0;
+ odd_row[l] = 0;
+ }
+ }
+}
+
+// Build interleaved B panel for TensorFMA16A32 (weights @ V).
+static inline void __attribute__((always_inline)) pack_v_interleaved(et_fp16_t * out,
+ const char * v_head,
+ int64_t kv_base,
+ int64_t dv_start,
+ int64_t kv_count,
+ int64_t nb1_v) {
+ for (int k = 0; k < TILE_KV; ++k) {
+ const int l = k >> 1;
+ const int r = k & 1;
+ et_fp16_t * const dst = out + l * 32 + r;
+ if (k < (int) kv_count) {
+ const et_fp16_t * v_row = (const et_fp16_t *) (v_head + (kv_base + k) * nb1_v) + dv_start;
+ for (int n = 0; n < 16; ++n) {
+ dst[n * 2] = v_row[n];
+ }
+ } else {
+ for (int n = 0; n < 16; ++n) {
+ dst[n * 2] = 0;
+ }
+ }
+ }
+}
+
+// Prefetch KV rows for one chunk into L2.
+static inline void __attribute__((always_inline)) prefetch_kv_to_l2(const char * head,
+ int64_t kv_start,
+ int64_t d_start,
+ int64_t kv_count,
+ int64_t nb1) {
+ const void * base = (const void *) (head + kv_start * nb1 + d_start * 2);
+ l2_prefetch(base, (uint64_t) kv_count, (uint64_t) nb1);
+}
+
+static inline void __attribute__((always_inline)) convert_q_row_f32_to_f16(et_fp16_t * dst,
+ const float * src,
+ int64_t n) {
+ static const int32_t __attribute__((aligned(32))) offsets[8] = { 0, 2, 4, 6, 8, 10, 12, 14 };
+
+ unsigned long old_mask;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f1, 0(%[offs]) \n\t"
+ : [ms] "=&r"(old_mask)
+ : [offs] "r"(offsets)
+ : "f1");
+
+ for (int64_t d = 0; d < n; d += 8) {
+ __asm__ volatile(
+ "flw.ps f2, 0(%[src]) \n\t"
+ "fcvt.f16.ps f3, f2 \n\t"
+ "fsch.ps f3, f1(%[dst]) \n\t"
+ :
+ : [src] "r"(src + d), [dst] "r"(dst + d)
+ : "f2", "f3", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %[ms] \n\t" : : [ms] "r"(old_mask));
+}
+
+static inline void __attribute__((always_inline)) zero_acc_vec(float * acc, int64_t dv) {
+ const float zero = 0.0f;
+ unsigned long old_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(old_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+ __asm__ volatile("fbc.ps f2, 0(%[z])" ::[z] "r"(&zero) : "f2");
+
+ for (int64_t d = 0; d < dv; d += 8) {
+ __asm__ volatile("fsw.ps f2, 0(%[a]) \n\t" ::[a] "r"(acc + d) : "f2", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(old_mask));
+}
+
+static inline void __attribute__((always_inline)) scale_acc_vec(float * acc, int64_t dv, float scale) {
+ unsigned long old_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(old_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ for (int64_t d = 0; d < dv; d += 8) {
+ __asm__ volatile(
+ "fbc.ps f2, 0(%[s]) \n\t"
+ "flw.ps f3, 0(%[a]) \n\t"
+ "fmul.ps f3, f3, f2 \n\t"
+ "fsw.ps f3, 0(%[a]) \n\t"
+ :
+ : [s] "r"(&scale), [a] "r"(acc + d)
+ : "f2", "f3", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(old_mask));
+}
+
+static inline void __attribute__((always_inline)) normalize_store_vec(float * out,
+ float * acc,
+ int64_t dv,
+ float inv,
+ int use_fast_store) {
+ unsigned long old_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(old_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ for (int64_t d = 0; d < dv; d += 8) {
+ __asm__ volatile(
+ "fbc.ps f2, 0(%[inv]) \n\t"
+ "flw.ps f3, 0(%[a]) \n\t"
+ "fmul.ps f3, f3, f2 \n\t"
+ "fsw.ps f3, 0(%[a]) \n\t"
+ :
+ : [inv] "r"(&inv), [a] "r"(acc + d)
+ : "f2", "f3", "memory");
+ if (use_fast_store) {
+ __asm__ volatile(
+ "flw.ps f4, 0(%[a]) \n\t"
+ "fsw.ps f4, 0(%[o]) \n\t"
+ :
+ : [a] "r"(acc + d), [o] "r"(out + d)
+ : "f4", "memory");
+ } else {
+ atomic_store_f32((volatile float *) &out[d + 0], acc[d + 0]);
+ atomic_store_f32((volatile float *) &out[d + 1], acc[d + 1]);
+ atomic_store_f32((volatile float *) &out[d + 2], acc[d + 2]);
+ atomic_store_f32((volatile float *) &out[d + 3], acc[d + 3]);
+ atomic_store_f32((volatile float *) &out[d + 4], acc[d + 4]);
+ atomic_store_f32((volatile float *) &out[d + 5], acc[d + 5]);
+ atomic_store_f32((volatile float *) &out[d + 6], acc[d + 6]);
+ atomic_store_f32((volatile float *) &out[d + 7], acc[d + 7]);
+ }
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(old_mask));
+}
+
+static inline size_t tensor_bytes_fa(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+// Evict a byte range from L1D to L2 SCP, splitting into batches of ≤16
+// cache lines (the hw limit for evict_to_l2). Use before a barrier when
+// another minion in the shire needs to read the region, or after a barrier
+// on the reader side to drop stale L1D copies before reading peer data.
+static inline void __attribute__((always_inline)) evict_range_to_l2(const void * addr, int64_t bytes) {
+ if (bytes <= 0) {
+ return;
+ }
+ int64_t lines = (bytes + 63) / 64;
+ const char * p = (const char *) addr;
+ while (lines > 0) {
+ int64_t batch = lines > 16 ? 16 : lines;
+ evict_to_l2((const void *) p, (uint64_t) batch, 64);
+ p += batch * 64;
+ lines -= batch;
+ }
+}
+
+// Split-KV online merge inner loop:
+//
+// for d in [0, dv) step 8:
+// acc[d..d+8] = alpha_own * acc[d..d+8] + alpha_peer * peer_acc[d..d+8]
+//
+// Runs on the reducer (k_split == 0) after all tensor_fma ops for the row are
+// complete, so f0..f31 are dead at entry. We still bracket the loop in inline
+// asm with explicit f2/f3/f4/f5 clobbers to lock register usage down — per the
+// MM register lifetime rule, never let the compiler mingle FP ops into code
+// that sits anywhere near a tensor engine output window.
+static inline void __attribute__((always_inline)) merge_rescale_add_asm(float * acc,
+ const float * peer_acc,
+ int64_t dv,
+ float alpha_own,
+ float alpha_peer) {
+ unsigned long old_mask;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f4, 0(%[ao]) \n\t" // broadcast alpha_own
+ "fbc.ps f5, 0(%[ap]) \n\t" // broadcast alpha_peer
+ : [ms] "=&r"(old_mask)
+ : [ao] "r"(&alpha_own), [ap] "r"(&alpha_peer)
+ : "f4", "f5");
+
+ for (int64_t d = 0; d < dv; d += 8) {
+ __asm__ volatile(
+ "flw.ps f2, 0(%[a]) \n\t" // own
+ "flw.ps f3, 0(%[p]) \n\t" // peer
+ "fmul.ps f2, f2, f4 \n\t" // own *= alpha_own
+ "fmul.ps f3, f3, f5 \n\t" // peer *= alpha_peer
+ "fadd.ps f2, f2, f3 \n\t"
+ "fsw.ps f2, 0(%[a]) \n\t"
+ :
+ : [a] "r"(acc + d), [p] "r"(peer_acc + d)
+ : "f2", "f3", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(old_mask));
+}
+
+int entry_point(struct ggml_et_flash_attn_ext_params * params, void * env) {
+ (void) env;
+
+ uint64_t hart_id = get_hart_id();
+ uint64_t shire_id = get_shire_id();
+
+ if (shire_id >= NUM_COMPUTE_SHIRES) {
+ return 0;
+ }
+
+ const int is_hart1 = hart_id & 1;
+ uint64_t local_minion = (hart_id >> 1) & 0x1F;
+
+ struct ggml_tensor * q = ¶ms->src0;
+ struct ggml_tensor * k = ¶ms->src1;
+ struct ggml_tensor * v = ¶ms->src2;
+ struct ggml_tensor * dst = ¶ms->dst;
+ const int32_t has_mask = params->has_mask;
+ struct ggml_tensor * mask = has_mask ? ¶ms->mask : (struct ggml_tensor *) 0;
+
+ const char * q_data = (const char *) q->data;
+ const char * k_data = (const char *) k->data;
+ const char * v_data = (const char *) v->data;
+ char * dst_data = (char *) dst->data;
+
+ // et_barrier(ET_BARRIER_GLOBAL);
+ evict_region_past_l2(q->data, tensor_bytes_fa(q));
+ evict_region_past_l2(k->data, tensor_bytes_fa(k));
+ evict_region_past_l2(v->data, tensor_bytes_fa(v));
+ if (mask) {
+ evict_region_past_l2(mask->data, tensor_bytes_fa(mask));
+ }
+ et_barrier(ET_BARRIER_GLOBAL);
+
+ const int64_t dk = q->ne[0];
+ const int64_t nq = q->ne[1];
+ const int64_t nhq = q->ne[2];
+ const int64_t no = q->ne[3];
+ const int64_t nk = k->ne[1];
+ const int64_t nhk = k->ne[2];
+ const int64_t dv = v->ne[0];
+
+ if (dv > FA_DV_MAX || dk > FA_DK_MAX) {
+ return -1;
+ }
+ if (k->nb[0] != 2 || v->nb[0] != 2) {
+ return -1;
+ }
+ if ((dk % 8) != 0 || (dv % 16) != 0) {
+ return -1;
+ }
+
+ const int64_t gqa_ratio = nhq / nhk;
+ const int64_t total_rows = nq * nhq * no;
+ const float scale = params->scale;
+ const int use_fast_store = (dv % 16 == 0);
+
+ // Split-KV team layout (mirrors mul_mat_f16_matrix_engine.c)
+ //
+ // When total_rows is small compared to the total minion count (typical
+ // for decode: nq=1, nhq small), we group k_splits minions within the
+ // same shire into a team that cooperates on one row by splitting the
+ // KV dimension. Each team member computes a partial (M_p, S_p, acc_p)
+ // over its KV slab; the k_split==0 member merges the partials with the
+ // softmax combine rule.
+ //
+ // k_splits is a power of two, capped at MINIONS_PER_SHIRE (so a team
+ // never spans shires — L2 SCP is shire-local) and at nk_tiles (so each
+ // team member gets at least one KV tile).
+ const int64_t nk_tiles = (nk + TILE_KV - 1) / TILE_KV;
+ const int64_t total_minions = 2 * NUM_COMPUTE_SHIRES * MINIONS_PER_SHIRE;
+ int64_t k_splits = 1;
+ if (total_rows < total_minions) {
+ int64_t target = total_minions / total_rows;
+ int64_t ks = 1;
+ while (ks * 2 <= target && ks * 2 <= MINIONS_PER_SHIRE && ks * 2 <= nk_tiles) {
+ ks *= 2;
+ }
+ k_splits = ks;
+ }
+
+ const int64_t tiles_per_shire = MINIONS_PER_SHIRE / k_splits;
+ const int64_t k_split = (int64_t) local_minion % k_splits;
+ const int64_t local_tile_idx = (int64_t) local_minion / k_splits;
+ const int64_t tiles_stride = (int64_t) NUM_COMPUTE_SHIRES * tiles_per_shire;
+
+ // KV slab for this k_split. With k_splits=1 this is the full range.
+ const int64_t tiles_per_split_rounded = (nk_tiles + k_splits - 1) / k_splits;
+ const int64_t tile_start = k_split * tiles_per_split_rounded;
+ int64_t tile_end = tile_start + tiles_per_split_rounded;
+ if (tile_end > nk_tiles) {
+ tile_end = nk_tiles;
+ }
+ const int64_t kv_start = tile_start * TILE_KV;
+ int64_t kv_end = tile_end * TILE_KV;
+ if (kv_end > nk) {
+ kv_end = nk;
+ }
+
+ // L2 SCP pointers for this minion
+ uint64_t scp_base = local_minion * SCP_PER_MINION;
+ et_fp16_t * scp_kp[2] = {
+ (et_fp16_t *) et_shire_l2scp_local(scp_base + SCP_KP0_OFF),
+ (et_fp16_t *) et_shire_l2scp_local(scp_base + SCP_KP1_OFF),
+ };
+
+ // Hart 1 does K-panel packing
+ //
+ // When k_splits > 1, hart 1 must also participate in the two shire
+ // barriers that bracket the merge phase (one before and one after, so
+ // the reducer can read peer partials safely and the writers know when
+ // their acc/stats slab is free to reuse). Hart 1 has no useful work
+ // between those barriers.
+ //
+ // All teams in a shire must iterate the same number of times so the
+ // per-iter shire barriers stay balanced. Teams whose assigned row is
+ // past total_rows still call the barriers but skip the packing work.
+ et_barrier(ET_BARRIER_SHIRE);
+ // et_barrier(ET_BARRIER_GLOBAL);
+ if (is_hart1) {
+ uint32_t chunk_id = 0;
+ const int64_t row_base = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES;
+
+ int64_t max_iters;
+ if (k_splits > 1) {
+ max_iters = (total_rows + tiles_stride - 1) / tiles_stride;
+ } else {
+ max_iters = (row_base >= total_rows) ? 0 : ((total_rows - row_base - 1) / tiles_stride + 1);
+ }
+
+ for (int64_t iter = 0; iter < max_iters; iter++) {
+ const int64_t row = row_base + iter * tiles_stride;
+ const int has_work = (row < total_rows);
+
+ if (has_work) {
+ const int64_t iq3 = row / (nhq * nq);
+ const int64_t rem = row % (nhq * nq);
+ const int64_t iq2 = rem / nq;
+ const int64_t ik2 = iq2 / gqa_ratio;
+
+ const char * k_head = k_data + ik2 * k->nb[2] + iq3 * k->nb[3];
+
+ for (int64_t kv_base = kv_start; kv_base < kv_end; kv_base += TILE_KV) {
+ const int64_t kv_count = (kv_base + TILE_KV <= nk) ? TILE_KV : (nk - kv_base);
+
+ for (int64_t dk_chunk = 0; dk_chunk < dk; dk_chunk += TILE_K) {
+ int buf = chunk_id & 1;
+
+ // Back-pressure: before overwriting buf[buf] on chunk N
+ // (which will displace chunk N-2), wait for hart 0 to
+ // post that it's done with chunk N-2. Gates both
+ // directions of double-buffering.
+ //
+ // NOTE: we use et_sem_* (FCC 0 only) rather than
+ // et_barrier(ET_BARRIER_MINION) here because the
+ // minion barrier for minion 0 shares FLB 0 with
+ // ET_BARRIER_SHIRE. Mixing them deadlocks. See
+ // feedback_flb_collision.
+ if (chunk_id >= 2) {
+ et_sem_wait(ET_BARRIER_MINION);
+ }
+
+ // Prefetch K data for this chunk
+ prefetch_kv_to_l2(k_head, kv_base, dk_chunk, kv_count, k->nb[1]);
+
+ pack_k_for_transpose16(scp_kp[buf], k_head, kv_base, dk_chunk, kv_count, k->nb[1]);
+
+ FENCE;
+ flush_to_l2(scp_kp[buf], 16, 64);
+ flush_to_l2((et_fp16_t *) ((char *) scp_kp[buf] + 1024), 16, 64);
+ WAIT_CACHEOPS;
+
+ // Signal: this buf is ready for hart 0 to consume.
+ et_sem_post(ET_BARRIER_MINION);
+
+ chunk_id++;
+ }
+ }
+ }
+
+ // Shire barriers for split-KV merge (hart 1 is a passive arrival).
+ if (k_splits > 1) {
+ et_barrier(ET_BARRIER_SHIRE); // A: team has written its partial
+ et_barrier(ET_BARRIER_SHIRE); // B: reducer has finished merge
+ }
+ }
+
+ // Self-drain phantom FCC 0 credits left by the wait-skip on the
+ // first 2 chunks. Hart 1 issued chunk_id posts but only
+ // (chunk_id - 2) waits (when chunk_id >= 2), so hart 1's FCC 0
+ // carries +min(chunk_id,2) credits from hart 0's matching posts
+ // that hart 1 never consumed.
+ uint32_t drain = (chunk_id < 2) ? chunk_id : 2;
+ for (uint32_t d = 0; d < drain; d++) {
+ et_sem_wait(ET_BARRIER_MINION);
+ }
+
+ // FENCE;
+ // et_barrier(ET_BARRIER_GLOBAL);
+ return 0;
+ }
+
+ // Hart 0: tensor engine compute
+#ifndef UBERKERNEL_SUPPRESS_SCP_SETUP
+ setup_cache_scp();
+#endif
+ CLEAR_TENSOR_ERROR;
+
+ // Q converted to F16 (one row at a time)
+ et_fp16_t q_f16[FA_DK_MAX] __attribute__((aligned(64)));
+
+ // Score buffer for QK^T output (16 scores per KV tile)
+ float scores[TILE_KV] __attribute__((aligned(64)));
+
+ // Small buffers for V accumulation
+ et_fp16_t w_f16_buf[32] __attribute__((aligned(64))); // 64 bytes
+ et_fp16_t vpanel_buf[8 * 32] __attribute__((aligned(64))); // 512 bytes
+
+ float * acc = (float *) et_shire_l2scp_local(scp_base + SCP_ACC_OFF);
+
+ uint32_t chunk_id = 0;
+
+ // Iter-based outer loop (matches hart 1). When k_splits > 1 all teams
+ // in a shire iterate the same number of times so the per-row shire
+ // barriers stay balanced; iterations with row >= total_rows skip the
+ // compute but still participate in the barriers.
+ const int64_t hart0_row_base = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES;
+ int64_t hart0_max_iters;
+ if (k_splits > 1) {
+ hart0_max_iters = (total_rows + tiles_stride - 1) / tiles_stride;
+ } else {
+ hart0_max_iters = (hart0_row_base >= total_rows) ? 0 : ((total_rows - hart0_row_base - 1) / tiles_stride + 1);
+ }
+
+ for (int64_t iter = 0; iter < hart0_max_iters; iter++) {
+ const int64_t row = hart0_row_base + iter * tiles_stride;
+ if (row >= total_rows) {
+ // No-work iteration: only participate in barriers (k_splits > 1).
+ if (k_splits > 1) {
+ et_barrier(ET_BARRIER_SHIRE); // A
+ et_barrier(ET_BARRIER_SHIRE); // B
+ }
+ continue;
+ }
+
+ const int64_t iq3 = row / (nhq * nq);
+ const int64_t rem = row % (nhq * nq);
+ const int64_t iq2 = rem / nq;
+ const int64_t iq1 = rem % nq;
+ const int64_t ik2 = iq2 / gqa_ratio;
+
+ // Read Q row (F32) and convert to F16
+ const float * pq = (const float *) (q_data + iq1 * q->nb[1] + iq2 * q->nb[2] + iq3 * q->nb[3]);
+ convert_q_row_f32_to_f16(q_f16, pq, dk);
+
+ // V base for this head + batch (K packing handled by hart 1)
+ const char * v_head = v_data + ik2 * v->nb[2] + iq3 * v->nb[3];
+
+ // Output pointer
+ float * out = (float *) (dst_data + iq2 * dst->nb[1] + iq1 * dst->nb[2] + iq3 * dst->nb[3]);
+
+ zero_acc_vec(acc, dv);
+ float M = ET_NEG_INF_F;
+ float S = 0.0f;
+ const char * mask_base = has_mask ? get_mask_row_base(mask, iq1, iq2, iq3) : (const char *) 0;
+
+ // Flush Q_f16 to L2 so tensor_load can see it
+ FENCE;
+ flush_to_l2(q_f16, (dk * 2 + 63) / 64, 64);
+ WAIT_CACHEOPS;
+
+ for (int64_t kv_base = kv_start; kv_base < kv_end; kv_base += TILE_KV) {
+ const int64_t kv_count = (kv_base + TILE_KV <= nk) ? TILE_KV : (nk - kv_base);
+
+ // Set tensor_mask for partial tiles
+ if (kv_count < TILE_KV) {
+ uint64_t tmask = (1ULL << kv_count) - 1;
+ __asm__ __volatile__("csrw 0x805, %0" : : "r"(tmask));
+ }
+
+ // ============================================================
+ // QK^T via TensorFMA16A32
+ // ============================================================
+
+ // Pipelined QK^T:
+ // - Q for the whole row is preloaded once into A_L1[0..n-1].
+ // Each FMA picks its chunk via scp_loc_a = chunk_idx.
+ // - K is double-buffered in L1: K_BUFS[0]=lines 16..31,
+ // K_BUFS[1]=lines 32..47.
+ // - In iteration i (1..N-1), the K[i] load runs concurrently
+ // with the FMA on chunk i-1: they touch disjoint L1 regions
+ // (FMA reads K_BUFS[(i-1)&1], load writes K_BUFS[i&1]; FMA
+ // reads A_L1[i-1], load doesn't touch A_L1).
+ //
+ // L1 footprint: max dk=512 → Q uses 16 lines (0..15), K uses 32
+ // lines (16..47). Within ET-SoC-1 L1 SCP (≥128 lines per minion).
+ const int64_t n_dk_chunks = dk / TILE_K;
+ const uint64_t K_BUFS[2] = {
+ (uint64_t) B_L1_START, // 16..31
+ (uint64_t) (B_L1_START + 16), // 32..47
+ };
+
+ // Preload entire Q row into A_L1[0..n_dk_chunks-1] (one tensor_load,
+ // one wait, regardless of dk).
+ tensor_load(false, false, A_L1_START, TENSOR_LOAD_PLAIN, 0, (uint64_t) q_f16, 0,
+ (uint64_t) (n_dk_chunks - 1), 64, 0);
+
+ // Prologue: wait hart 1's K[0], issue K[0] load, wait both loads.
+ {
+ int buf = chunk_id & 1;
+ et_sem_wait(ET_BARRIER_MINION);
+ tensor_load(false, false, K_BUFS[0], TENSOR_LOAD_TRANSPOSE16, 0, (uint64_t) scp_kp[buf], 0, 15, 64, 1);
+ tensor_wait(TENSOR_LOAD_WAIT_0); // Q row complete
+ tensor_wait(TENSOR_LOAD_WAIT_1); // K[0] complete
+ et_sem_post(ET_BARRIER_MINION);
+ chunk_id++;
+ }
+
+ // Main loop: in iter i, issue K[i] load and FMA chunk i-1 in
+ // parallel. The matrix engine is busy on FMA[i-1] while the
+ // load unit fetches K[i] from L2 SCP.
+ //
+ // Order of waits matters: wait K[i] load first, then sem_post
+ // immediately (frees scp_kp[buf] for hart 1 to refill chunk i+2),
+ // then wait FMA. Putting sem_post after FMA wait would stall
+ // hart 1 by a full FMA latency — defeating the producer pipeline.
+ for (int64_t i = 1; i < n_dk_chunks; i++) {
+ int buf = chunk_id & 1;
+ int k_slot_prev = (int) ((i - 1) & 1);
+ int k_slot = (int) (i & 1);
+
+ et_sem_wait(ET_BARRIER_MINION);
+ tensor_load(false, false, K_BUFS[k_slot], TENSOR_LOAD_TRANSPOSE16, 0, (uint64_t) scp_kp[buf], 0, 15, 64,
+ 1);
+
+ tensor_fma((kv_count < TILE_KV), 3, 0, 15, 0, false, false, false, false, K_BUFS[k_slot_prev],
+ (uint64_t) (i - 1), TENSOR_FMA_OP_FP16, (i == 1));
+
+ tensor_wait(TENSOR_LOAD_WAIT_1); // K[i] in L1
+ et_sem_post(ET_BARRIER_MINION); // release scp_kp[buf] EARLY
+ tensor_wait(TENSOR_FMA_WAIT); // then wait FMA[i-1]
+ chunk_id++;
+ }
+
+ // Epilogue: FMA on the last chunk (no overlapping load).
+ {
+ int k_slot_last = (int) ((n_dk_chunks - 1) & 1);
+ tensor_fma((kv_count < TILE_KV), 3, 0, 15, 0, false, false, false, false, K_BUFS[k_slot_last],
+ (uint64_t) (n_dk_chunks - 1), TENSOR_FMA_OP_FP16, (n_dk_chunks == 1));
+ tensor_wait(TENSOR_FMA_WAIT);
+ }
+
+ // Prefetch V rows for this tile.
+ // Only useful for the partial-tile path below
+ if (kv_count < TILE_KV) {
+ for (int64_t d = 0; d < dv; d += 32) {
+ prefetch_kv_to_l2(v_head, kv_base, d, kv_count, v->nb[1]);
+ }
+ }
+
+ // Extract QK^T scores from vector register file
+ __asm__ volatile("" ::: "f0", "f1");
+ {
+ unsigned long _ms;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f2, 0(%[p_scale]) \n\t"
+ "fmul.ps f0, f0, f2 \n\t"
+ "fmul.ps f1, f1, f2 \n\t"
+ "fsw.ps f0, 0(%[dst]) \n\t"
+ "fsw.ps f1, 32(%[dst]) \n\t"
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(_ms)
+ : [dst] "r"(scores), [p_scale] "r"(&scale)
+ : "f0", "f1", "f2", "memory");
+ }
+
+ // ============================================================
+ // Two-phase softmax + V accumulation
+ // ============================================================
+
+ float weights[TILE_KV] __attribute__((aligned(64)));
+ {
+ // A1: apply mask to scores, pad unused slots
+ for (int64_t j = 0; j < kv_count; ++j) {
+ if (has_mask) {
+ float mv = get_mask_val_from_base(mask, mask_base, kv_base + j);
+ if (mv == ET_NEG_INF_F || mv != mv) {
+ scores[j] = ET_NEG_INF_F;
+ } else {
+ scores[j] += mv;
+ }
+ }
+ }
+ for (int64_t j = kv_count; j < TILE_KV; ++j) {
+ scores[j] = ET_NEG_INF_F;
+ }
+
+ // A1b: SIMD horizontal max across all 16 scores
+ float tile_max;
+ {
+ unsigned long _ms;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f2, 0(%[sc]) \n\t"
+ "flw.ps f3, 32(%[sc]) \n\t"
+ "fmax.ps f2, f2, f3 \n\t"
+ "fswizz.ps f3, f2, 0xB1 \n\t"
+ "fmax.ps f2, f2, f3 \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fmax.ps f2, f2, f3 \n\t"
+ "fmvz.x.ps t0, f2, 4 \n\t"
+ "fbcx.ps f3, t0 \n\t"
+ "fmax.ps %[tm], f2, f3 \n\t"
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(_ms), [tm] "=f"(tile_max)
+ : [sc] "r"(scores)
+ : "f2", "f3", "t0", "memory");
+ }
+
+ if (tile_max > ET_NEG_INF_F) {
+ // A2: rescale accumulator if this tile has a new global max
+ if (tile_max > M) {
+ float rescale = et_exp2f((M - tile_max) * 1.4426950408889634f);
+ scale_acc_vec(acc, dv, rescale);
+ S *= rescale;
+ M = tile_max;
+ }
+
+ // A3: SIMD exp2 + horizontal sum
+ // Interleaved: f2/f3 chains alternate to hide ALU latency.
+ // fexp.ps has multi-cycle latency — the two independent
+ // exp2 calls naturally pipeline.
+ {
+ const float log2e = 1.4426950408889634f;
+ float S_tile;
+ unsigned long _ms;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f2, 0(%[sc]) \n\t"
+ "fbc.ps f4, 0(%[pM]) \n\t"
+ "flw.ps f3, 32(%[sc]) \n\t"
+ "fbc.ps f5, 0(%[pL]) \n\t"
+ "fsub.ps f2, f2, f4 \n\t"
+ "fsub.ps f3, f3, f4 \n\t"
+ "fmul.ps f2, f2, f5 \n\t"
+ "fmul.ps f3, f3, f5 \n\t"
+ "fexp.ps f2, f2 \n\t"
+ "fexp.ps f3, f3 \n\t"
+ "fsw.ps f2, 0(%[wt]) \n\t"
+ "fsw.ps f3, 32(%[wt]) \n\t"
+ "fadd.ps f2, f2, f3, rne \n\t"
+ "fswizz.ps f3, f2, 0xB1 \n\t"
+ "fadd.ps f2, f2, f3, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f2, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f2, 4 \n\t"
+ "fbcx.ps f3, t0 \n\t"
+ "fadd.ps %[st], f2, f3, rne \n\t"
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(_ms), [st] "=f"(S_tile)
+ : [pM] "r"(&M), [pL] "r"(&log2e), [sc] "r"(scores), [wt] "r"(weights)
+ : "f2", "f3", "f4", "f5", "t0", "memory");
+ S += S_tile;
+ }
+
+ // Phase B: weights @ V via TensorFMA16A32
+ {
+ // B1: convert weights F32 → F16
+ convert_q_row_f32_to_f16(w_f16_buf, weights, TILE_KV);
+
+ FENCE;
+ flush_to_l2(w_f16_buf, 1, 64);
+ WAIT_CACHEOPS;
+
+ // Issue weights load (wait_id=0) and the first V chunk
+ // load (wait_id=1) concurrently. Weights comes from
+ // L2 SCP (just flushed); V[0] comes from DRAM via
+ // INTERLEAVE16 — running them in parallel hides the
+ // shorter load behind the longer one. For partial
+ // tiles, V is software-packed below — we only kick
+ // off the early V load on the full-tile fast path.
+ tensor_load(false, false, A_L1_START, TENSOR_LOAD_PLAIN, 0, (uint64_t) w_f16_buf, 0, 0, 64, 0);
+
+ const int v_full_tile = (kv_count == TILE_KV);
+ const uintptr_t v_base = (uintptr_t) v_head + kv_base * v->nb[1];
+ const uint64_t nb1_v = (uint64_t) v->nb[1];
+ uint64_t b_cur = 8;
+
+ if (v_full_tile) {
+ tensor_load(false, false, b_cur, TENSOR_LOAD_INTERLEAVE16, 0, (uint64_t) v_base, 0, 7,
+ nb1_v, 1);
+ }
+
+ tensor_wait(TENSOR_LOAD_WAIT_0); // weights in A_L1
+ if (v_full_tile) {
+ tensor_wait(TENSOR_LOAD_WAIT_1); // V[0] in b_cur
+ }
+
+ // B2: process dv in chunks of 16
+ if (v_full_tile) {
+ for (int64_t dv_off = 0; dv_off < dv; dv_off += 16) {
+ const uint64_t b_nxt = b_cur ^ 24;
+
+ if (dv_off + 16 < dv) {
+ tensor_load(false, false, b_nxt, TENSOR_LOAD_INTERLEAVE16, 0,
+ (uint64_t) (v_base + (dv_off + 16) * 2), 0, 7, nb1_v, 1);
+ }
+
+ tensor_fma(false, 3, 0, 7, 0, false, false, false, false, b_cur, A_L1_START,
+ TENSOR_FMA_OP_FP16, true);
+ tensor_wait(TENSOR_FMA_WAIT);
+
+ __asm__ volatile("" ::: "f0", "f1");
+ {
+ unsigned long _ms;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f2, 0(%[pa]) \n\t"
+ "flw.ps f3, 32(%[pa]) \n\t"
+ "fadd.ps f0, f0, f2 \n\t"
+ "fadd.ps f1, f1, f3 \n\t"
+ "fsw.ps f0, 0(%[pa]) \n\t"
+ "fsw.ps f1, 32(%[pa]) \n\t"
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(_ms)
+ : [pa] "r"(acc + dv_off)
+ : "f0", "f1", "f2", "f3", "memory");
+ }
+
+ if (dv_off + 16 < dv) {
+ tensor_wait(TENSOR_LOAD_WAIT_1);
+ b_cur = b_nxt;
+ }
+ }
+ } else {
+ // Partial tile: software pack, no pipeline
+ for (int64_t dv_off = 0; dv_off < dv; dv_off += 16) {
+ pack_v_interleaved(vpanel_buf, v_head, kv_base, dv_off, kv_count, v->nb[1]);
+ FENCE;
+ flush_to_l2(vpanel_buf, 8, 64);
+ WAIT_CACHEOPS;
+ tensor_load(false, false, B_L1_START, TENSOR_LOAD_PLAIN, 0, (uint64_t) vpanel_buf, 0, 7,
+ 64, 0);
+ tensor_wait(TENSOR_LOAD_WAIT_0);
+
+ tensor_fma(false, 3, 0, 7, 0, false, false, false, false, B_L1_START, A_L1_START,
+ TENSOR_FMA_OP_FP16, true);
+ tensor_wait(TENSOR_FMA_WAIT);
+
+ __asm__ volatile("" ::: "f0", "f1");
+ {
+ unsigned long _ms;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f2, 0(%[pa]) \n\t"
+ "flw.ps f3, 32(%[pa]) \n\t"
+ "fadd.ps f0, f0, f2 \n\t"
+ "fadd.ps f1, f1, f3 \n\t"
+ "fsw.ps f0, 0(%[pa]) \n\t"
+ "fsw.ps f1, 32(%[pa]) \n\t"
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(_ms)
+ : [pa] "r"(acc + dv_off)
+ : "f0", "f1", "f2", "f3", "memory");
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ // Finalize row
+ //
+ // k_splits == 1: this minion computed the full row. Normalize in
+ // place and store to DRAM.
+ //
+ // k_splits > 1: this minion computed a KV slab. Publish the
+ // partial (M, S, acc) to L2 SCP, sync with the
+ // team, and let the k_split==0 member do the
+ // softmax combine and the final store. All tensor
+ // engine ops are complete before this block, so
+ // f0..f31 are free to use.
+ if (k_splits > 1) {
+ // Publish our partial.
+ volatile float * my_stats = (volatile float *) et_shire_l2scp_local(scp_base + SCP_STATS_OFF);
+ my_stats[0] = M;
+ my_stats[1] = S;
+ FENCE;
+ evict_range_to_l2(acc, (int64_t) dv * (int64_t) sizeof(float));
+ evict_to_l2((const void *) my_stats, 1, 64);
+ WAIT_CACHEOPS;
+
+ // A: team members have all written their partials.
+ et_barrier(ET_BARRIER_SHIRE);
+
+ if (k_split == 0) {
+ // Online softmax merge: fold peers 1..k_splits-1 into our
+ // own (M_running, S_running, acc). For each peer p:
+ // M_new = max(M_running, M_p)
+ // α_own = exp2((M_running - M_new) * log2e)
+ // α_p = exp2((M_p - M_new) * log2e)
+ // acc[d] = α_own * acc[d] + α_p * peer_acc[d]
+ // S_running = α_own * S_running + α_p * S_p
+ float M_running = M;
+ float S_running = S;
+ const float log2e = 1.4426950408889634f;
+
+ for (int64_t p = 1; p < k_splits; p++) {
+ uint64_t peer_scp = (local_tile_idx * k_splits + p) * SCP_PER_MINION;
+ volatile float * peer_stats = (volatile float *) et_shire_l2scp_local(peer_scp + SCP_STATS_OFF);
+ float * peer_acc = (float *) et_shire_l2scp_local(peer_scp + SCP_ACC_OFF);
+
+ // Drop stale L1D copies before reading peer's data.
+ evict_to_l2((const void *) peer_stats, 1, 64);
+ evict_range_to_l2(peer_acc, (int64_t) dv * (int64_t) sizeof(float));
+ WAIT_CACHEOPS;
+
+ const float M_p = peer_stats[0];
+ const float S_p = peer_stats[1];
+
+ const float M_new = (M_p > M_running) ? M_p : M_running;
+ const float alpha_own = (M_running == ET_NEG_INF_F) ? 0.0f : et_exp2f((M_running - M_new) * log2e);
+ const float alpha_p = (M_p == ET_NEG_INF_F) ? 0.0f : et_exp2f((M_p - M_new) * log2e);
+
+ merge_rescale_add_asm(acc, peer_acc, dv, alpha_own, alpha_p);
+
+ S_running = alpha_own * S_running + alpha_p * S_p;
+ M_running = M_new;
+ }
+
+ const float S_inv = (S_running == 0.0f) ? 0.0f : et_fdiv(1.0f, S_running);
+ normalize_store_vec(out, acc, dv, S_inv, use_fast_store);
+ }
+
+ // B: reducer is done, team may reuse its acc/stats slabs.
+ et_barrier(ET_BARRIER_SHIRE);
+ } else {
+ // k_splits == 1 fast path — this minion owns the full row.
+ const float S_inv = S == 0.0f ? 0.0f : et_fdiv(1.0f, S);
+ normalize_store_vec(out, acc, dv, S_inv, use_fast_store);
+ }
+ }
+
+ FENCE;
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// F32 Flash Attention for ET backend
+//
+// Supports:
+// - arbitrary dk/dv (up to 128)
+// - GQA (n_head_q can differ from n_head_kv)
+// - mask (F16 or F32, causal pattern)
+// - F16 or F32 K and V (with non-contiguous strides from KV cache permute)
+//
+// Limitations:
+// - Q and dst must be F32
+// - no sinks, ALiBi, logit softcap
+//
+// Parallelization strategy:
+// - flatten [query position, head, outer batch] into independent rows
+// - assign rows round-robin across ET threads
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdbool.h>
+#include <stdint.h>
+
+struct ggml_et_flash_attn_ext_params {
+ struct ggml_tensor src0; // Q tensor (F32)
+ struct ggml_tensor src1; // K tensor (F16 or F32)
+ struct ggml_tensor src2; // V tensor (F16 or F32)
+ struct ggml_tensor mask; // mask tensor (F16 or F32), zeroed when absent
+ struct ggml_tensor dst; // Output tensor (F32)
+ float scale; // Scale factor applied to QK
+ int32_t has_mask; // nonzero if mask is present
+};
+
+// Maximum head dimension supported (128 covers all common LLMs).
+#define FA_DV_MAX 128
+
+// Read element d from a row, handling F16 or F32 type.
+// row_base points to the start of the row (byte address).
+// nb0 is the stride per element (2 for F16, 4 for F32).
+static inline float read_kv_f32(const char * row_base, int64_t d, int64_t nb0, int type) {
+ if (type == GGML_TYPE_F32) {
+ return *(const float *) (row_base + d * nb0);
+ }
+ // F16
+ return fp16_to_fp32(*(const uint16_t *) (row_base + d * nb0));
+}
+
+// Dot product of F32 query vector with a K row (F16 or F32).
+static inline float dot_qk(const float * q, const char * k_row, int64_t dk, int64_t k_nb0, int k_type) {
+ float acc = 0.0f;
+ if (k_type == GGML_TYPE_F32) {
+ const float * kf = (const float *) k_row;
+ for (int64_t i = 0; i < dk; ++i) {
+ acc += q[i] * kf[i];
+ }
+ } else {
+ // F16 stride-aware read
+ for (int64_t i = 0; i < dk; ++i) {
+ acc += q[i] * fp16_to_fp32(*(const uint16_t *) (k_row + i * k_nb0));
+ }
+ }
+ return acc;
+}
+
+static inline float get_mask_val(const struct ggml_tensor * mask, int64_t iq1, int64_t ik1, int64_t iq2, int64_t iq3) {
+ // mask layout: [nk, nq, ne2, ne3] -> broadcast via modulo
+ const char * base = (const char *) mask->data + iq1 * mask->nb[1] + (iq2 % mask->ne[2]) * mask->nb[2] +
+ (iq3 % mask->ne[3]) * mask->nb[3];
+
+ if (mask->type == GGML_TYPE_F32) {
+ return *(const float *) (base + ik1 * mask->nb[0]);
+ }
+ // F16
+ return fp16_to_fp32(*(const uint16_t *) (base + ik1 * mask->nb[0]));
+}
+
+int entry_point(struct ggml_et_flash_attn_ext_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env || !params) {
+ return -1;
+ }
+
+ const int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ const int num_threads = get_num_threads(kernel_env->shire_mask);
+ if (thread_id < 0 || num_threads <= 0) {
+ return 0;
+ }
+
+ struct ggml_tensor * q = ¶ms->src0;
+ struct ggml_tensor * k = ¶ms->src1;
+ struct ggml_tensor * v = ¶ms->src2;
+ struct ggml_tensor * dst = ¶ms->dst;
+ const int32_t has_mask = params->has_mask;
+ struct ggml_tensor * mask = has_mask ? ¶ms->mask : (struct ggml_tensor *) 0;
+
+ const char * q_data = (const char *) q->data;
+ const char * k_data = (const char *) k->data;
+ const char * v_data = (const char *) v->data;
+ char * dst_data = (char *) dst->data;
+
+ const int k_type = k->type;
+ const int v_type = v->type;
+ const int64_t k_nb0 = k->nb[0];
+ const int64_t v_nb0 = v->nb[0];
+
+ const int64_t dk = q->ne[0]; // head dim for keys/queries
+ const int64_t nq = q->ne[1]; // number of query positions
+ const int64_t nhq = q->ne[2]; // number of query heads
+ const int64_t no = q->ne[3]; // outer batch
+
+ const int64_t nk = k->ne[1]; // number of key/value positions
+ const int64_t nhk = k->ne[2]; // number of kv heads
+ const int64_t dv = v->ne[0]; // head dim for values
+
+ if (dv > FA_DV_MAX) {
+ return -1;
+ }
+
+ // GQA: query heads per kv head
+ const int64_t gqa_ratio = nhq / nhk;
+
+ const int64_t total_rows = nq * nhq * no;
+ const float scale = params->scale;
+
+ // When dv is a multiple of 16 (64 bytes = cache line), output rows are
+ // cache-line aligned and we can use fast normal stores. Otherwise we must
+ // use atomic stores to avoid cache-line sharing corruption.
+ const int use_fast_store = (dv % 16 == 0);
+
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ const int64_t iq3 = row / (nhq * nq);
+ const int64_t rem = row % (nhq * nq);
+ const int64_t iq2 = rem / nq; // query head index
+ const int64_t iq1 = rem % nq; // query position
+
+ // Map query head -> kv head for GQA
+ const int64_t ik2 = iq2 / gqa_ratio;
+
+ // Q is always F32
+ const float * pq = (const float *) (q_data + iq1 * q->nb[1] + iq2 * q->nb[2] + iq3 * q->nb[3]);
+
+ // dst layout: [dv, nhq, nq, no]
+ float * out = (float *) (dst_data + iq2 * dst->nb[1] + iq1 * dst->nb[2] + iq3 * dst->nb[3]);
+
+ // Base byte offsets for K and V head+batch slice
+ const int64_t kv_base = ik2 * k->nb[2] + iq3 * k->nb[3];
+ const int64_t vv_base = ik2 * v->nb[2] + iq3 * v->nb[3];
+
+ float acc[FA_DV_MAX];
+ for (int64_t d = 0; d < dv; ++d) {
+ acc[d] = 0.0f;
+ }
+
+ float M = -3.402823466e+38f;
+ float S = 0.0f;
+
+ for (int64_t ik1 = 0; ik1 < nk; ++ik1) {
+ // If mask is present, check for -inf (skip masked positions)
+ float mask_val = 0.0f;
+ if (has_mask) {
+ mask_val = get_mask_val(mask, iq1, ik1, iq2, iq3);
+ // llama.cpp uses -inf for masked positions
+ if (mask_val == -3.402823466e+38f || mask_val != mask_val) {
+ continue;
+ }
+ }
+
+ const char * pk = k_data + ik1 * k->nb[1] + kv_base;
+ const char * pv = v_data + ik1 * v->nb[1] + vv_base;
+
+ float s = dot_qk(pq, pk, dk, k_nb0, k_type) * scale + mask_val;
+ const float Mold = M;
+
+ float ms = 1.0f;
+ float vs = 1.0f;
+ if (s > M) {
+ M = s;
+ ms = et_expf(Mold - M);
+ for (int64_t d = 0; d < dv; ++d) {
+ acc[d] *= ms;
+ }
+ } else {
+ vs = et_expf(s - M);
+ }
+
+ // Accumulate weighted V
+ if (v_type == GGML_TYPE_F32) {
+ const float * pvf = (const float *) pv;
+ for (int64_t d = 0; d < dv; ++d) {
+ acc[d] += pvf[d] * vs;
+ }
+ } else {
+ for (int64_t d = 0; d < dv; ++d) {
+ acc[d] += fp16_to_fp32(*(const uint16_t *) (pv + d * v_nb0)) * vs;
+ }
+ }
+
+ S = S * ms + vs;
+ }
+
+ const float S_inv = S == 0.0f ? 0.0f : et_fdiv(1.0f, S);
+ if (use_fast_store) {
+ for (int64_t d = 0; d < dv; ++d) {
+ out[d] = acc[d] * S_inv;
+ }
+ } else {
+ for (int64_t d = 0; d < dv; ++d) {
+ atomic_store_f32((volatile float *) &out[d], acc[d] * S_inv);
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Gated Delta Net F32 Kernel
+//
+// Implements the gated delta rule recurrence:
+// For each head h, timestep t:
+// 1. Gate decay: S *= exp(g) (scalar or per-element KDA)
+// 2. Delta update: delta[j] = (v[j] - dot(S_row_j, k)) * beta
+// 3. Outer product: S_row_j += k * delta[j]
+// 4. Attention: attn[j] = dot(S_row_j, q) * scale
+//
+// State is stored transposed: s_out[j*S_v + i] = S[i][j]
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_gated_delta_net_params {
+ struct ggml_tensor q; // [S_v, H_q, n_tokens, n_seqs_q]
+ struct ggml_tensor k; // [S_v, H_k, n_tokens, n_seqs_k]
+ struct ggml_tensor v; // [S_v, H, n_tokens, n_seqs]
+ struct ggml_tensor g; // [1 or S_v, H, n_tokens, n_seqs]
+ struct ggml_tensor beta; // [1, H, n_tokens, n_seqs]
+ struct ggml_tensor state_in; // [S_v*S_v*H, K, n_seqs]
+ struct ggml_tensor dst; // [S_v*H, n_tokens*n_seqs + S_v*n_seqs*K]
+ int32_t S_v; // head dimension
+ int32_t H; // number of value heads
+ int32_t H_q; // number of Q heads
+ int32_t H_k; // number of K heads
+ int32_t n_tokens; // total tokens
+ int32_t n_seqs; // number of sequences
+ int32_t n_seqs_q; // Q sequence count
+ int32_t n_seqs_k; // K sequence count
+ int32_t kda; // 1 if per-element gate, 0 if scalar
+ int32_t K; // snapshot slot count
+ float scale; // 1/sqrt(S_v)
+};
+
+static inline float hsum_f10(void) {
+ float result;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result)::"t0", "f1", "f2", "f3", "f4", "f5");
+ return result;
+}
+
+int entry_point(struct ggml_et_gated_delta_net_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ const struct ggml_tensor * q_tsr = ¶ms->q;
+ const struct ggml_tensor * k_tsr = ¶ms->k;
+ const struct ggml_tensor * v_tsr = ¶ms->v;
+ const struct ggml_tensor * g_tsr = ¶ms->g;
+ const struct ggml_tensor * beta_tsr = ¶ms->beta;
+ const struct ggml_tensor * state_tsr = ¶ms->state_in;
+ const struct ggml_tensor * dst_tsr = ¶ms->dst;
+
+ const float * q = (const float *) q_tsr->data;
+ const float * k = (const float *) k_tsr->data;
+ const float * v = (const float *) v_tsr->data;
+ const float * g = (const float *) g_tsr->data;
+ const float * beta = (const float *) beta_tsr->data;
+ const float * state_in = (const float *) state_tsr->data;
+ float * dst_data = (float *) dst_tsr->data;
+
+ const int32_t S_v = params->S_v;
+ const int32_t H = params->H;
+ const int32_t H_q = params->H_q;
+ const int32_t H_k = params->H_k;
+ const int32_t n_tokens = params->n_tokens;
+ const int32_t n_seqs = params->n_seqs;
+ const int32_t n_seqs_q = params->n_seqs_q;
+ const int32_t n_seqs_k = params->n_seqs_k;
+ const int32_t kda = params->kda;
+ const int32_t K = params->K;
+ const float scale = params->scale;
+
+ if (!q || !k || !v || !g || !beta || !state_in || !dst_data) {
+ return -1;
+ }
+
+ // Preserve the original contract for every tensor except q, k, and v, which may be
+ // row-contiguous with strided higher dimensions.
+ if (q_tsr->nb[0] != sizeof(float) || k_tsr->nb[0] != sizeof(float) || v_tsr->nb[0] != sizeof(float) ||
+ g_tsr->nb[0] != sizeof(float) || beta_tsr->nb[0] != sizeof(float) || state_tsr->nb[0] != sizeof(float) ||
+ dst_tsr->nb[0] != sizeof(float)) {
+ return -1;
+ }
+
+ const int32_t attn_elems = S_v * H * n_tokens * n_seqs;
+ float * attn_out_base = dst_data;
+ float * state_out_base = dst_data + attn_elems;
+
+ const int32_t state_plane_floats = S_v * S_v * H * n_seqs;
+
+ const int32_t G0 = kda ? S_v : 1;
+
+ const size_t q_nb1 = q_tsr->nb[1];
+ const size_t q_nb2 = q_tsr->nb[2];
+ const size_t q_nb3 = q_tsr->nb[3];
+ const size_t k_nb1 = k_tsr->nb[1];
+ const size_t k_nb2 = k_tsr->nb[2];
+ const size_t k_nb3 = k_tsr->nb[3];
+ const size_t v_nb1 = v_tsr->nb[1];
+ const size_t v_nb2 = v_tsr->nb[2];
+ const size_t v_nb3 = v_tsr->nb[3];
+ const int32_t g_stride_h = G0;
+ const int32_t g_stride_t = G0 * H;
+ const int32_t g_stride_s = G0 * H * n_tokens;
+ const int32_t b_stride_t = H;
+ const int32_t b_stride_s = H * n_tokens;
+
+ float exp_g_buf[128];
+
+ // FP and SIMD share the same register file. Scalar FP needs the default
+ // mask; 8-wide .ps blocks need m0=255. Save once, toggle at boundaries.
+ unsigned long default_mask;
+ __asm__ volatile("mova.x.m %[ms]\n" : [ms] "=r"(default_mask));
+
+ // Parallelize over (j_block, head, seq). J_BLK must satisfy two separate
+ // cache-line alignment constraints at once:
+ // (a) State: J_BLK consecutive rows of s_out (each S_v floats) span an
+ // integer number of cache lines. For S_v * sizeof(float) >= 64 this
+ // is trivially any J_BLK >= 1.
+ // (b) Attention output: each j writes exactly one float into
+ // attn_ptr[j], which is densely packed. If J_BLK * sizeof(float) is
+ // less than a cache line, distinct threads will share a line and
+ // race on scalar stores — ET's L1 isn't coherent so we lose writes.
+ //
+ // (b) dominates: J_BLK must be at least ET_CACHE_LINE_SIZE_BYTES / 4 so
+ // that each thread owns a whole cache line of attn_ptr. That's 16 on
+ // ET-SoC-1, and it's also a whole number of state rows for every
+ // S_v >= 1, so (a) is automatically satisfied.
+ const int32_t J_BLK = ET_CACHE_LINE_SIZE_BYTES / (int32_t) sizeof(float);
+ const int32_t n_j_blocks = (S_v + J_BLK - 1) / J_BLK;
+ const int32_t total_work = n_j_blocks * H * n_seqs;
+
+ for (int32_t ir = thread_id; ir < total_work; ir += num_threads) {
+ const int32_t jb = ir % n_j_blocks;
+ const int32_t head = (ir / n_j_blocks) % H;
+ const int32_t seq = ir / (n_j_blocks * H);
+
+ const int32_t j_start = jb * J_BLK;
+ const int32_t j_end = (j_start + J_BLK < S_v) ? j_start + J_BLK : S_v;
+
+ const int32_t h_q = head % H_q;
+ const int32_t h_k = head % H_k;
+ const int32_t seq_q = (n_seqs_q == n_seqs) ? seq : (seq * n_seqs_q / n_seqs);
+ const int32_t seq_k = (n_seqs_k == n_seqs) ? seq : (seq * n_seqs_k / n_seqs);
+
+ const int32_t head_state_off = (seq * H + head) * S_v * S_v;
+ // Live RMW buffer = first snapshot plane (slot 0).
+ float * s_out = state_out_base + head_state_off;
+ // Input state: seq `seq`, head `head`.
+ const float * s_in = state_in + head_state_off;
+
+ // Skip the explicit s_in -> s_out copy. At t=0 pass A/B read through
+ // src_state = s_in; pass B writes the first new row to s_out. From
+ // t=1 onward src_state flips to s_out (read-modify-write in place).
+ const float * src_state = s_in;
+
+ const int32_t attn_stride_t = S_v * H;
+ float * attn_ptr = attn_out_base + (seq * n_tokens * H + head) * S_v;
+
+ const float zero = 0.0f;
+
+ for (int32_t t = 0; t < n_tokens; t++) {
+ const float * q_t = (const float *) ((const char *) q + seq_q * q_nb3 + t * q_nb2 + h_q * q_nb1);
+ const float * k_t = (const float *) ((const char *) k + seq_k * k_nb3 + t * k_nb2 + h_k * k_nb1);
+ const float * v_t = (const float *) ((const char *) v + seq * v_nb3 + t * v_nb2 + head * v_nb1);
+ const float * g_t = g + seq * g_stride_s + t * g_stride_t + head * g_stride_h;
+ const float beta_val = beta[seq * b_stride_s + t * b_stride_t + head];
+
+ // Precompute per-element gate for the kda path; scalar decay
+ // otherwise. Decay is fused into per-j pass A/B below, not
+ // applied to state in a separate pre-pass.
+ float decay = 0.0f; // only used when !kda
+ if (kda) {
+ const float log2e = 1.4426950408889634f;
+ __asm__ volatile("mov.m.x m0, x0, 255\n" :::);
+ __asm__ volatile("fbc.ps f20, %[l2e]\n" : : [l2e] "m"(log2e) : "f20");
+ for (int32_t i = 0; i < S_v; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[g_vec]\n"
+ "fmul.ps f10, f10, f20, rne\n"
+ "fexp.ps f10, f10\n"
+ "fsw.ps f10, %[out]\n"
+ : [out] "=m"(*(float (*)[8]) & exp_g_buf[i])
+ : [g_vec] "m"(*(const float (*)[8]) & g_t[i])
+ : "f10");
+ }
+ __asm__ volatile("mova.m.x %[ms]\n" : : [ms] "r"(default_mask));
+ } else {
+ decay = et_expf(g_t[0]);
+ }
+
+ for (int32_t j = j_start; j < j_end; j++) {
+ const float * src_row = src_state + j * S_v;
+ float * dst_row = s_out + j * S_v;
+
+ __asm__ volatile("mov.m.x m0, x0, 255\n" :::);
+ if (kda) {
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+ for (int32_t i = 0; i < S_v; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[s_vec]\n"
+ "flw.ps f12, %[g_vec]\n"
+ "flw.ps f13, %[k_vec]\n"
+ "fmul.ps f11, f11, f12\n" // row_dec = row * g
+ "fmadd.ps f10, f11, f13, f10\n" // acc += row_dec * k
+ :
+ : [s_vec] "m"(*(const float (*)[8]) & src_row[i]),
+ [g_vec] "m"(*(const float (*)[8]) & exp_g_buf[i]),
+ [k_vec] "m"(*(const float (*)[8]) & k_t[i])
+ : "f10", "f11", "f12", "f13");
+ }
+ } else {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "fbc.ps f22, %[d]\n"
+ :
+ : [z] "m"(zero), [d] "m"(decay)
+ : "f10", "f22");
+ for (int32_t i = 0; i < S_v; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[s_vec]\n"
+ "flw.ps f13, %[k_vec]\n"
+ "fmul.ps f11, f11, f22\n" // row_dec = row * decay
+ "fmadd.ps f10, f11, f13, f10\n" // acc += row_dec * k
+ :
+ : [s_vec] "m"(*(const float (*)[8]) & src_row[i]), [k_vec] "m"(*(const float (*)[8]) &
+ k_t[i])
+ : "f10", "f11", "f13");
+ }
+ }
+
+ float dot_sk = hsum_f10();
+ __asm__ volatile("mova.m.x %[ms]\n" : : [ms] "r"(default_mask));
+
+ float delta_j = (v_t[j] - dot_sk) * beta_val;
+
+ // -------- Pass B: decay + outer product + attn --------
+ __asm__ volatile("mov.m.x m0, x0, 255\n" :::);
+ if (kda) {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "fbc.ps f21, %[dj]\n"
+ :
+ : [z] "m"(zero), [dj] "m"(delta_j)
+ : "f10", "f21");
+ for (int32_t i = 0; i < S_v; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[s_vec]\n"
+ "flw.ps f12, %[g_vec]\n"
+ "flw.ps f13, %[k_vec]\n"
+ "flw.ps f14, %[q_vec]\n"
+ "fmul.ps f11, f11, f12\n" // row_dec = row * g
+ "fmadd.ps f11, f13, f21, f11\n" // row_new = row_dec + k*delta_j
+ "fsw.ps f11, %[s_out]\n"
+ "fmadd.ps f10, f11, f14, f10\n" // attn_acc += row_new * q
+ : [s_out] "=m"(*(float (*)[8]) & dst_row[i])
+ : [s_vec] "m"(*(const float (*)[8]) & src_row[i]),
+ [g_vec] "m"(*(const float (*)[8]) & exp_g_buf[i]),
+ [k_vec] "m"(*(const float (*)[8]) & k_t[i]), [q_vec] "m"(*(const float (*)[8]) & q_t[i])
+ : "f10", "f11", "f12", "f13", "f14");
+ }
+ } else {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "fbc.ps f21, %[dj]\n"
+ "fbc.ps f22, %[d]\n"
+ :
+ : [z] "m"(zero), [dj] "m"(delta_j), [d] "m"(decay)
+ : "f10", "f21", "f22");
+ for (int32_t i = 0; i < S_v; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[s_vec]\n"
+ "flw.ps f13, %[k_vec]\n"
+ "flw.ps f14, %[q_vec]\n"
+ "fmul.ps f11, f11, f22\n" // row_dec = row * decay
+ "fmadd.ps f11, f13, f21, f11\n" // row_new = row_dec + k*delta_j
+ "fsw.ps f11, %[s_out]\n"
+ "fmadd.ps f10, f11, f14, f10\n" // attn_acc += row_new * q
+ : [s_out] "=m"(*(float (*)[8]) & dst_row[i])
+ : [s_vec] "m"(*(const float (*)[8]) & src_row[i]),
+ [k_vec] "m"(*(const float (*)[8]) & k_t[i]), [q_vec] "m"(*(const float (*)[8]) & q_t[i])
+ : "f10", "f11", "f13", "f14");
+ }
+ }
+
+ float attn_val = hsum_f10();
+ __asm__ volatile("mova.m.x %[ms]\n" : : [ms] "r"(default_mask));
+
+ attn_ptr[j] = attn_val * scale;
+ }
+
+ // n-way merge snapshot: live state lives in slot 0 (== s_out).
+ // Copies state to target snapshot slots [1, K-1] in reverse chronological order.
+ // target_slot == 0 is the live buffer itself => no copy.
+ // target_slot >= K (when n_tokens > K) => older slots are discarded.
+ if (K > 1) {
+ const int32_t target_slot = (n_tokens - 1) - t;
+ if (target_slot > 0 && target_slot < K) {
+ float * snap = state_out_base + target_slot * state_plane_floats + head_state_off;
+ for (int32_t j = j_start; j < j_end; j++) {
+ const float * src = s_out + j * S_v;
+ float * dst = snap + j * S_v;
+ for (int32_t i = 0; i < S_v; i++) {
+ dst[i] = src[i];
+ }
+ }
+ }
+ }
+
+ // After t=0, state lives in s_out; flip src_state so subsequent
+ // timesteps read-modify-write in place.
+ src_state = s_out;
+ attn_ptr += attn_stride_t;
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal GET_ROWS F32 Kernel
+// Extracts specific rows from a source tensor based on row indices
+//
+// Algorithm:
+// 1. Read row indices from src1 (int32 tensor)
+// 2. For each index, extract the corresponding row from src0
+// 3. Copy the row data to the output tensor dst
+// 4. Handle different input types: F32, Q8_0, Q4_0, and Q4_K (quantized)
+//
+// Operation: dst[i] = src0[indices[i]] for i = 0..num_indices
+//
+// Features supported:
+// - F32 input data (direct copy)
+// - Q4_0 quantized input data (dequantized to F32)
+// - Q8_0 quantized input data (dequantized to F32)
+// - Q4_K quantized input data (dequantized to F32)
+// - Int32 row indices
+// - Multi-dimensional tensor support
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+#define CACHE_LINE_SIZE_BYTES 64
+
+struct ggml_et_get_rows_params {
+ struct ggml_tensor src0; // Data tensor (F32, Q4_0, Q8_0, or Q4_K)
+ struct ggml_tensor src1; // Row indices tensor (I32)
+ struct ggml_tensor dst; // Output tensor (F32)
+};
+
+#define CACHE_LINE_SIZE_BYTES 64
+#define CACHE_ELEMENTS(elem_size) (CACHE_LINE_SIZE_BYTES / (elem_size))
+
+// Copy a row of F32 data from source to destination
+static void copy_f32_row(float * dst, const float * src, int64_t num_elements) {
+ // Simple memcpy for F32 data - no conversion needed
+ for (int64_t i = 0; i < num_elements; i++) {
+ dst[i] = src[i];
+ }
+}
+
+static void copy_f16_row(float * dst, const uint16_t * src, int64_t num_elements) {
+ for (int64_t i = 0; i < num_elements; i++) {
+ dst[i] = fp16_to_fp32(src[i]);
+ }
+}
+
+// Copy a row of F32 data from source to destination, aligned to cache line boundaries
+// using FP32 load/store instructions. They don't perform data conversion so is fine.
+// Requirement: n_bytes is a multiple of CACHE_LINE_SIZE (64 bytes)
+static void copy_row_cache_align(float * dst, const float * src, int64_t n_bytes) {
+ int num_f32_elem = n_bytes / sizeof(float);
+
+ // Unrolled to do an entire cache line at a time
+ __asm__ volatile(
+ "1: \n\t"
+ // --- Process 64 Bytes (1 Cache Line) ---
+ // Load 256 bits (32 bytes) into f0 and the other into f1
+ "flq2 f0, 0(%[src]) \n\t"
+ "flq2 f1, 32(%[src]) \n\t"
+
+ // Store 256 bits (32 bytes) from f0 and f1
+ "fsq2 f0, 0(%[dst]) \n\t"
+ "fsq2 f1, 32(%[dst]) \n\t"
+
+ // Increment Pointers by 64 bytes
+ "addi %[src], %[src], 64 \n\t"
+ "addi %[dst], %[dst], 64 \n\t"
+
+ // Decrement count by 16 elements
+ "addi %[n], %[n], -16 \n\t"
+
+ // Loop if at least 16 elements remain
+ "bge %[n], %[stride_count], 1b \n\t"
+
+ : [dst] "+r"(dst), [src] "+r"(src), [n] "+r"(num_f32_elem)
+ : [stride_count] "r"(16L)
+ : "f0", "f1", "memory");
+}
+
+// Copied from GGML: copy a row of Q4_0 data to F32 destination (with dequantization)
+static void copy_q4_0_row(float * dst, const block_q4_0 * src_blocks, int64_t num_elements) {
+ const int64_t num_blocks = (num_elements + QK4_0 - 1) / QK4_0;
+
+ for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+ const int64_t elements_in_block = (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK4_0) : QK4_0;
+
+ float temp_buffer[QK4_0];
+ dequantize_q4_0_block(&src_blocks[block_idx], temp_buffer);
+
+ for (int64_t i = 0; i < elements_in_block; i++) {
+ dst[block_idx * QK4_0 + i] = temp_buffer[i];
+ }
+ }
+}
+
+// Copy a row of Q8_0 data to F32 destination (with dequantization)
+static void copy_q8_0_row(float * dst, const block_q8_0 * src_blocks, int64_t num_elements) {
+ // Number of Q8_0 blocks needed for this row
+ const int64_t num_blocks = (num_elements + QK8_0 - 1) / QK8_0; // Round up to handle partial blocks
+
+ for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+ const int64_t elements_in_block =
+ (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK8_0) : QK8_0; // Handle last partial block
+
+ // Dequantize the block
+ float temp_buffer[QK8_0];
+ dequantize_q8_0_block(&src_blocks[block_idx], temp_buffer);
+
+ // Copy dequantized values to destination
+ for (int64_t i = 0; i < elements_in_block; i++) {
+ dst[block_idx * QK8_0 + i] = temp_buffer[i];
+ }
+ }
+}
+
+// Copy a row of Q4_K data to F32 destination (with dequantization)
+static void copy_q4_K_row(float * dst, const block_q4_K * src_blocks, int64_t num_elements) {
+ const int64_t num_blocks = (num_elements + QK_K - 1) / QK_K;
+
+ for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+ const int64_t elements_in_block = (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK_K) : QK_K;
+
+ float temp_buffer[QK_K];
+ dequantize_q4_K_block(&src_blocks[block_idx], temp_buffer);
+
+ for (int64_t i = 0; i < elements_in_block; i++) {
+ dst[block_idx * QK_K + i] = temp_buffer[i];
+ }
+ }
+}
+
+static void dequantize_q8_0_block_cache_aligned(const block_q8_0 * block, float * dst) {
+ const int8_t * qs_ptr = block->qs;
+
+ uint64_t temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+
+ const int32_t __attribute__((aligned(32))) vec_indices[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+ float scale = fp16_to_fp32(block->d);
+ __asm__ volatile(
+ "fbcx.ps f0, %0 \n\t" // Broadcast integer scale to all lanes
+ "flq2 f1, 0(%1) \n\t" // Load gether indicies
+ ::"r"(scale),
+ "r"(vec_indices)
+ : "f0", "f1");
+
+ for (int i = 0; i < 4; i++) {
+ __asm__ volatile(
+ "fgb.ps f2, f1(%0) \n\t" // Loads 8 bytes from (qs_ptr + indices) and sign-extends to 32-bit int.
+ "fcvt.ps.pw f2, f2, rne \n\t" // Convert Int32 to Float32
+ "fmul.ps f2, f2, f0 \n\t" // f2 = f2 * f0 (scale)
+ "fsq2 f2, 0(%1) \n\t" // Store 256 bits (8 floats) to dst.
+
+ ::"r"(qs_ptr),
+ "r"(dst)
+ : "f2", "memory");
+
+ // Advance pointers in C
+ qs_ptr += 8;
+ dst += 8;
+ }
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+}
+
+// Copy a row of Q4_0 data to F32 destination (with dequantization), cache-aligned
+static void copy_q4_0_row_cache_aligned(float * dst, const block_q4_0 * src_blocks, int64_t num_elements) {
+ const int64_t num_blocks = (num_elements + QK4_0 - 1) / QK4_0;
+
+ // Scatter byte offsets: even lanes -> dst[j], odd lanes -> dst[j + QK4_0/2]
+ // For 4 consecutive packed bytes producing [low0, high0, low1, high1, low2, high2, low3, high3]:
+ // low_i -> byte offset i*4 (positions 0,1,2,3 in first half)
+ // high_i -> byte offset (16+i)*4 (positions 16,17,18,19 in second half)
+ const int32_t __attribute__((aligned(32))) scatter_offsets[8] = { 0 * 4, 16 * 4, 1 * 4, 17 * 4,
+ 2 * 4, 18 * 4, 3 * 4, 19 * 4 };
+
+ // Gather indices: each byte loaded twice for low/high nibble extraction
+ const int32_t __attribute__((aligned(32))) gather_indices[8] = { 0, 0, 1, 1, 2, 2, 3, 3 };
+
+ uint64_t temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+
+ // Load constant vectors once — shared across all blocks and iterations
+ __asm__ volatile(
+ "flq2 f4, 0(%0) \n\t" // f4 = scatter offsets
+ "flq2 f1, 0(%1) \n\t" // f1 = gather indices {0,0,1,1,2,2,3,3}
+ ::"r"(scatter_offsets),
+ "r"(gather_indices)
+ : "f1", "f4");
+
+ for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+ const block_q4_0 * block = &src_blocks[block_idx];
+ const uint8_t * qs = block->qs;
+ float * block_dst = dst + block_idx * QK4_0;
+
+ float scale = fp16_to_fp32(block->d);
+ float bias = -8.0f * scale;
+
+ // Per-block: broadcast scale and bias
+ __asm__ volatile(
+ "fbcx.ps f0, %0 \n\t" // f0 = broadcast(scale)
+ "fbcx.ps f3, %1 \n\t" // f3 = broadcast(-8 * scale)
+ ::"r"(scale),
+ "r"(bias)
+ : "f0", "f3");
+
+ // 4 iterations x 4 packed bytes = 16 bytes = full block -> 32 floats
+ for (int i = 0; i < 4; i++) {
+ __asm__ volatile(
+ "fgb.ps f2, f1(%0) \n\t" // Gather: [b0,b0,b1,b1,b2,b2,b3,b3]
+ "mov.m.x m0, x0, 0xAA \n\t" // Odd lanes only (fills gather latency)
+ "fsrli.pi f2, f2, 4 \n\t" // Odd lanes: byte >> 4 (high nibble)
+ "mov.m.x m0, x0, 0xFF \n\t" // Restore full mask
+ "fslli.pi f2, f2, 28 \n\t" // Isolate low 4 bits: shift left 28
+ "fsrli.pi f2, f2, 28 \n\t" // then right 28 -> nibble in [3:0]
+ "fcvt.ps.pw f2, f2, rne \n\t" // Int32 -> Float32
+ "fmul.ps f2, f2, f0 \n\t" // * scale
+ "fadd.ps f2, f2, f3 \n\t" // + bias -> (nibble - 8) * scale
+ "fscw.ps f2, f4(%1) \n\t" // Scatter to GGML positions
+
+ ::"r"(qs),
+ "r"(block_dst)
+ : "f2", "memory");
+
+ qs += 4; // 4 packed bytes consumed
+ block_dst += 4; // Advance base by 4 float positions
+ }
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask)); // Restore mask
+}
+
+// Copy a row of Q8_0 data to F32 destination (with dequantization)
+static void copy_q8_0_row_cache_aligned(float * dst, const block_q8_0 * src_blocks, int64_t num_elements) {
+ // Number of Q8_0 blocks needed for this row
+ const int64_t num_blocks = (num_elements + QK8_0 - 1) / QK8_0; // Round up to handle partial blocks
+
+ for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+ const int64_t elements_in_block =
+ (block_idx == num_blocks - 1) ? (num_elements - block_idx * QK8_0) : QK8_0; // Handle last partial block
+
+ // Dequantize the block
+ float temp_buffer[QK8_0];
+ dequantize_q8_0_block_cache_aligned(&src_blocks[block_idx], temp_buffer);
+
+ // Copy dequantized values to destination
+ for (int64_t i = 0; i < elements_in_block; i++) {
+ dst[block_idx * QK8_0 + i] = temp_buffer[i];
+ }
+ }
+}
+
+// Vectorized dequantization of a Q4_K super-block (256 elements) to F32
+// Processes 8 groups of 32 elements, using ET SIMD for the inner loops.
+// Output is sequential (no scatter needed unlike Q4_0).
+static void copy_q4_K_row_cache_aligned(float * dst, const block_q4_K * src_blocks, int64_t num_elements) {
+ const int64_t num_blocks = (num_elements + QK_K - 1) / QK_K;
+
+ // Gather indices for sequential byte access: {0,1,2,3,4,5,6,7}
+ const int32_t __attribute__((aligned(32))) gather_indices[8] = { 0, 1, 2, 3, 4, 5, 6, 7 };
+
+ uint64_t temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+
+ // Load gather indices once — shared across all blocks
+ __asm__ volatile("flq2 f1, 0(%0) \n\t" // f1 = gather indices {0,1,2,3,4,5,6,7}
+ ::"r"(gather_indices)
+ : "f1");
+
+ for (int64_t block_idx = 0; block_idx < num_blocks; block_idx++) {
+ const block_q4_K * block = &src_blocks[block_idx];
+ const uint8_t * qs = block->qs;
+ float * block_dst = dst + block_idx * QK_K;
+
+ const float d = fp16_to_fp32(block->d);
+ const float min = fp16_to_fp32(block->dmin);
+
+ int is = 0;
+ for (int j = 0; j < QK_K; j += 64) {
+ // Extract per-group scales and mins (scalar — only 8 pairs per super-block)
+ uint8_t sc, m;
+ get_scale_min_k4(is + 0, block->scales, &sc, &m);
+ const float d1 = d * sc;
+ const float neg_m1 = -(min * m);
+ get_scale_min_k4(is + 1, block->scales, &sc, &m);
+ const float d2 = d * sc;
+ const float neg_m2 = -(min * m);
+
+ // Low nibbles: 32 elements using d1, neg_m1
+ __asm__ volatile(
+ "fbcx.ps f0, %0 \n\t" // f0 = broadcast(d1)
+ "fbcx.ps f3, %1 \n\t" // f3 = broadcast(-m1)
+ ::"r"(d1),
+ "r"(neg_m1)
+ : "f0", "f3");
+
+ const uint8_t * qs_lo = qs;
+ float * dst_lo = block_dst + j;
+ for (int k = 0; k < 4; k++) {
+ __asm__ volatile(
+ "fgb.ps f2, f1(%0) \n\t" // Gather 8 bytes, sign-extend to int32
+ "fandi.pi f2, f2, 0xF \n\t" // Mask low nibble (imm10=15)
+ "fcvt.ps.pw f2, f2, rne \n\t" // Int32 -> Float32
+ "fmadd.ps f2, f2, f0, f3\n\t" // d1 * nibble + (-m1)
+ "fsq2 f2, 0(%1) \n\t" // Store 8 floats
+ ::"r"(qs_lo),
+ "r"(dst_lo)
+ : "f2", "memory");
+ qs_lo += 8;
+ dst_lo += 8;
+ }
+
+ // High nibbles: 32 elements using d2, neg_m2
+ __asm__ volatile(
+ "fbcx.ps f0, %0 \n\t" // f0 = broadcast(d2)
+ "fbcx.ps f3, %1 \n\t" // f3 = broadcast(-m2)
+ ::"r"(d2),
+ "r"(neg_m2)
+ : "f0", "f3");
+
+ const uint8_t * qs_hi = qs;
+ float * dst_hi = block_dst + j + 32;
+ for (int k = 0; k < 4; k++) {
+ __asm__ volatile(
+ "fgb.ps f2, f1(%0) \n\t" // Gather 8 bytes, sign-extend to int32
+ "fsrli.pi f2, f2, 4 \n\t" // Shift right 4: high nibble
+ "fandi.pi f2, f2, 0xF \n\t" // Mask to 4 bits (clean any sign-ext artifacts)
+ "fcvt.ps.pw f2, f2, rne \n\t" // Int32 -> Float32
+ "fmadd.ps f2, f2, f0, f3\n\t" // d2 * nibble + (-m2)
+ "fsq2 f2, 0(%1) \n\t" // Store 8 floats
+ ::"r"(qs_hi),
+ "r"(dst_hi)
+ : "f2", "memory");
+ qs_hi += 8;
+ dst_hi += 8;
+ }
+
+ qs += 32; // Advance to next 32 packed bytes
+ is += 2;
+ }
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask)); // Restore mask
+}
+
+// Determine the number of F32 elements per work unit for a given source type.
+// For F32: 1 cacheline (16 elements)
+// For quantized types: 1 quant block
+static int64_t get_elements_per_work_unit(int type) {
+ const int64_t elements_per_cacheline = CACHE_LINE_SIZE_BYTES / sizeof(float); // 16
+ switch (type) {
+ case GGML_TYPE_Q8_0:
+ return QK8_0; // 32 elements = 2 cachelines
+ case GGML_TYPE_Q4_0:
+ return QK4_0; // 32 elements = 2 cachelines
+ case GGML_TYPE_Q4_K:
+ return QK_K; // 256 elements = 16 cachelines
+ default:
+ return elements_per_cacheline; // 16 elements = 1 cacheline
+ }
+}
+
+static int get_row_f32_mc_cacheline_aligned(struct ggml_et_get_rows_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ struct ggml_tensor * src0 = ¶ms->src0; // Data tensor
+ struct ggml_tensor * src1 = ¶ms->src1; // Row indices tensor (I32)
+ struct ggml_tensor * dst = ¶ms->dst; // Output tensor (F32)
+
+ const int64_t ne00 = src0->ne[0]; // Source columns (row width)
+ const int64_t ne01 = src0->ne[1]; // Source rows (total available rows)
+ const int64_t ne02 = src0->ne[2]; // Source batch dimension
+ const int64_t ne03 = src0->ne[3]; // Source outer batch dimension
+
+ const int64_t ne10 = src1->ne[0]; // Number of indices in dimension 0
+ const int64_t ne11 = src1->ne[1]; // Number of indices in dimension 1
+ const int64_t ne12 = src1->ne[2]; // Batch dimension for indices
+ const int64_t ne13 = src1->ne[3]; // Outer batch dimension for indices
+
+ const int64_t total_rows_to_extract = ne10 * ne11 * ne12 * ne13;
+
+ // Determine work unit size based on source type
+ const int64_t elements_per_wu = get_elements_per_work_unit(src0->type);
+ const int64_t wus_per_row = ne00 / elements_per_wu;
+ const int64_t total_wus = total_rows_to_extract * wus_per_row;
+
+ // Distribute work units across threads (contiguous ranges)
+ const int64_t wus_per_thread = (total_wus + num_threads - 1) / num_threads;
+ const int64_t wu_start = thread_id * wus_per_thread;
+ int64_t wu_end = wu_start + wus_per_thread;
+ if (wu_end > total_wus) {
+ wu_end = total_wus;
+ }
+
+ void * src0_data = src0->data;
+ int32_t * src1_data = (int32_t *) src1->data;
+ float * dst_data = (float *) dst->data;
+
+ int64_t wu = wu_start;
+ while (wu < wu_end) {
+ // Determine which row this work unit belongs to and offset within row
+ const int64_t row_idx = wu / wus_per_row;
+ const int64_t wu_in_row = wu % wus_per_row;
+
+ // How many work units to process in this row (batch contiguous WUs in same row)
+ int64_t wus_remaining_in_row = wus_per_row - wu_in_row;
+ int64_t wus_to_process = wu_end - wu;
+ if (wus_remaining_in_row < wus_to_process) {
+ wus_to_process = wus_remaining_in_row;
+ }
+
+ // Calculate multi-dimensional index for this row
+ const int64_t i = row_idx;
+ const int64_t i13_idx = i / (ne12 * ne11 * ne10);
+ const int64_t i12_idx = (i - i13_idx * ne12 * ne11 * ne10) / (ne11 * ne10);
+ const int64_t i11_idx = (i - i13_idx * ne12 * ne11 * ne10 - i12_idx * ne11 * ne10) / ne10;
+ const int64_t i10_idx = i - i13_idx * ne12 * ne11 * ne10 - i12_idx * ne11 * ne10 - i11_idx * ne10;
+
+ // Get the row index from src1
+ const int64_t index_offset = i13_idx * ne12 * ne11 * ne10 + i12_idx * ne11 * ne10 + i11_idx * ne10 + i10_idx;
+ const int32_t row_index = src1_data[index_offset];
+
+ if (row_index < 0 || row_index >= ne01) {
+ return -1; // Index out of bounds
+ }
+
+ const int64_t batch_offset =
+ i11_idx * ne01 * ne00 + i12_idx * ne02 * ne01 * ne00 + i13_idx * ne03 * ne02 * ne01 * ne00;
+
+ const int64_t elem_offset_in_row = wu_in_row * elements_per_wu;
+ const int64_t num_elements = wus_to_process * elements_per_wu;
+
+ float * dst_row = dst_data + row_idx * ne00 + elem_offset_in_row;
+
+ if (src0->type == GGML_TYPE_F32) {
+ // F32 source: direct copy of cacheline-aligned chunk
+ const float * src_row = (const float *) src0_data + row_index * ne00 + batch_offset + elem_offset_in_row;
+ copy_row_cache_align(dst_row, src_row, num_elements * sizeof(float));
+ } else if (src0->type == GGML_TYPE_F16) {
+ // F16 source: scalar conversion over a destination-aligned write chunk.
+ const uint16_t * src_row =
+ (const uint16_t *) src0_data + row_index * ne00 + batch_offset + elem_offset_in_row;
+ copy_f16_row(dst_row, src_row, num_elements);
+ } else if (src0->type == GGML_TYPE_Q8_0) {
+ // Q8_0 source: dequantize work-unit-aligned blocks
+ const int64_t blocks_per_row = (ne00 + QK8_0 - 1) / QK8_0;
+ const int64_t src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+ const int64_t block_start = elem_offset_in_row / QK8_0;
+ const block_q8_0 * src_blocks = (const block_q8_0 *) src0_data + src_block_offset + block_start;
+ copy_q8_0_row_cache_aligned(dst_row, src_blocks, num_elements);
+ } else if (src0->type == GGML_TYPE_Q4_0) {
+ // Q4_0 source: dequantize work-unit-aligned blocks
+ const int64_t blocks_per_row = (ne00 + QK4_0 - 1) / QK4_0;
+ const int64_t src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+ const int64_t block_start = elem_offset_in_row / QK4_0;
+ const block_q4_0 * src_blocks = (const block_q4_0 *) src0_data + src_block_offset + block_start;
+ copy_q4_0_row_cache_aligned(dst_row, src_blocks, num_elements);
+ } else if (src0->type == GGML_TYPE_Q4_K) {
+ // Q4_K source: dequantize work-unit-aligned blocks
+ const int64_t blocks_per_row = (ne00 + QK_K - 1) / QK_K;
+ const int64_t src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+ const int64_t block_start = elem_offset_in_row / QK_K;
+ const block_q4_K * src_blocks = (const block_q4_K *) src0_data + src_block_offset + block_start;
+ copy_q4_K_row_cache_aligned(dst_row, src_blocks, num_elements);
+ }
+
+ wu += wus_to_process;
+ }
+
+ return 0;
+}
+
+static inline size_t tensor_bytes(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+int entry_point(struct ggml_et_get_rows_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+ if (!kernel_env) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0; // Data tensor (F32, Q4_0, Q8_0, or Q4_K)
+ struct ggml_tensor * src1 = ¶ms->src1; // Row indices tensor (I32)
+ struct ggml_tensor * dst = ¶ms->dst; // Output tensor (F32)
+
+ // Fast path - we know how to deal with them multi-core
+ if ((src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_Q8_0 ||
+ src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K) &&
+ src1->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_F32 && dst->ne[0] % CACHE_ELEMENTS(sizeof(float)) == 0) {
+ return get_row_f32_mc_cacheline_aligned(params, env);
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (thread_id != 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ if (dst->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_I32) {
+ return -1; // Invalid output or index type
+ }
+
+ if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_Q8_0 &&
+ src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K) {
+ return -1; // Unsupported input type
+ }
+
+ void * src0_data = src0->data;
+ int32_t * src1_data = (int32_t *) src1->data;
+ float * dst_data = (float *) dst->data;
+#ifdef ET_UBERKERNEL
+ evict_region_past_l2(src0_data, tensor_bytes(src0));
+ evict_region_past_l2(src1_data, tensor_bytes(src1));
+ evict_region_past_l2(dst_data, tensor_bytes(dst));
+#endif
+
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ const int64_t ne00 = src0->ne[0]; // Source columns (row width)
+ const int64_t ne01 = src0->ne[1]; // Source rows (total available rows)
+ const int64_t ne02 = src0->ne[2]; // Source batch dimension
+ const int64_t ne03 = src0->ne[3]; // Source outer batch dimension
+
+ const int64_t ne10 = src1->ne[0]; // Number of indices in dimension 0
+ const int64_t ne11 = src1->ne[1]; // Number of indices in dimension 1
+ const int64_t ne12 = src1->ne[2]; // Batch dimension for indices
+ const int64_t ne13 = src1->ne[3]; // Outer batch dimension for indices
+
+ const int64_t total_rows_to_extract = ne10 * ne11 * ne12 * ne13;
+#ifdef ET_UBERKERNEL
+ et_barrier(ET_BARRIER_GLOBAL);
+#endif
+ // Naive single-threaded implementation - process all rows sequentially
+ // XXX: Do we really need a single-threaded implementation?
+ for (int64_t i = 0; i < total_rows_to_extract; i++) {
+ // Calculate multi-dimensional index for the current output position
+ const int64_t i13_idx = i / (ne12 * ne11 * ne10);
+ const int64_t i12_idx = (i - i13_idx * ne12 * ne11 * ne10) / (ne11 * ne10);
+ const int64_t i11_idx = (i - i13_idx * ne12 * ne11 * ne10 - i12_idx * ne11 * ne10) / ne10;
+ const int64_t i10_idx = i - i13_idx * ne12 * ne11 * ne10 - i12_idx * ne11 * ne10 - i11_idx * ne10;
+
+ // Get the row index from src1
+ const int64_t index_offset = i13_idx * ne12 * ne11 * ne10 + i12_idx * ne11 * ne10 + i11_idx * ne10 + i10_idx;
+ const int32_t row_index = src1_data[index_offset];
+
+ if (row_index < 0 || row_index >= ne01) {
+ return -1; // Index out of bounds
+ }
+
+ const int64_t batch_offset =
+ i11_idx * ne01 * ne00 + i12_idx * ne02 * ne01 * ne00 + i13_idx * ne03 * ne02 * ne01 * ne00;
+
+ const int64_t dst_offset = i;
+
+ if (src0->type == GGML_TYPE_F32) {
+ // F32 source: direct copy
+ const float * src_row = (const float *) src0_data + row_index * ne00 + batch_offset;
+ float * dst_row = dst_data + dst_offset * ne00;
+ copy_f32_row(dst_row, src_row, ne00);
+ } else if (src0->type == GGML_TYPE_F16) {
+ // F16 source: scalar conversion
+ const uint16_t * src_row = (const uint16_t *) src0_data + row_index * ne00 + batch_offset;
+ float * dst_row = dst_data + dst_offset * ne00;
+ copy_f16_row(dst_row, src_row, ne00);
+ } else if (src0->type == GGML_TYPE_Q8_0) {
+ // Q8_0 source: dequantize while copying
+ const int64_t blocks_per_row = (ne00 + QK8_0 - 1) / QK8_0;
+ const int64_t src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+ const block_q8_0 * src_blocks = (const block_q8_0 *) src0_data + src_block_offset;
+ float * dst_row = dst_data + dst_offset * ne00;
+ copy_q8_0_row(dst_row, src_blocks, ne00);
+ } else if (src0->type == GGML_TYPE_Q4_0) {
+ // Q4_0 source: dequantize while copying
+ const int64_t blocks_per_row = (ne00 + QK4_0 - 1) / QK4_0;
+ const int64_t src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+ const block_q4_0 * src_blocks = (const block_q4_0 *) src0_data + src_block_offset;
+ float * dst_row = dst_data + dst_offset * ne00;
+ copy_q4_0_row(dst_row, src_blocks, ne00);
+ } else if (src0->type == GGML_TYPE_Q4_K) {
+ // Q4_K source: dequantize while copying
+ const int64_t blocks_per_row = (ne00 + QK_K - 1) / QK_K;
+ const int64_t src_block_offset = (row_index * blocks_per_row) + (batch_offset / ne00) * blocks_per_row;
+ const block_q4_K * src_blocks = (const block_q4_K *) src0_data + src_block_offset;
+ float * dst_row = dst_data + dst_offset * ne00;
+ copy_q4_K_row(dst_row, src_blocks, ne00);
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+// ET kernel entry-point parameter structs and tensor helpers.
+
+#ifndef GGML_TENSOR_H
+#define GGML_TENSOR_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "ggml.h"
+
+struct ggml_et_binary_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+// bias.data == NULL -> unfused MUL_MAT; otherwise dst = mat_mul(...) + bias.
+struct ggml_et_mm_q8_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+ struct ggml_tensor bias;
+};
+
+struct ggml_et_mul_mat_id_params {
+ struct ggml_tensor src0; // [K, M, n_expert]
+ struct ggml_tensor src1; // [K, n_expert_used, batch]
+ struct ggml_tensor src2; // [n_expert_used, batch] (I32 expert indices)
+ struct ggml_tensor dst; // [M, n_expert_used, batch, 1]
+};
+
+// ne[i] == 1 axes are skipped: their stride is unobservable.
+static inline int ggml_tensor_is_contiguous(const struct ggml_tensor * t, int type_size) {
+ int64_t expected = type_size;
+ for (int i = 0; i < GGML_MAX_DIMS; i++) {
+ if (t->ne[i] > 1 && (int64_t) t->nb[i] != expected) {
+ return 0;
+ }
+ expected *= t->ne[i];
+ }
+ return 1;
+}
+
+#endif // GGML_TENSOR_H
--- /dev/null
+//******************************************************************************
+// GLU F32 Kernel (SwiGLU specifically)
+// Gated Linear Unit: y[i] = silu(x[i]) * g[i] where silu(x) = x * sigmoid(x)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+// GLU kernel parameters structure (from ET backend ops)
+struct ggml_et_glu_params {
+ struct ggml_tensor src0; // F32 input tensor A (or combined tensor if src1 is null)
+ struct ggml_tensor src1; // F32 input tensor B (null for single tensor mode)
+ struct ggml_tensor dst; // F32 output tensor (n/2 columns)
+ int32_t glu_op_type; // GLU operation type (REGLU=0, GEGLU=1, SWIGLU=2, etc.)
+ int32_t swapped; // Whether gate and value are swapped
+ float alpha; // SWIGLU_OAI: sigmoid scaling factor
+ float limit; // SWIGLU_OAI: clamp limit
+};
+
+// SiLU activation function: silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+static inline float silu_f32(float x) {
+ // For numerical stability, use the mathematically equivalent form:
+ // silu(x) = x / (1 + exp(-x)) = x * sigmoid(x)
+ // For large negative x, exp(-x) -> inf, so silu(x) -> 0
+ // For large positive x, exp(-x) -> 0, so silu(x) -> x
+
+ if (x > 20.0f) {
+ // For x > 20, exp(-x) is negligible, silu(x) ~ x
+ return x;
+ } else if (x < -20.0f) {
+ // For x < -20, silu(x) ~ 0
+ return 0.0f;
+ } else {
+ // Use standard formula: silu(x) = x / (1 + exp(-x))
+ // Optimized using ET hardware division
+ float exp_neg_x = et_expf(-x);
+ float denominator = 1.0f + exp_neg_x;
+ return et_fdiv(x, denominator);
+ }
+}
+
+// Vectorized GeGLU block processing (8 elements = 1 cache line, 64B aligned)
+// gelu(x) = 0.5*x*(1 + tanh(z)) = x * (1 - 1/(exp(2z)+1))
+// where z = sqrt(2/pi) * x * (1 + 0.044715*x^2)
+// Reformulated to avoid inf*0 NaN: uses x * sigmoid(2z) identity
+static inline void block_geglu(float * dst_block, const float * x_block, const float * g_block, int elements) {
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ float one_const = 1.0f;
+ float coef_a_const = 0.044715f;
+ float sqrt2pi_const = 0.79788456080286535587989211986876f; // sqrt(2/pi)
+ float two_log2e_const = 2.8853900817779268f; // 2 * log2(e)
+
+ for (int32_t i = 0; i < elements; i += 8) {
+ __asm__ volatile(
+ // Load inputs
+ "flw.ps f10, %[x_vec]\n" // f10 = x
+ "flw.ps f11, %[g_vec]\n" // f11 = g
+
+ // Broadcast constants
+ "fbc.ps f20, %[one_ptr]\n" // f20 = 1.0
+ "fbc.ps f22, %[coef_ptr]\n" // f22 = 0.044715
+ "fbc.ps f23, %[sqrt2pi_ptr]\n" // f23 = sqrt(2/pi)
+ "fbc.ps f24, %[two_log2e_ptr]\n" // f24 = 2*log2(e)
+
+ // inner = 1 + 0.044715 * x^2
+ "fmul.ps f12, f10, f10\n" // f12 = x^2
+ "fmadd.ps f13, f22, f12, f20\n" // f13 = 1 + 0.044715*x^2
+
+ // z = sqrt(2/pi) * x * inner
+ "fmul.ps f14, f23, f10\n" // f14 = sqrt(2/pi) * x
+ "fmul.ps f14, f14, f13\n" // f14 = z
+
+ // exp(2z) via fexp.ps: feed z * 2*log2(e) since fexp computes 2^input
+ "fmul.ps f15, f14, f24\n" // f15 = 2z * log2(e)
+ "fexp.ps f15, f15\n" // f15 = exp(2z)
+
+ // gelu(x) = x * (1 - 1/(exp(2z)+1)) [NaN-safe: no inf*0]
+ // exp(2z)->inf: rcp(inf)=0, 1-0=1, gelu=x
+ // exp(2z)->0: rcp(1)=1, 1-1=0, gelu=0
+ "fadd.ps f16, f15, f20\n" // f16 = exp(2z) + 1
+ "frcp.ps f16, f16\n" // f16 = 1/(exp(2z) + 1)
+ "fsub.ps f16, f20, f16\n" // f16 = 1 - 1/(exp(2z)+1)
+ "fmul.ps f16, f10, f16\n" // f16 = gelu(x)
+
+ // Final result
+ "fmul.ps f18, f16, f11\n" // f18 = gelu(x) * g
+
+ "fsw.ps f18, %[dst_out]\n"
+
+ : [dst_out] "=m"(*(float (*)[8]) & dst_block[i])
+ : [x_vec] "m"(*(const float (*)[8]) & x_block[i]), [g_vec] "m"(*(const float (*)[8]) & g_block[i]),
+ [one_ptr] "m"(one_const), [coef_ptr] "m"(coef_a_const), [sqrt2pi_ptr] "m"(sqrt2pi_const),
+ [two_log2e_ptr] "m"(two_log2e_const)
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f18", "f20", "f22", "f23", "f24");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+}
+
+// Vectorized SwiGLU block processing (16 elements = 1 cache line)
+static inline void block_swiglu(float * dst_block, const float * x_block, const float * g_block, int elements) {
+ // Process 8 elements at a time using vector instructions
+ int32_t vec_end = (elements / 8) * 8;
+
+ // Set mask register to enable all 8 vector elements
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask)); // Save current mask
+ __asm__ volatile("mov.m.x m0, x0, 0xFF"); // Enable all 8 elements
+
+ // Constants for broadcasting
+ float zero_const = 0.0f;
+ float one_const = 1.0f;
+ float log2e_const = 1.4426950408889634f; // log2(e)
+
+ for (int32_t i = 0; i < vec_end; i += 8) {
+ // Vectorized SwiGLU: dst = silu(x) * g = (x / (1 + exp(-x))) * g
+ // Using ET hardware: exp, reciprocal, multiply operations
+ __asm__ volatile(
+ // Load input vectors
+ "flw.ps f10, %[x_vec]\n" // f10 = x[0..7]
+ "flw.ps f11, %[g_vec]\n" // f11 = g[0..7]
+
+ // Broadcast constants to vector registers
+ "fbc.ps f20, %[zero_ptr]\n" // f20 = broadcast(0.0f) to all 8 elements
+ "fbc.ps f21, %[one_ptr]\n" // f21 = broadcast(1.0f) to all 8 elements
+
+ // Compute -x (negate x by subtracting from zero)
+ "fsub.ps f12, f20, f10\n" // f12 = 0 - x = -x
+
+ // Convert to base-2 exponent: -x * log2(e) = -x * 1.44269504
+ // Load log2(e) constant
+ "fbc.ps f22, %[log2e_ptr]\n" // f22 = broadcast(1.44269504f)
+ "fmul.ps f13, f12, f22\n" // f13 = -x * log2(e)
+
+ // Compute 2^(-x * log2(e)) = exp(-x)
+ "fexp.ps f14, f13\n" // f14 = 2^(-x * log2(e)) = exp(-x)
+
+ // Compute 1 + exp(-x)
+ "fadd.ps f15, f14, f21\n" // f15 = exp(-x) + 1
+
+ // Compute 1 / (1 + exp(-x)) using reciprocal
+ "frcp.ps f16, f15\n" // f16 = 1 / (1 + exp(-x))
+
+ // Compute silu(x) = x * (1 / (1 + exp(-x)))
+ "fmul.ps f17, f10, f16\n" // f17 = x * (1 / (1 + exp(-x))) = silu(x)
+
+ // Compute final result: silu(x) * g
+ "fmul.ps f18, f17, f11\n" // f18 = silu(x) * g
+
+ // Store result
+ "fsw.ps f18, %[dst_out]\n" // Store 8 results to destination
+
+ : [dst_out] "=m"(*(float (*)[8]) & dst_block[i])
+ : [x_vec] "m"(*(const float (*)[8]) & x_block[i]), [g_vec] "m"(*(const float (*)[8]) & g_block[i]),
+ [zero_ptr] "m"(zero_const), // Memory reference to 0.0f for broadcasting
+ [one_ptr] "m"(one_const), // Memory reference to 1.0f for broadcasting
+ [log2e_ptr] "m"(log2e_const) // Memory reference to log2(e) for broadcasting
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17", "f18", "f20", "f21", "f22");
+ }
+
+ // Restore original mask
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ // Handle remaining elements (< 8) with scalar operations
+ for (int32_t i = vec_end; i < elements; i++) {
+ dst_block[i] = silu_f32(x_block[i]) * g_block[i];
+ }
+}
+
+// Vectorized ReGLU block: dst = max(0, x) * g
+static inline void block_reglu(float * dst_block, const float * x_block, const float * g_block, int elements) {
+ int32_t vec_end = (elements / 8) * 8;
+
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ float zero_const = 0.0f;
+
+ for (int32_t i = 0; i < vec_end; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x_vec]\n" // f10 = x
+ "flw.ps f11, %[g_vec]\n" // f11 = g
+ "fbc.ps f20, %[zero_ptr]\n" // f20 = 0.0
+ "fmax.ps f12, f10, f20\n" // f12 = max(x, 0)
+ "fmul.ps f13, f12, f11\n" // f13 = relu(x) * g
+ "fsw.ps f13, %[dst_out]\n"
+ : [dst_out] "=m"(*(float (*)[8]) & dst_block[i])
+ : [x_vec] "m"(*(const float (*)[8]) & x_block[i]), [g_vec] "m"(*(const float (*)[8]) & g_block[i]),
+ [zero_ptr] "m"(zero_const)
+ : "f10", "f11", "f12", "f13", "f20");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ for (int32_t i = vec_end; i < elements; i++) {
+ float xv = x_block[i];
+ dst_block[i] = (xv > 0.0f) ? xv * g_block[i] : 0.0f;
+ }
+}
+
+// Vectorized GeGLU-Quick block: dst = x * sigmoid(1.702 * x) * g
+// Using gelu_quick(x) = x / (1 + exp(-1.702*x))
+static inline void block_geglu_quick(float * dst_block, const float * x_block, const float * g_block, int elements) {
+ int32_t vec_end = (elements / 8) * 8;
+
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ float zero_const = 0.0f;
+ float one_const = 1.0f;
+ // -1.702 * log2(e), so that fexp.ps(x * neg_k_log2e) = exp(-1.702*x)
+ float neg_k_log2e_const = -1.702f * 1.4426950408889634f;
+
+ for (int32_t i = 0; i < vec_end; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x_vec]\n" // f10 = x
+ "flw.ps f11, %[g_vec]\n" // f11 = g
+ "fbc.ps f20, %[zero_ptr]\n" // f20 = 0
+ "fbc.ps f21, %[one_ptr]\n" // f21 = 1
+ "fbc.ps f22, %[k_ptr]\n" // f22 = -1.702*log2(e)
+ "fmul.ps f13, f10, f22\n" // f13 = -1.702*x*log2(e)
+ "fexp.ps f14, f13\n" // f14 = exp(-1.702*x)
+ "fadd.ps f15, f14, f21\n" // f15 = 1 + exp(-1.702*x)
+ "frcp.ps f16, f15\n" // f16 = sigmoid(1.702*x)
+ "fmul.ps f17, f10, f16\n" // f17 = gelu_quick(x)
+ "fmul.ps f18, f17, f11\n" // f18 = gelu_quick(x) * g
+ "fsw.ps f18, %[dst_out]\n"
+ : [dst_out] "=m"(*(float (*)[8]) & dst_block[i])
+ : [x_vec] "m"(*(const float (*)[8]) & x_block[i]), [g_vec] "m"(*(const float (*)[8]) & g_block[i]),
+ [zero_ptr] "m"(zero_const), [one_ptr] "m"(one_const), [k_ptr] "m"(neg_k_log2e_const)
+ : "f10", "f11", "f13", "f14", "f15", "f16", "f17", "f18", "f20", "f21", "f22");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ for (int32_t i = vec_end; i < elements; i++) {
+ float xv = x_block[i];
+ // Reuse silu reciprocal path: sigmoid(1.702*x) = 1/(1+exp(-1.702*x))
+ float e = et_expf(-1.702f * xv);
+ dst_block[i] = et_fdiv(xv, 1.0f + e) * g_block[i];
+ }
+}
+
+// Vectorized SwiGLU-OAI block (OpenAI gpt-oss variant):
+// x_c = min(x, limit)
+// y_c = clamp(g, -limit, limit)
+// out = (x_c / (1 + exp(-alpha * x_c))) * (y_c + 1)
+static inline void block_swiglu_oai(float * dst_block,
+ const float * x_block,
+ const float * g_block,
+ int elements,
+ float alpha,
+ float limit) {
+ int32_t vec_end = (elements / 8) * 8;
+
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ float zero_const = 0.0f;
+ float one_const = 1.0f;
+ float limit_pos = limit;
+ float limit_neg = -limit;
+ // -alpha * log2(e): feed (x * neg_alpha_log2e) into fexp.ps to get exp(-alpha*x)
+ float neg_alpha_l2e = -alpha * 1.4426950408889634f;
+
+ for (int32_t i = 0; i < vec_end; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x_vec]\n" // f10 = x raw
+ "flw.ps f11, %[g_vec]\n" // f11 = g raw
+
+ "fbc.ps f20, %[zero_ptr]\n" // f20 = 0
+ "fbc.ps f21, %[one_ptr]\n" // f21 = 1
+ "fbc.ps f23, %[lim_pos]\n" // f23 = +limit
+ "fbc.ps f24, %[lim_neg]\n" // f24 = -limit
+ "fbc.ps f25, %[k_ptr]\n" // f25 = -alpha*log2(e)
+
+ // x_c = min(x, +limit) (no lower bound on x per OAI spec)
+ "fmin.ps f12, f10, f23\n" // f12 = x_c
+
+ // y_c = clamp(g, -limit, +limit) = min(max(g, -limit), +limit)
+ "fmax.ps f13, f11, f24\n" // f13 = max(g, -limit)
+ "fmin.ps f13, f13, f23\n" // f13 = y_c
+
+ // sigmoid(alpha * x_c) = 1 / (1 + exp(-alpha * x_c))
+ "fmul.ps f14, f12, f25\n" // f14 = -alpha*x_c*log2(e)
+ "fexp.ps f15, f14\n" // f15 = exp(-alpha*x_c)
+ "fadd.ps f15, f15, f21\n" // f15 = 1 + exp(-alpha*x_c)
+ "frcp.ps f16, f15\n" // f16 = sigmoid(alpha*x_c)
+
+ // out_glu = x_c * sigmoid(alpha*x_c)
+ "fmul.ps f17, f12, f16\n" // f17 = swiglu_oai gate output
+
+ // dst = out_glu * (y_c + 1)
+ "fadd.ps f18, f13, f21\n" // f18 = y_c + 1
+ "fmul.ps f19, f17, f18\n" // f19 = final
+ "fsw.ps f19, %[dst_out]\n"
+
+ : [dst_out] "=m"(*(float (*)[8]) & dst_block[i])
+ : [x_vec] "m"(*(const float (*)[8]) & x_block[i]), [g_vec] "m"(*(const float (*)[8]) & g_block[i]),
+ [zero_ptr] "m"(zero_const), [one_ptr] "m"(one_const), [lim_pos] "m"(limit_pos), [lim_neg] "m"(limit_neg),
+ [k_ptr] "m"(neg_alpha_l2e)
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16", "f17", "f18", "f19", "f20", "f21", "f23", "f24", "f25");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ // Scalar tail (mirrors CPU reference exactly)
+ for (int32_t i = vec_end; i < elements; i++) {
+ float xv = x_block[i];
+ float yv = g_block[i];
+ if (xv > limit) {
+ xv = limit;
+ }
+ if (yv > limit) {
+ yv = limit;
+ }
+ if (yv < -limit) {
+ yv = -limit;
+ }
+ float e = et_expf(-alpha * xv);
+ float out_glu = et_fdiv(xv, 1.0f + e);
+ dst_block[i] = out_glu * (yv + 1.0f);
+ }
+}
+
+// Scalar erf approximation (Abramowitz & Stegun 7.1.26, max error ~1.5e-7)
+static inline float erf_approx(float x) {
+ const float a1 = 0.254829592f;
+ const float a2 = -0.284496736f;
+ const float a3 = 1.421413741f;
+ const float a4 = -1.453152027f;
+ const float a5 = 1.061405429f;
+ const float p = 0.3275911f;
+
+ float sign = (x < 0.0f) ? -1.0f : 1.0f;
+ float ax = (x < 0.0f) ? -x : x;
+ float t = et_fdiv(1.0f, 1.0f + p * ax);
+ float t2 = t * t;
+ float t3 = t2 * t;
+ float t4 = t3 * t;
+ float t5 = t4 * t;
+ float poly = a1 * t + a2 * t2 + a3 * t3 + a4 * t4 + a5 * t5;
+ float y = 1.0f - poly * et_expf(-ax * ax);
+ return sign * y;
+}
+
+// GeGLU-Erf block: dst = 0.5 * x * (1 + erf(x / sqrt(2))) * g
+// Scalar implementation — variant is rarely used so we keep complexity low.
+static inline void block_geglu_erf(float * dst_block, const float * x_block, const float * g_block, int elements) {
+ const float sqrt_2_inv = 0.70710678118654752440f;
+ for (int32_t i = 0; i < elements; i++) {
+ float xv = x_block[i];
+ dst_block[i] = 0.5f * xv * (1.0f + erf_approx(xv * sqrt_2_inv)) * g_block[i];
+ }
+}
+
+// Main entry point for GLU kernel
+int entry_point(struct ggml_et_glu_params * params, void * env) {
+ // Cast env to proper type
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ // Validate environment pointer
+ if (!kernel_env) {
+ return -1;
+ }
+
+ // Get thread info using shire mask from environment
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ // Basic safety check on params
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ // Supported variants: SwiGLU, SwiGLU-OAI, GeGLU, GeGLU-Erf, GeGLU-Quick, ReGLU
+ switch (params->glu_op_type) {
+ case GGML_GLU_OP_SWIGLU:
+ case GGML_GLU_OP_SWIGLU_OAI:
+ case GGML_GLU_OP_GEGLU:
+ case GGML_GLU_OP_GEGLU_ERF:
+ case GGML_GLU_OP_GEGLU_QUICK:
+ case GGML_GLU_OP_REGLU:
+ break;
+ default:
+ return -1; // Unsupported GLU operation
+ }
+
+ // Extract tensor references
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = params->src1.data ? ¶ms->src1 : 0;
+ struct ggml_tensor * dst = ¶ms->dst;
+ int32_t swapped = params->swapped;
+
+ // Validate tensor types (F32 only)
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ if (src1 && src1->type != GGML_TYPE_F32) {
+ return -1; // Unsupported src1 type
+ }
+
+ // Get data pointers
+ float * src0_data = (float *) src0->data;
+ float * src1_data = src1 ? (float *) src1->data : src0_data;
+ float * dst_data = (float *) dst->data;
+
+ // Validate data pointers
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ // Get tensor dimensions
+ const int64_t nc = dst->ne[0]; // Output columns (input columns / 2)
+ const int64_t nr = dst->ne[1] * dst->ne[2] * dst->ne[3]; // Total rows
+
+ // Get strides
+ const size_t src0_stride = src0->nb[1]; // Stride between rows in src0
+ const size_t src1_stride = src1 ? src1->nb[1] : src0->nb[1]; // Stride between rows in src1
+ const size_t dst_stride = dst->nb[1]; // Stride between rows in dst
+
+ // Validate dimensions for split SwiGLU
+ if (src1) {
+ // Split tensor mode: src0 and src1 should have same shape as dst
+ if (src0->ne[0] != nc || src1->ne[0] != nc) {
+ return -1; // Dimension mismatch in split mode
+ }
+ } else {
+ // Single tensor mode: src0 should have 2*nc columns
+ if (src0->ne[0] != 2 * nc) {
+ return -1; // Dimension mismatch in single tensor mode
+ }
+ }
+
+ // Calculate total elements for cache line distribution
+ const int64_t elements_per_cacheline = 16; // 64 bytes / 4 bytes per float
+ const int64_t total_elements = nr * nc;
+ const int64_t total_cachelines = (total_elements + elements_per_cacheline - 1) / elements_per_cacheline;
+
+ // Distribute cache lines across threads
+ int64_t cachelines_per_thread = (total_cachelines + num_threads - 1) / num_threads;
+ int64_t start_cacheline = thread_id * cachelines_per_thread;
+ int64_t end_cacheline = start_cacheline + cachelines_per_thread;
+
+ // Clamp end_cacheline to actual number of cache lines
+ if (end_cacheline > total_cachelines) {
+ end_cacheline = total_cachelines;
+ }
+
+ // Thread should return if no work to do
+ if (start_cacheline >= total_cachelines) {
+ return 0;
+ }
+
+ // Process cache lines assigned to this thread
+ for (int64_t cl = start_cacheline; cl < end_cacheline; cl++) {
+ // Map cache line back to element coordinates
+ int64_t global_element_start = cl * elements_per_cacheline;
+ int64_t row = global_element_start / nc;
+ int64_t col = global_element_start % nc;
+
+ // Skip if we're past the end of data
+ if (global_element_start >= total_elements) {
+ break;
+ }
+
+ // Calculate how many elements to process in this cache line
+ int64_t elements_remaining = total_elements - global_element_start;
+ int elements_this_block =
+ (int) ((elements_remaining < elements_per_cacheline) ? elements_remaining : elements_per_cacheline);
+
+ // Process elements that span across rows
+ int64_t elements_processed = 0;
+ while (elements_processed < elements_this_block && row < nr) {
+ // Calculate elements to process in current row
+ int64_t elements_in_row = nc - col;
+ int64_t elements_to_process = elements_this_block - elements_processed;
+ if (elements_to_process > elements_in_row) {
+ elements_to_process = elements_in_row;
+ }
+
+ // Get pointers for current row and column range
+ float * dst_ptr = (float *) ((char *) dst_data + row * dst_stride) + col;
+
+ float * x_ptr;
+ float * g_ptr;
+
+ if (src1) {
+ // Split tensor mode
+ x_ptr = (float *) ((char *) src0_data + row * src0_stride) + col;
+ g_ptr = (float *) ((char *) src1_data + row * src1_stride) + col;
+ } else {
+ // Single tensor mode - src0 contains both x and g
+ float * src0_row = (float *) ((char *) src0_data + row * src0_stride);
+ if (swapped) {
+ g_ptr = src0_row + col; // First half is gate
+ x_ptr = src0_row + nc + col; // Second half is value
+ } else {
+ x_ptr = src0_row + col; // First half is value
+ g_ptr = src0_row + nc + col; // Second half is gate
+ }
+ }
+
+ // Process this segment
+ switch (params->glu_op_type) {
+ case GGML_GLU_OP_GEGLU:
+ block_geglu(dst_ptr, x_ptr, g_ptr, (int) elements_to_process);
+ break;
+ case GGML_GLU_OP_SWIGLU:
+ block_swiglu(dst_ptr, x_ptr, g_ptr, (int) elements_to_process);
+ break;
+ case GGML_GLU_OP_REGLU:
+ block_reglu(dst_ptr, x_ptr, g_ptr, (int) elements_to_process);
+ break;
+ case GGML_GLU_OP_GEGLU_QUICK:
+ block_geglu_quick(dst_ptr, x_ptr, g_ptr, (int) elements_to_process);
+ break;
+ case GGML_GLU_OP_GEGLU_ERF:
+ block_geglu_erf(dst_ptr, x_ptr, g_ptr, (int) elements_to_process);
+ break;
+ case GGML_GLU_OP_SWIGLU_OAI:
+ block_swiglu_oai(dst_ptr, x_ptr, g_ptr, (int) elements_to_process, params->alpha, params->limit);
+ break;
+ default:
+ return -1;
+ }
+
+ // Update counters
+ elements_processed += elements_to_process;
+ col += elements_to_process;
+
+ // Move to next row if current row is complete
+ if (col >= nc) {
+ row++;
+ col = 0;
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// GROUP_NORM F32 Kernel
+// Baseline scalar implementation:
+// normalize over (ne0 * ne1 * channels_in_group) for each (group, batch).
+//
+// Parallelization:
+// - Work is partitioned across (group, batch) pairs.
+// - For non-cache-aligned ne0, writes are emitted in row-groups so each thread's
+// destination write footprint still spans an integer number of cache lines.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_group_norm_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+ int32_t n_groups;
+ float eps;
+};
+
+int entry_point(struct ggml_et_group_norm_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * src0_data = (const float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int32_t n_groups = params->n_groups;
+ const float eps = params->eps;
+
+ if (n_groups <= 0 || eps < 0.0f) {
+ return -1;
+ }
+
+ const int64_t ne0 = dst->ne[0];
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t ne3 = dst->ne[3];
+
+ if (src0->ne[0] != ne0 || src0->ne[1] != ne1 || src0->ne[2] != ne2 || src0->ne[3] != ne3) {
+ return -1;
+ }
+
+ const int64_t nb1 = dst->nb[1];
+ const int64_t nb2 = dst->nb[2];
+ const int64_t nb3 = dst->nb[3];
+ const int64_t nb01 = src0->nb[1];
+ const int64_t nb02 = src0->nb[2];
+ const int64_t nb03 = src0->nb[3];
+
+ const int64_t channels_per_group = (ne2 + n_groups - 1) / n_groups;
+ if (channels_per_group <= 0) {
+ return -1;
+ }
+
+ const int64_t active_groups = (ne2 + channels_per_group - 1) / channels_per_group;
+ const int64_t total_work = active_groups * ne3;
+ const int64_t rows_per_write_group = et_rows_per_cacheline_group(ne0, sizeof(float));
+
+ for (int64_t work = thread_id; work < total_work; work += num_threads) {
+ const int64_t i3 = work / active_groups;
+ const int64_t group_idx = work % active_groups;
+
+ const int64_t channel_start = group_idx * channels_per_group;
+ int64_t channel_end = channel_start + channels_per_group;
+ if (channel_end > ne2) {
+ channel_end = ne2;
+ }
+
+ const int64_t channel_count = channel_end - channel_start;
+ if (channel_count <= 0) {
+ continue;
+ }
+
+ float sum = 0.0f;
+ float denom = 0.0f;
+ for (int64_t i2 = channel_start; i2 < channel_end; ++i2) {
+ for (int64_t i1 = 0; i1 < ne1; ++i1) {
+ const float * src_row = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ sum += src_row[i0];
+ denom += 1.0f;
+ }
+ }
+ }
+
+ const float mean = et_fdiv(sum, denom);
+
+ float var_sum = 0.0f;
+ for (int64_t i2 = channel_start; i2 < channel_end; ++i2) {
+ for (int64_t i1 = 0; i1 < ne1; ++i1) {
+ const float * src_row = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ const float centered = src_row[i0] - mean;
+ var_sum += centered * centered;
+ }
+ }
+ }
+
+ const float variance = et_fdiv(var_sum, denom);
+ const float scale = et_fdiv(1.0f, et_sqrtf(variance + eps));
+
+ if (ne0 % 16 == 0) {
+ for (int64_t i2 = channel_start; i2 < channel_end; ++i2) {
+ for (int64_t i1 = 0; i1 < ne1; ++i1) {
+ const float * src_row =
+ (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_row = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ dst_row[i0] = (src_row[i0] - mean) * scale;
+ }
+ }
+ }
+ } else {
+ const int64_t total_rows_in_group = channel_count * ne1;
+ const int64_t total_write_groups = (total_rows_in_group + rows_per_write_group - 1) / rows_per_write_group;
+
+ for (int64_t write_group = 0; write_group < total_write_groups; ++write_group) {
+ const int64_t row_start = write_group * rows_per_write_group;
+ int64_t row_end = row_start + rows_per_write_group;
+ if (row_end > total_rows_in_group) {
+ row_end = total_rows_in_group;
+ }
+
+ for (int64_t row = row_start; row < row_end; ++row) {
+ const int64_t local_i2 = row / ne1;
+ const int64_t i1 = row % ne1;
+ const int64_t i2 = channel_start + local_i2;
+
+ const float * src_row =
+ (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_row = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+ for (int64_t i0 = 0; i0 < ne0; ++i0) {
+ dst_row[i0] = (src_row[i0] - mean) * scale;
+ }
+ }
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// IM2COL Kernel
+// Baseline scalar implementation for:
+// src1: [N, IC, IH, IW] -> dst: [N, OH, OW, IC*KH*KW] (2D)
+// src1: [N, IC, IW] -> dst: [N, 1, OW, IC* KW] (1D)
+//
+// Work is distributed by row-groups so threads own cache-line-aligned chunks of
+// destination rows even when ne0 is not cache aligned.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+static inline void im2col_store_elem(void * dst_base, enum ggml_type dst_type, int64_t idx, float value) {
+ if (dst_type == GGML_TYPE_F32) {
+ ((float *) dst_base)[idx] = value;
+ } else {
+ ((uint16_t *) dst_base)[idx] = fp32_to_fp16(value);
+ }
+}
+
+static inline float im2col_load_src_elem(const void * src_base, enum ggml_type src_type, int64_t idx) {
+ if (src_type == GGML_TYPE_F32) {
+ return ((const float *) src_base)[idx];
+ }
+
+ return fp16_to_fp32(((const uint16_t *) src_base)[idx]);
+}
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env || params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (!src1->data || !dst->data) {
+ return -1;
+ }
+
+ if (!((dst->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) ||
+ (dst->type == GGML_TYPE_F16 && (src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32)))) {
+ return -1;
+ }
+
+ const int32_t s0 = ((const int32_t *) dst->op_params)[0];
+ const int32_t s1 = ((const int32_t *) dst->op_params)[1];
+ const int32_t p0 = ((const int32_t *) dst->op_params)[2];
+ const int32_t p1 = ((const int32_t *) dst->op_params)[3];
+ const int32_t d0 = ((const int32_t *) dst->op_params)[4];
+ const int32_t d1 = ((const int32_t *) dst->op_params)[5];
+ const int32_t is_2d = ((const int32_t *) dst->op_params)[6];
+
+ const int64_t N = is_2d ? src1->ne[3] : src1->ne[2];
+ const int64_t IC = is_2d ? src1->ne[2] : src1->ne[1];
+ const int64_t IH = is_2d ? src1->ne[1] : 1;
+ const int64_t IW = src1->ne[0];
+
+ const int64_t KH = is_2d ? src0->ne[1] : 1;
+ const int64_t KW = src0->ne[0];
+
+ const int64_t OH = is_2d ? dst->ne[2] : 1;
+ const int64_t OW = dst->ne[1];
+ const int64_t row_elems = dst->ne[0];
+ const int64_t total_rows = OW * OH * N;
+
+ const size_t src_batch_stride = is_2d ? src1->nb[3] : src1->nb[2];
+ const size_t src_channel_stride = is_2d ? src1->nb[2] : src1->nb[1];
+
+ const size_t dst_row_stride = dst->nb[1];
+ const size_t dst_plane_stride = is_2d ? dst->nb[2] : 0;
+ const size_t dst_batch_stride = is_2d ? dst->nb[3] : dst->nb[2];
+
+ const int64_t dst_elem_size = (dst->type == GGML_TYPE_F32) ? (int64_t) sizeof(float) : (int64_t) sizeof(uint16_t);
+ const int64_t rows_per_group = et_rows_per_cacheline_group(row_elems, dst_elem_size);
+ const int64_t total_groups = (total_rows + rows_per_group - 1) / rows_per_group;
+
+ for (int64_t grp = thread_id; grp < total_groups; grp += num_threads) {
+ const int64_t row_start = grp * rows_per_group;
+ int64_t row_end = row_start + rows_per_group;
+ if (row_end > total_rows) {
+ row_end = total_rows;
+ }
+
+ for (int64_t row = row_start; row < row_end; ++row) {
+ const int64_t in = row / (OH * OW);
+ const int64_t rem = row % (OH * OW);
+ const int64_t ioh = rem / OW;
+ const int64_t iow = rem % OW;
+
+ void * dst_row = (char *) dst->data + in * dst_batch_stride + ioh * dst_plane_stride + iow * dst_row_stride;
+
+ for (int64_t iic = 0; iic < IC; ++iic) {
+ const void * src_channel = (const char *) src1->data + in * src_batch_stride + iic * src_channel_stride;
+
+ for (int64_t ikh = 0; ikh < KH; ++ikh) {
+ for (int64_t ikw = 0; ikw < KW; ++ikw) {
+ const int64_t iiw = iow * s0 + ikw * d0 - p0;
+ const int64_t iih = ioh * s1 + ikh * d1 - p1;
+ const int64_t dst_idx = iic * (KH * KW) + ikh * KW + ikw;
+
+ if (iiw < 0 || iiw >= IW || iih < 0 || iih >= IH) {
+ im2col_store_elem(dst_row, dst->type, dst_idx, 0.0f);
+ } else {
+ const int64_t src_idx = iih * IW + iiw;
+ const float value = im2col_load_src_elem(src_channel, src1->type, src_idx);
+ im2col_store_elem(dst_row, dst->type, dst_idx, value);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// L2 Norm F32 Kernel (L2 Normalization)
+// y[i] = x[i] / max(sqrt(sum(x^2)), eps)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <stdint.h>
+#include <string.h>
+
+// L2 Norm kernel parameters structure
+struct ggml_et_l2_norm_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon parameter for numerical stability
+};
+
+int entry_point(struct ggml_et_l2_norm_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+ float eps = params->eps;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ if (eps < 0.0f) {
+ return -1; // Invalid epsilon
+ }
+
+ const int64_t ne0 = dst->ne[0];
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t ne3 = dst->ne[3];
+
+ const size_t nb0 = dst->nb[0], nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+ const size_t nb00 = src0->nb[0], nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+
+ (void) nb0;
+ (void) nb00;
+
+ if (src0->ne[0] != ne0 || src0->ne[1] != ne1 || src0->ne[2] != ne2 || src0->ne[3] != ne3) {
+ return -1; // Shape mismatch
+ }
+
+ const int32_t total_rows = (int32_t) (ne1 * ne2 * ne3);
+ const int shire_threads = SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+
+ if (total_rows >= shire_threads) {
+ // Row-parallel: each thread processes whole rows
+ for (int64_t i3 = 0; i3 < ne3; i3++) {
+ for (int64_t i2 = 0; i2 < ne2; i2++) {
+ for (int64_t i1 = thread_id; i1 < ne1; i1 += num_threads) {
+ const float * src_ptr =
+ (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fmadd.ps f10, f11, f11, f10\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ float sum_sq;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(sum_sq)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ float l2_norm = et_powf(sum_sq, 0.5f);
+ if (l2_norm < eps) {
+ l2_norm = eps;
+ }
+ const float scale = et_fdiv(1.0f, l2_norm);
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fbc.ps f12, %[scale_ptr]\n"
+ "fmul.ps f13, f11, f12\n"
+ "fsw.ps f13, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0]), [scale_ptr] "m"(scale)
+ : "f11", "f12", "f13");
+ }
+ }
+ }
+ }
+ } else {
+ // Intra-row: threads within each shire cooperate via L2 SCP
+ int shire_tid = thread_id % shire_threads;
+ int threads_per_row = shire_threads / total_rows;
+ int my_row = shire_tid / threads_per_row;
+ int local_tid = shire_tid % threads_per_row;
+ int group_base = my_row * threads_per_row;
+
+ if (my_row >= total_rows) {
+ FENCE;
+ et_barrier(ET_BARRIER_SHIRE);
+ return 0;
+ }
+
+ int64_t i1 = my_row % ne1;
+ int64_t i2 = (my_row / ne1) % ne2;
+ int64_t i3 = my_row / (ne1 * ne2);
+
+ const float * src_ptr = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ const int32_t elems_per_cl = 16;
+ int32_t total_cls = ((int32_t) ne0 + elems_per_cl - 1) / elems_per_cl;
+ int32_t cls_per_thread = (total_cls + threads_per_row - 1) / threads_per_row;
+ int32_t my_start = local_tid * cls_per_thread * elems_per_cl;
+ int32_t my_end = my_start + cls_per_thread * elems_per_cl;
+ if (my_end > (int32_t) ne0) {
+ my_end = (int32_t) ne0;
+ }
+ if (my_start >= (int32_t) ne0) {
+ my_start = 0;
+ my_end = 0;
+ }
+
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Phase 1: partial sum of squares
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fmadd.ps f10, f11, f11, f10\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ float partial_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(partial_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ // Phase 2: L2SCP exchange
+ volatile float * my_slot = (volatile float *) et_shire_l2scp_local((uint64_t) shire_tid * 64);
+ *my_slot = partial_sum;
+ FENCE;
+ evict_to_l2((const void *) my_slot, 1, 64);
+ WAIT_CACHEOPS;
+
+ et_barrier(ET_BARRIER_SHIRE);
+
+ // Phase 3: all threads reduce + apply scale to own chunk
+ int workers = threads_per_row < total_cls ? threads_per_row : total_cls;
+
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ evict_to_l2((const void *) slot, 1, 64);
+ }
+ WAIT_CACHEOPS;
+
+ float total_sum_sq = 0.0f;
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ total_sum_sq += *slot;
+ }
+
+ float l2_norm = et_powf(total_sum_sq, 0.5f);
+ if (l2_norm < eps) {
+ l2_norm = eps;
+ }
+ const float scale = et_fdiv(1.0f, l2_norm);
+
+ if (my_start < my_end) {
+ uint32_t scale_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(scale_bits) : "f"(scale));
+ __asm__ volatile("fbcx.ps f13, %[sb]\n" : : [sb] "r"(scale_bits) : "f13");
+
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[x_vec]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f12", "f14");
+ }
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+
+ return 0;
+}
--- /dev/null
+/*-------------------------------------------------------------------------
+ * Independent Linker Script for GGML Bare Metal Kernels
+ *-------------------------------------------------------------------------
+ */
+
+/* Define maximum number of harts (threads) - simplified for bare metal */
+MAX_HARTS = 2112;
+
+SECTIONS
+{
+ /* Start at the base address passed by -Wl,--defsym=BASE_ADDRESS=... */
+ . = BASE_ADDRESS;
+
+ /* Export entry point symbol for runtime compatibility */
+ KERNEL_UMODE_ENTRY = BASE_ADDRESS;
+
+ /* Initialization section - must come first */
+ .text.init :
+ {
+ *(.text.init)
+ }
+
+ /* Align to cache line boundary */
+ . = ALIGN(64);
+
+ /* Main text section for code */
+ .text : {
+ *(.text)
+ *(.text.*)
+ }
+ . = ALIGN(64);
+
+ /* Data section with global pointer setup */
+ .data :
+ {
+ _data_start = .;
+ . = ALIGN(64);
+
+ /* Small data section and global pointer */
+ *(.sdata .sdata.*)
+ PROVIDE( __global_pointer$ = . + 0x800 );
+
+ /* Regular data */
+ *(.data .data.*)
+ . = ALIGN(64);
+ _data_end = .;
+ }
+ . = ALIGN(64);
+
+ /* BSS section for uninitialized data */
+ .bss(NOLOAD) :
+ {
+ _bss_start = .;
+ *(.sbss*);
+ *(.bss*);
+ . = ALIGN(64);
+ _bss_end = .;
+ }
+
+ /* Thread Local Storage (TLS) sections */
+ . = ALIGN(64);
+ .tdata :
+ {
+ *(.tdata*)
+ . = ALIGN(64);
+ }
+ __tdata_start = ADDR(.tdata);
+
+ .tbss : {
+ __tbss_start = .;
+ *(.tbss*)
+ }
+ . = . + SIZEOF(.tbss);
+ . = ALIGN(64);
+ __tbss_end = .;
+
+ /* TLS allocation area for all harts */
+ .tls-alloc ALIGN(64) (NOLOAD) : {
+ __tls_alloc_start = .;
+ . = . + (ABSOLUTE(__tbss_end) - ABSOLUTE(__tdata_start)) * MAX_HARTS;
+ }
+
+ /* End of kernel image */
+ _end = .;
+}
--- /dev/null
+//******************************************************************************
+// ET Floating Point Math Library
+// Provides ET hardware-specific math functions, FP16 conversion, and trig functions
+// for bare metal kernels
+//******************************************************************************
+
+#ifndef MATH_FP_H
+#define MATH_FP_H
+
+#include <stdint.h>
+
+//******************************************************************************
+// ET Hardware Math Functions
+//******************************************************************************
+
+// ET hardware division function (uses FRCP.PS instruction)
+static inline float et_fdiv(float a, float b) {
+ float d;
+ unsigned long temp;
+
+ __asm__ volatile(
+ "mova.x.m %[temp] \n\t"
+ "mov.m.x m0, x0, 1 \n\t"
+ "frcp.ps %[d], %[b] \n\t"
+ "fmul.s %[d], %[d], %[a] \n\t"
+ "mova.m.x %[temp] \n\t"
+ : [temp] "=&r"(temp), [d] "=&f"(d)
+ : [a] "f"(a), [b] "f"(b));
+
+ return d;
+}
+
+// Power function using ET hardware vector instructions
+// Implements pow(base, exp) = exp(exp * ln(base)) using FLOG.PS and FEXP.PS
+static inline float et_powf(float base, float exp) {
+ // Handle special cases
+ if (base <= 0.0f) {
+ if (base == 0.0f) {
+ if (exp > 0.0f) {
+ return 0.0f;
+ }
+
+ // For exp <= 0, return +infinity (IEEE 754: sign=0, exp=0xFF, mantissa=0)
+ union {
+ float f;
+ uint32_t i;
+ } inf = { .i = 0x7F800000 };
+
+ return inf.f;
+ }
+
+ // For negative base, return NaN (IEEE 754: exp=0xFF, mantissa!=0)
+ union {
+ float f;
+ uint32_t i;
+ } nan = { .i = 0x7FC00000 };
+
+ return nan.f;
+ }
+ if (base == 1.0f) {
+ return 1.0f;
+ }
+ if (exp == 0.0f) {
+ return 1.0f;
+ }
+ if (exp == 1.0f) {
+ return base;
+ }
+
+ // Use ET hardware instructions following DNN library pattern:
+ // pow(base, exp) = exp(exp * ln(base))
+ float result;
+ unsigned long temp;
+
+ __asm__ volatile(
+ "mova.x.m %[temp] \n\t" // Save current mask state
+ "mov.m.x m0, x0, 1 \n\t" // Set mask register m0 to enable element 0
+ "flog.ps %[result], %[base] \n\t" // result = ln(base)
+ "fmul.s %[result], %[result], %[exp]\n\t" // result = ln(base) * exp
+ "fexp.ps %[result], %[result] \n\t" // result = exp(ln(base) * exp) = base^exp
+ "mova.m.x %[temp] \n\t" // Restore mask state
+ : [temp] "=&r"(temp), [result] "=&f"(result)
+ : [base] "f"(base), [exp] "f"(exp));
+
+ return result;
+}
+
+// Natural logarithm.
+static inline float et_logf(float x) {
+ // Handle special cases
+ if (x < 0.0f) {
+ // Return NaN for negative input
+ union {
+ float f;
+ uint32_t i;
+ } nan = { .i = 0x7FC00000 };
+
+ return nan.f;
+ }
+ if (x == 0.0f) {
+ // Return -infinity for log(0)
+ union {
+ float f;
+ uint32_t i;
+ } inf = { .i = 0xFF800000 };
+
+ return inf.f;
+ }
+ if (x == 1.0f) {
+ return 0.0f;
+ }
+
+ float log2_result;
+ unsigned long temp;
+
+ __asm__ volatile(
+ "mova.x.m %[temp] \n\t" // Save current mask state
+ "mov.m.x m0, x0, 1 \n\t" // Set mask register m0 to enable element 0
+ "flog.ps %[result], %[x] \n\t" // result = log2(x)
+ "mova.m.x %[temp] \n\t" // Restore mask state
+ : [temp] "=&r"(temp), [result] "=&f"(log2_result)
+ : [x] "f"(x));
+
+ // Convert log2 to natural log: ln(x) = log2(x) * ln(2)
+ const float ln2 = 0.69314718055994530942f;
+ return log2_result * ln2;
+}
+
+// Square root function implemented as et_powf(x, 0.5)
+static inline float et_sqrtf(float x) {
+ // Handle special cases
+ if (x < 0.0f) {
+ // Return NaN for negative input (IEEE 754: exp=0xFF, mantissa!=0)
+ union {
+ float f;
+ uint32_t i;
+ } nan = { .i = 0x7FC00000 };
+
+ return nan.f;
+ }
+ if (x == 0.0f) {
+ return 0.0f;
+ }
+
+ return et_powf(x, 0.5f);
+}
+
+// Base-2 exponential: returns 2^x using the ET hardware FEXP.PS instruction.
+// No base conversion, no special-case clamping — this is the raw hardware op
+// with just the mask save/restore wrapper. Caller is responsible for ensuring
+// x is in a range that produces a useful result (roughly [-126, 128] for fp32).
+static inline float __attribute__((always_inline)) et_exp2f(float x) {
+ unsigned long old_mask;
+ float out;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 1 \n\t"
+ "fexp.ps %[out], %[x] \n\t"
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(old_mask), [out] "=&f"(out)
+ : [x] "f"(x));
+ return out;
+}
+
+// Exponential function using ET hardware FEXP.PS instruction
+// Note: FEXP.PS computes 2^x, so we need to convert: exp(x) = 2^(x * log2(e))
+static inline float et_expf(float x) {
+ // Handle special cases
+ if (x > 88.0f) {
+ // For x > 88, exp(x) would overflow, return +infinity
+ union {
+ float f;
+ uint32_t i;
+ } inf = { .i = 0x7F800000 };
+
+ return inf.f;
+ }
+ if (x < -87.0f) {
+ // For x < -87, exp(x) is essentially 0
+ return 0.0f;
+ }
+
+ // Convert to base-2 exponent: x * log2(e)
+ const float log2e = 1.4426950408889634f; // log2(e)
+ float x_log2e = x * log2e;
+
+ // Use ET hardware instruction: fexp.ps computes 2^x
+ float result;
+ unsigned long temp;
+
+ __asm__ volatile(
+ "mova.x.m %[temp] \n\t" // Save current mask state
+ "mov.m.x m0, x0, 1 \n\t" // Set mask register m0 to enable element 0
+ "fexp.ps %[result], %[x_log2e] \n\t" // result = 2^(x * log2(e)) = exp(x)
+ "mova.m.x %[temp] \n\t" // Restore mask state
+ : [temp] "=&r"(temp), [result] "=&f"(result)
+ : [x_log2e] "f"(x_log2e));
+
+ return result;
+}
+
+//******************************************************************************
+// Trigonometric Functions
+//******************************************************************************
+
+// FSIN.PS
+
+// Sine function using Taylor series
+static inline float et_sinf(float x) {
+ const float pi = 3.14159265358979323846f;
+ const float two_pi = 6.28318530717958647693f;
+ const float pi_over_2 = 1.57079632679489661923f;
+
+ if (x > pi || x < -pi) {
+ float cycles = x * et_fdiv(1.0f, two_pi);
+ int n = (int) cycles;
+ if (x < 0.0f) {
+ n--; // Floor for negative
+ }
+ x = x - (float) n * two_pi;
+ }
+
+ // sin(x) = sin(π - x) for x in [π/2, π]
+ // sin(x) = -sin(-π - x) for x in [-π, -π/2]
+ int negate = 0;
+ if (x > pi_over_2) {
+ x = pi - x;
+ } else if (x < -pi_over_2) {
+ x = -pi - x;
+ negate = 1;
+ }
+
+ // sin(x) ≈ x - x^3/3! + x^5/5! - x^7/7! + x^9/9! - x^11/11!
+ const float x2 = x * x;
+ const float x3 = x2 * x;
+ const float x5 = x3 * x2;
+ const float x7 = x5 * x2;
+ const float x9 = x7 * x2;
+ const float x11 = x9 * x2;
+
+ float result = x - x3 * et_fdiv(1.0f, 6.0f) // x^3/3!
+ + x5 * et_fdiv(1.0f, 120.0f) // x^5/5!
+ - x7 * et_fdiv(1.0f, 5040.0f) // x^7/7!
+ + x9 * et_fdiv(1.0f, 362880.0f) // x^9/9!
+ - x11 * et_fdiv(1.0f, 39916800.0f); // x^11/11!
+
+ return negate ? -result : result;
+}
+
+// Cosine function using identity cos(x) = sin(x + π/2)
+static inline float et_cosf(float x) {
+ const float pi_over_2 = 1.57079632679489661923f;
+ return et_sinf(x + pi_over_2);
+}
+
+//******************************************************************************
+// FP16 <-> FP32 Conversion Functions
+//******************************************************************************
+
+// Convert FP16 (IEEE 754 half precision) to FP32 (single precision)
+// Uses ET hardware FCVT.PS.F16 instruction for accurate conversion
+static inline float fp16_to_fp32(uint16_t h) {
+ float result;
+ unsigned long temp;
+ uint32_t raw = (uint32_t) h;
+
+ __asm__ volatile(
+ "mova.x.m %[temp] \n\t" // Save current mask state
+ "mov.m.x m0, x0, 1 \n\t" // Set mask register m0 to enable element 0
+ "fbcx.ps %[result], %[raw] \n\t" // Broadcast raw FP16 bits into vector register
+ "fcvt.ps.f16 %[result], %[result] \n\t" // Convert FP16 to FP32
+ "mova.m.x %[temp] \n\t" // Restore mask state
+ : [temp] "=&r"(temp), [result] "=&f"(result)
+ : [raw] "r"(raw));
+
+ return result;
+}
+
+// Convert FP32 (single precision) to FP16 (IEEE 754 half precision)
+// Uses ET hardware FCVT.F16.PS instruction for accurate conversion
+static inline uint16_t fp32_to_fp16(float f) {
+ float result_f;
+ unsigned long temp;
+
+ __asm__ volatile(
+ "mova.x.m %[temp] \n\t" // Save current mask state
+ "mov.m.x m0, x0, 1 \n\t" // Set mask register m0 to enable element 0
+ "fcvt.f16.ps %[result], %[f] \n\t" // Convert FP32 to FP16 (result in lower 16 bits)
+ "mova.m.x %[temp] \n\t" // Restore mask state
+ : [temp] "=&r"(temp), [result] "=&f"(result_f)
+ : [f] "f"(f));
+
+ // Extract lower 16 bits containing the FP16 value
+ // The instruction zero-extends to 32 bits, so upper 16 bits are 0
+ uint32_t result_bits = *(uint32_t *) &result_f;
+ return (uint16_t) result_bits;
+}
+
+#endif // MATH_FP_H
--- /dev/null
+//******************************************************************************
+// MEAN F32 Kernel
+// Row-wise mean reduction: dst[0, i1, i2, i3] = mean(src0[0..ne00-1, i1, i2, i3])
+//
+// Modes:
+// - total_rows >= shire_threads: row-parallel, each thread handles whole rows.
+// - total_rows < shire_threads: intra-row reduction within a shire. Threads
+// within a shire cooperate via shire-local L2 SCP slots. All shires
+// duplicate the work because L2 SCP is per-shire (no cross-shire coherency).
+//
+// ne00 may be any positive size and rows may have any 4-byte alignment. We
+// take the 8-wide vector path only when the row pointer is 32B-aligned and
+// fall back to scalar for the leftover tail (or for the entire row when the
+// row start is not 32B-aligned).
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_mean_params {
+ struct ggml_tensor src0; // F32 input [ne00, ne01, ne02, ne03]
+ struct ggml_tensor dst; // F32 output [1, ne01, ne02, ne03]
+};
+
+// Sum a contiguous F32 slice [base+i_lo, base+i_hi). Uses the 8-wide vector
+// path only when `base + i_lo` is 32B-aligned; the tail (and the whole slice
+// when misaligned) is summed with scalar fadd.s.
+static inline float partial_sum_slice(const float * base, int32_t i_lo, int32_t i_hi) {
+ if (i_lo >= i_hi) {
+ return 0.0f;
+ }
+
+ const float * p = base + i_lo;
+ int32_t n = i_hi - i_lo;
+ float acc = 0.0f;
+ int32_t i = 0;
+
+ if (n >= 8 && (((uintptr_t) p) & 31) == 0) {
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (; i + 8 <= n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x]\n"
+ "fadd.ps f10, f10, f11\n"
+ :
+ : [x] "m"(*(const float (*)[8]) & p[i])
+ : "f10", "f11");
+ }
+
+ float vec_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(vec_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+ acc = vec_sum;
+ }
+
+ for (; i < n; i++) {
+ acc += p[i];
+ }
+ return acc;
+}
+
+int entry_point(struct ggml_et_mean_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne00 = src0->ne[0];
+ const int64_t ne01 = src0->ne[1];
+ const int64_t ne02 = src0->ne[2];
+ const int64_t ne03 = src0->ne[3];
+
+ const size_t nb01 = src0->nb[1];
+ const size_t nb02 = src0->nb[2];
+ const size_t nb03 = src0->nb[3];
+
+ const size_t nb1 = dst->nb[1];
+ const size_t nb2 = dst->nb[2];
+ const size_t nb3 = dst->nb[3];
+
+ if (ne00 <= 0) {
+ return 0;
+ }
+
+ const int32_t total_rows = (int32_t) (ne01 * ne02 * ne03);
+ const int shire_threads = SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+ const float inv_ne00 = et_fdiv(1.0f, (float) (int32_t) ne00);
+
+ // Row-parallel: each thread owns whole rows.
+ if (total_rows >= shire_threads) {
+ for (int64_t ir = thread_id; ir < total_rows; ir += num_threads) {
+ const int64_t i03 = ir / (ne02 * ne01);
+ const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01;
+ const int64_t i01 = ir - i03 * ne02 * ne01 - i02 * ne01;
+
+ const float * src_row = (const float *) ((const char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ float * dst_ptr = (float *) ((char *) dst_data + i01 * nb1 + i02 * nb2 + i03 * nb3);
+
+ float row_sum = partial_sum_slice(src_row, 0, (int32_t) ne00);
+ atomic_store_f32(dst_ptr, row_sum * inv_ne00);
+ }
+ // Shire co-work
+ } else {
+ int shire_tid = thread_id % shire_threads;
+ int threads_per_row = shire_threads / total_rows;
+ int my_row = shire_tid / threads_per_row;
+ int local_tid = shire_tid % threads_per_row;
+ int group_base = my_row * threads_per_row;
+
+ if (my_row >= total_rows) {
+ FENCE;
+ et_barrier(ET_BARRIER_SHIRE);
+ return 0;
+ }
+
+ int64_t i1 = my_row % ne01;
+ int64_t i2 = (my_row / ne01) % ne02;
+ int64_t i3 = my_row / (ne01 * ne02);
+
+ const float * src_ptr = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ // Chunk size in elements, rounded up to a multiple of 8 so that every
+ // thread's slice start stays 32B-aligned relative to src_ptr (which
+ // matters for the vector path inside partial_sum_slice).
+ int32_t chunk = ((int32_t) ne00 + threads_per_row - 1) / threads_per_row;
+ chunk = (chunk + 7) & ~7;
+ if (chunk < 8) {
+ chunk = 8;
+ }
+
+ int32_t my_start = local_tid * chunk;
+ int32_t my_end = my_start + chunk;
+ if (my_end > (int32_t) ne00) {
+ my_end = (int32_t) ne00;
+ }
+ if (my_start > (int32_t) ne00) {
+ my_start = my_end = (int32_t) ne00;
+ }
+
+ int workers = ((int32_t) ne00 + chunk - 1) / chunk;
+ if (workers > threads_per_row) {
+ workers = threads_per_row;
+ }
+
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ float partial_sum = partial_sum_slice(src_ptr, my_start, my_end);
+
+ // Publish partial to shire-local L2 SCP slot (64B per slot, one per
+ // hart). evict_to_l2 is required on the WRITER because scalar stores
+ // land in L1D first; readers must also evict before reading.
+ volatile float * my_slot = (volatile float *) et_shire_l2scp_local((uint64_t) shire_tid * 64);
+ *my_slot = partial_sum;
+ FENCE;
+ evict_to_l2((const void *) my_slot, 1, 64);
+ WAIT_CACHEOPS;
+
+ et_barrier(ET_BARRIER_SHIRE);
+
+ if (local_tid == 0) {
+ // Reader-side evictions for every contributing peer slot.
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ evict_to_l2((const void *) slot, 1, 64);
+ }
+ WAIT_CACHEOPS;
+
+ float total_sum = 0.0f;
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ total_sum += *slot;
+ }
+ atomic_store_f32(dst_ptr, total_sum * inv_ne00);
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Memory Operations Kernel — tensor_store based memset
+//
+// Uses the tensor engine's store path (bypasses L1+L2 caches) to achieve hiher
+// performance. Unrolled vector writes can write at ~25GB/s and tensor writes
+// can so ~71 GB/s. Only even harts (hart 0 per minion) participate, as due to
+// hardware design (only thye have matrix engine access and co-op stores seems
+// slower)
+//******************************************************************************
+
+#include "platform.h"
+#include "tensor.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+
+// Operation identifiers for memops kernel
+enum ggml_et_memop_type {
+ GGML_ET_MEMOP_MEMSET = 0,
+};
+
+// Memset operation parameters (must match host-side struct in ggml-et-memops.cpp)
+struct memset_params {
+ uint32_t op_type;
+ uint32_t value;
+ void * dst_ptr;
+ size_t size;
+};
+
+// Fill all 32 f-regs with a replicated byte pattern
+static inline void __attribute__((always_inline)) fill_fregs(uint32_t fill32) {
+ register uint64_t val __asm__("a2") = fill32;
+ __asm__ __volatile__(
+ "fbcx.ps f0, %[v]\n\t"
+ "fbcx.ps f1, %[v]\n\t"
+ "fbcx.ps f2, %[v]\n\t"
+ "fbcx.ps f3, %[v]\n\t"
+ "fbcx.ps f4, %[v]\n\t"
+ "fbcx.ps f5, %[v]\n\t"
+ "fbcx.ps f6, %[v]\n\t"
+ "fbcx.ps f7, %[v]\n\t"
+ "fbcx.ps f8, %[v]\n\t"
+ "fbcx.ps f9, %[v]\n\t"
+ "fbcx.ps f10, %[v]\n\t"
+ "fbcx.ps f11, %[v]\n\t"
+ "fbcx.ps f12, %[v]\n\t"
+ "fbcx.ps f13, %[v]\n\t"
+ "fbcx.ps f14, %[v]\n\t"
+ "fbcx.ps f15, %[v]\n\t"
+ "fbcx.ps f16, %[v]\n\t"
+ "fbcx.ps f17, %[v]\n\t"
+ "fbcx.ps f18, %[v]\n\t"
+ "fbcx.ps f19, %[v]\n\t"
+ "fbcx.ps f20, %[v]\n\t"
+ "fbcx.ps f21, %[v]\n\t"
+ "fbcx.ps f22, %[v]\n\t"
+ "fbcx.ps f23, %[v]\n\t"
+ "fbcx.ps f24, %[v]\n\t"
+ "fbcx.ps f25, %[v]\n\t"
+ "fbcx.ps f26, %[v]\n\t"
+ "fbcx.ps f27, %[v]\n\t"
+ "fbcx.ps f28, %[v]\n\t"
+ "fbcx.ps f29, %[v]\n\t"
+ "fbcx.ps f30, %[v]\n\t"
+ "fbcx.ps f31, %[v]\n\t" ::[v] "r"(val)
+ : "f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "f9", "f10", "f11", "f12", "f13", "f14", "f15", "f16",
+ "f17", "f18", "f19", "f20", "f21", "f22", "f23", "f24", "f25", "f26", "f27", "f28", "f29", "f30", "f31");
+}
+
+// Fill a partial region [start, end) using tensor_store for 16-byte-aligned
+// chunks and byte stores for any remainder < 16 bytes.
+// Assumes f-regs are already loaded with the fill pattern.
+static void memset_tail(uint8_t * start, uint8_t * end, uint8_t val) {
+ uint8_t * cur = start;
+
+ // Full 64-byte rows via tensor_store (up to 16 at a time = 1KB)
+ while (cur + 64 <= end) {
+ size_t rows = (end - cur) / 64;
+ if (rows > 16) {
+ rows = 16;
+ }
+ tensor_store(0, 0, 3, rows - 1, (uintptr_t) cur, 0, 64);
+ cur += rows * 64;
+ }
+
+ // Remaining 16-byte aligned chunk (16, 32, or 48 bytes)
+ if (cur + 16 <= end) {
+ size_t cols = (end - cur) / 16;
+ tensor_store(0, 0, cols - 1, 0, (uintptr_t) cur, 0, 64);
+ cur += cols * 16;
+ }
+
+ tensor_wait(TENSOR_STORE_WAIT);
+
+ // Final < 16 bytes with byte stores
+ while (cur < end) {
+ *(volatile uint8_t *) cur = val;
+ cur++;
+ }
+}
+
+#define ALIGN_UP(ptr, align) ((uint8_t *) (((uintptr_t) (ptr) + (align) - 1) & ~((uintptr_t) (align) - 1)))
+
+int entry_point(struct memset_params * params, kernel_environment_t * env) {
+ uint64_t hart_id = get_hart_id();
+
+ // Only even harts have tensor engine access
+ if (hart_id & 1) {
+ return 0;
+ }
+
+ if (!params || ((uintptr_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ if (params->op_type != GGML_ET_MEMOP_MEMSET) {
+ return -1;
+ }
+
+ uint8_t * dst = (uint8_t *) params->dst_ptr;
+ size_t size = params->size;
+
+ if (!dst || size == 0) {
+ return -1;
+ }
+
+ // Dynamic hart count from shire_mask
+ int num_even_harts = manual_popcountll(env->shire_mask) * SOC_MINIONS_PER_SHIRE;
+
+ // global_id: shire * 32 + minion (for even harts)
+ uint64_t global_id = ((hart_id >> 6) << 5) + ((hart_id >> 1) & 0x1F);
+
+ uint8_t val = params->value & 0xFF;
+ uint32_t fill32 = val | ((uint32_t) val << 8) | ((uint32_t) val << 16) | ((uint32_t) val << 24);
+
+ uint8_t * end = dst + size;
+
+ setup_cache_scp();
+ CLEAR_TENSOR_ERROR;
+ fill_fregs(fill32);
+
+ // Align to 16 bytes (tensor_store minimum alignment)
+ uint8_t * base = ALIGN_UP(dst, 16);
+ if (base > end) {
+ base = end;
+ }
+
+ // Hart 0 handles head bytes before alignment
+ if (global_id == 0) {
+ volatile uint8_t * p = dst;
+ while (p < (volatile uint8_t *) base) {
+ *p++ = val;
+ }
+ }
+
+ // Bulk: 1KB blocks distributed across all harts (base is already 16-byte aligned)
+ size_t aligned_size = end - base;
+ size_t total_blocks = aligned_size / 1024;
+
+ if (total_blocks > 0) {
+ size_t blocks_per_hart = total_blocks / num_even_harts;
+ size_t extra = total_blocks % num_even_harts;
+ size_t my_start = blocks_per_hart * global_id + (global_id < extra ? global_id : extra);
+ size_t my_count = blocks_per_hart + (global_id < extra ? 1 : 0);
+
+ uint8_t * addr = base + my_start * 1024;
+ for (size_t b = 0; b < my_count; b++) {
+ tensor_store(0, 0, 3, 15, (uintptr_t) addr, 0, 64);
+ addr += 1024;
+ }
+ tensor_wait(TENSOR_STORE_WAIT);
+ }
+
+ // Hart 0 handles the tail after the last full 1KB block
+ if (global_id == 0) {
+ memset_tail(base + total_blocks * 1024, end, val);
+ }
+
+ FENCE;
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// MUL_MAT Kernel
+// Matrix multiplication: C[M,N] = A[M,K] * B[K,N]
+//******************************************************************************
+
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <stdint.h>
+
+#define STRIDE_M 2048 /* 32 shires x 32 minions x 2 harts */
+#define STRIDE_M_KSPLIT 1024 /* 32 shires x 32 minions (both harts share rows) */
+#define KSPLIT_MIN_K_BLOCKS 256 /* K >= 8192 elements */
+#define KSPLIT_SMALL_ROWS_K_BLOCKS 64 /* K >= 2048 elements for very small M */
+#define KSPLIT_MAX_ROWS 8 /* max rows per minion for K-split */
+#define TILE_KB 256 /* K-tile size in Q4_0 blocks (8192 elems, 32KB B data) */
+#define KSPLIT_GROUP_ROWS 4
+#define SIMPLE_X2_ROWS 2
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ uint64_t hart_id = get_hart_id();
+
+ // Matrix dimensions
+ const int64_t K = params->src0.ne[0];
+ const int64_t M = params->src0.ne[1];
+ const int64_t N = params->src1.ne[1];
+ const int64_t ne02 = params->src0.ne[2];
+ const int64_t ne03 = params->src0.ne[3];
+ const int64_t ne12 = params->src1.ne[2];
+ const int64_t ne13 = params->src1.ne[3];
+
+ // Strides (in bytes)
+ const size_t nb01 = params->src0.nb[1];
+ const size_t nb02 = params->src0.nb[2];
+ const size_t nb03 = params->src0.nb[3];
+
+ const size_t nb11 = params->src1.nb[1];
+ const size_t nb12 = params->src1.nb[2];
+ const size_t nb13 = params->src1.nb[3];
+
+ const size_t nbd1 = params->dst.nb[1];
+ const size_t nbd2 = params->dst.nb[2];
+ const size_t nbd3 = params->dst.nb[3];
+
+ // Q4_0 block size is 32
+ const int64_t K_blocks = K / 32;
+ const int use_simple_x2 = ((nb01 & 31) == 0);
+
+ // Broadcasting ratios
+ const int64_t r2 = ne12 / ne02;
+ const int64_t r3 = ne13 / ne03;
+
+ // K-split decision
+ const int64_t minion_id = hart_id >> 1; /* 0..1023 global */
+ const int64_t local_minion = (hart_id >> 1) & 0x1F; /* 0..31 within shire */
+ const int is_hart1 = hart_id & 1;
+ const int64_t rows_per_minion = (M + STRIDE_M_KSPLIT - 1) / STRIDE_M_KSPLIT;
+ const int64_t k_half = K_blocks / 2;
+ const int use_ksplit_small_rows = (rows_per_minion <= 2) && (K_blocks >= KSPLIT_SMALL_ROWS_K_BLOCKS);
+ /*
+ * K-split when K is large enough to benefit, and either:
+ * - few rows (≤4): always safe, proven working
+ * - more rows (5-8): only if each hart's half fits in one tile,
+ * otherwise L1 thrashing from 2 harts × 8 rows kills performance
+ *
+ * Also allow K-split earlier for the low-M regime (≤2 rows/minion). In
+ * that case the simple row-striped path leaves half the machine idle, so
+ * using both harts on each row pays off even for moderate K.
+ */
+ const int use_ksplit = ((K_blocks >= KSPLIT_MIN_K_BLOCKS) && (rows_per_minion <= KSPLIT_MAX_ROWS) &&
+ (rows_per_minion <= 4 || k_half <= TILE_KB)) ||
+ use_ksplit_small_rows;
+ const int use_ksplit_group = !use_ksplit && (K_blocks >= KSPLIT_MIN_K_BLOCKS) && (rows_per_minion > 4) &&
+ (rows_per_minion <= KSPLIT_MAX_ROWS);
+
+ if (use_ksplit) {
+ /* Each hart processes half the K dimension */
+ const int64_t k_start = is_hart1 ? k_half : 0;
+ const int64_t k_len = is_hart1 ? (K_blocks - k_half) : k_half;
+
+ /* One cache-line-aligned L2SCP slot per minion for exchange */
+ volatile float * l2scp_slot = (volatile float *) et_shire_l2scp_local(local_minion * 64);
+
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+
+ for (int64_t m = minion_id; m < M; m += STRIDE_M_KSPLIT) {
+ const block_q4_0 * q_row = (const block_q4_0 *) (src0_ptr2 + m * nb01);
+
+ float partial = compute_row_dot_q4_0(q_row + k_start, b_col_base + k_start * 32, k_len);
+
+ if (is_hart1) {
+ *l2scp_slot = partial;
+ FENCE;
+ flush_to_l2((const void *) l2scp_slot, 1, 64);
+ WAIT_CACHEOPS;
+ et_sem_post(ET_BARRIER_MINION);
+ et_sem_wait(ET_BARRIER_MINION);
+ } else {
+ et_sem_wait(ET_BARRIER_MINION);
+ float other = *l2scp_slot;
+ et_sem_post(ET_BARRIER_MINION);
+
+ float * dst_entry = (float *) (dst_ptr2 + n * nbd1 + m * sizeof(float));
+ atomic_store_f32((volatile float *) dst_entry, partial + other);
+ }
+ }
+ }
+ }
+ }
+ } else if (use_ksplit_group) {
+ /*
+ * Grouped K-split for the 5-8 rows/minion regime.
+ *
+ * Both harts process the same 4-row group, each on half of K, and
+ * exchange 4 partial sums once per group instead of once per row.
+ * This keeps the K-split bandwidth benefit while cutting semaphore
+ * traffic by 4x relative to the old per-row exchange.
+ */
+ const int64_t k_start = is_hart1 ? k_half : 0;
+ const int64_t k_len = is_hart1 ? (K_blocks - k_half) : k_half;
+ volatile float * l2scp_slot = (volatile float *) et_shire_l2scp_local(local_minion * 64);
+
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+
+ for (int64_t m_base = minion_id; m_base < M; m_base += STRIDE_M_KSPLIT * KSPLIT_GROUP_ROWS) {
+ const int64_t m0 = m_base;
+ const int64_t m1 = m0 + STRIDE_M_KSPLIT;
+ const int64_t m2 = m1 + STRIDE_M_KSPLIT;
+ const int64_t m3 = m2 + STRIDE_M_KSPLIT;
+
+ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
+
+ for (int64_t kb = 0; kb < K_blocks; kb += TILE_KB) {
+ int64_t tile_len = k_len - kb;
+ if (tile_len > TILE_KB) {
+ tile_len = TILE_KB;
+ }
+ if (tile_len <= 0) {
+ break;
+ }
+ const float * b_tile = b_col_base + (k_start + kb) * 32;
+ const int64_t row_kb = k_start + kb;
+
+ if (m0 < M) {
+ s0 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m0 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ if (m1 < M) {
+ s1 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m1 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ if (m2 < M) {
+ s2 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m2 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ if (m3 < M) {
+ s3 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m3 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ }
+
+ if (is_hart1) {
+ l2scp_slot[0] = s0;
+ l2scp_slot[1] = s1;
+ l2scp_slot[2] = s2;
+ l2scp_slot[3] = s3;
+ FENCE;
+ flush_to_l2((const void *) l2scp_slot, 1, 64);
+ WAIT_CACHEOPS;
+ et_sem_post(ET_BARRIER_MINION);
+ et_sem_wait(ET_BARRIER_MINION);
+ } else {
+ et_sem_wait(ET_BARRIER_MINION);
+ const float p0 = l2scp_slot[0];
+ const float p1 = l2scp_slot[1];
+ const float p2 = l2scp_slot[2];
+ const float p3 = l2scp_slot[3];
+ et_sem_post(ET_BARRIER_MINION);
+
+ float * c_base = (float *) (dst_ptr2 + n * nbd1);
+ if (m0 < M) {
+ atomic_store_f32((volatile float *) (c_base + m0), s0 + p0);
+ }
+ if (m1 < M) {
+ atomic_store_f32((volatile float *) (c_base + m1), s1 + p1);
+ }
+ if (m2 < M) {
+ atomic_store_f32((volatile float *) (c_base + m2), s2 + p2);
+ }
+ if (m3 < M) {
+ atomic_store_f32((volatile float *) (c_base + m3), s3 + p3);
+ }
+ }
+ }
+ }
+ }
+ }
+ } else if (K_blocks > TILE_KB) {
+ /*
+ * Tile-outer with scalar row groups: process up to 4 rows per
+ * hart sharing each B tile before advancing to the next tile.
+ * Uses scalar float variables (not an array) to accumulate across
+ * tiles — avoids the flw/fadd.s/fsw stack ops that corrupt vector
+ * register state on ET-SoC-1's MMX-style shared FP file.
+ */
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+
+ for (int64_t m0 = hart_id; m0 < M; m0 += STRIDE_M * 4) {
+ const int64_t m1 = m0 + STRIDE_M;
+ const int64_t m2 = m0 + STRIDE_M * 2;
+ const int64_t m3 = m0 + STRIDE_M * 3;
+
+ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
+
+ for (int64_t kb = 0; kb < K_blocks; kb += TILE_KB) {
+ int64_t tile_len = K_blocks - kb;
+ if (tile_len > TILE_KB) {
+ tile_len = TILE_KB;
+ }
+ const float * b_tile = b_col_base + kb * 32;
+
+ s0 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m0 * nb01) + kb, b_tile,
+ tile_len);
+ if (m1 < M) {
+ s1 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m1 * nb01) + kb, b_tile,
+ tile_len);
+ }
+ if (m2 < M) {
+ s2 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m2 * nb01) + kb, b_tile,
+ tile_len);
+ }
+ if (m3 < M) {
+ s3 += compute_row_dot_q4_0((const block_q4_0 *) (src0_ptr2 + m3 * nb01) + kb, b_tile,
+ tile_len);
+ }
+ }
+
+ float * dst_base = (float *) (dst_ptr2 + n * nbd1);
+ atomic_store_f32((volatile float *) (dst_base + m0), s0);
+ if (m1 < M) {
+ atomic_store_f32((volatile float *) (dst_base + m1), s1);
+ }
+ if (m2 < M) {
+ atomic_store_f32((volatile float *) (dst_base + m2), s2);
+ }
+ if (m3 < M) {
+ atomic_store_f32((volatile float *) (dst_base + m3), s3);
+ }
+ }
+ }
+ }
+ }
+ } else {
+ /*
+ * Simple path for small K.
+ *
+ * When `nb01` is 32-byte aligned, every row has the same block-alignment
+ * pattern. That lets us compute two rows together and reuse each loaded
+ * B chunk across both rows instead of reloading it in a second dot call.
+ */
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+ q4_dot_state q4_state;
+ q4_dot_begin(&q4_state);
+
+ if (use_simple_x2) {
+ for (int64_t m0 = hart_id; m0 < M; m0 += STRIDE_M * SIMPLE_X2_ROWS) {
+ const int64_t m1 = m0 + STRIDE_M;
+ const block_q4_0 * q_row0 = (const block_q4_0 *) (src0_ptr2 + m0 * nb01);
+
+ if (m1 < M) {
+ const block_q4_0 * q_row1 = (const block_q4_0 *) (src0_ptr2 + m1 * nb01);
+ float s0, s1;
+ q4_dot_compute_x2_aligned(q_row0, q_row1, b_col_base, K_blocks, &s0, &s1);
+
+ float * dst0 = (float *) (dst_ptr2 + n * nbd1 + m0 * sizeof(float));
+ float * dst1 = (float *) (dst_ptr2 + n * nbd1 + m1 * sizeof(float));
+ atomic_store_f32((volatile float *) dst0, s0);
+ atomic_store_f32((volatile float *) dst1, s1);
+ } else {
+ float sum = q4_dot_compute(q_row0, b_col_base, K_blocks);
+ float * dst = (float *) (dst_ptr2 + n * nbd1 + m0 * sizeof(float));
+ atomic_store_f32((volatile float *) dst, sum);
+ }
+ }
+ } else {
+ for (int64_t m = hart_id; m < M; m += STRIDE_M) {
+ const block_q4_0 * q_row = (const block_q4_0 *) (src0_ptr2 + m * nb01);
+
+ float sum = q4_dot_compute(q_row, b_col_base, K_blocks);
+
+ float * dst_entry = (float *) (dst_ptr2 + n * nbd1 + m * sizeof(float));
+ atomic_store_f32((volatile float *) dst_entry, sum);
+ }
+ }
+
+ q4_dot_end(&q4_state);
+ }
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+#include "tensor.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+
+// Q4_0 x F32 -> F32 MUL_MAT on the tensor (matrix) engine, TensorFMA32.
+// Hart 1: dequantize Q4_0 weights to FP32 into double-buffered L2 SCP.
+// Hart 0: tensor engine compute (FMA, reduce, store).
+
+#define NUM_COMPUTE_SHIRES 32
+#define MINIONS_PER_SHIRE 32
+
+#define TILE_M 16
+#define TILE_N 16
+#define BLOCK_K QK4_0 // 32 elements per Q4_0 block
+#define FMA_K 16 // tensor FMA k-width for FP32 (a_num_cols = FMA_K-1)
+
+#define CACHEOP_MAX 0
+#define REP_RATE 0
+
+#define A_L1_START 0 // L1 SCP lines 0..15 for A (activations)
+#define B_L1_START 16 // L1 SCP lines 16..31 for B (dequantized weights)
+
+// L2 SCP layout per minion (double-buffered dequant panel + sync counters).
+// panel = BLOCK_K k-lines x TILE_M m (FP32) = 32 * 64 = 2048 bytes, in TenB
+// [k][m] order: panel[k*TILE_M + m].
+#define SCP_PANEL_SIZE (BLOCK_K * TILE_M * (uint64_t) sizeof(float)) // 2048
+#define SCP_READY_OFF (2 * SCP_PANEL_SIZE) // 4096
+#define SCP_CONSUMED_OFF (SCP_READY_OFF + 64) // 4160
+#define SCP_PER_MINION (SCP_CONSUMED_OFF + 64) // 4224
+
+// Signal a counter value to the other hart via L2 SCP.
+static inline void __attribute__((always_inline)) scp_signal(volatile uint32_t * flag, uint32_t value) {
+ *flag = value;
+ FENCE;
+ evict_to_l2((const void *) flag, 1, 64);
+ WAIT_CACHEOPS;
+}
+
+// Wait for a counter in L2 SCP to reach the expected value.
+static inline void __attribute__((always_inline)) scp_wait(volatile uint32_t * flag, uint32_t expected) {
+ while (1) {
+ evict_to_l2((const void *) flag, 1, 64);
+ WAIT_CACHEOPS;
+ if (*flag >= expected) {
+ return;
+ }
+ }
+}
+
+// Dequantize one 32-element Q4_0 block of TILE_M weight rows into the FP32
+// panel, written directly in TenB [k][m] order: panel[k*TILE_M + m].
+// Low nibble of byte i -> k = i
+// High nibble of byte i -> k = i + 16
+// value = d * (nibble - 8)
+//
+// Vectorized: for each weight row m we gather 8 packed bytes at a time, expand
+// the low/high nibbles to FP32 (nibble-8), scale by the block's fp16 d, and
+// fscw.ps-scatter the 8 values down 8 panel lines (stride 64B) at column m.
+// 4 groups of 8 cover the 32 k-values (low 0..15, high 16..31).
+static inline void __attribute__((always_inline)) dequant_q4_0_panel(float * panel,
+ const char * src0_batch,
+ int64_t mb,
+ int64_t kb_block,
+ int64_t nb1_0) {
+ static const int32_t __attribute__((aligned(32))) scatter_idx[8] = {
+ 0, 64, 128, 192, 256, 320, 384, 448 // byte offsets: 8 lines apart
+ };
+ static const int32_t __attribute__((aligned(32))) gather_idx[8] = {
+ 0, 1, 2, 3, 4, 5, 6, 7 // 8 consecutive bytes
+ };
+
+ unsigned long old_mask;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t" // all 8 lanes active
+ "flw.ps f1, (%[sidx]) \n\t" // f1 = scatter offsets
+ "flw.ps f2, (%[gidx]) \n\t" // f2 = gather offsets
+ : [ms] "=&r"(old_mask)
+ : [sidx] "r"(scatter_idx), [gidx] "r"(gather_idx)
+ : "f1", "f2");
+
+ char * pbase = (char *) panel;
+ for (int j = 0; j < TILE_M; ++j) {
+ const block_q4_0 * blk = (const block_q4_0 *) (src0_batch + (mb + j) * nb1_0) + kb_block;
+ uint32_t scale_raw = (uint32_t) blk->d;
+ const uint8_t * qs = blk->qs;
+ char * col = pbase + j * 4; // column m=j of the panel
+
+ __asm__ volatile(
+ "fbcx.ps f3, %[sb] \n\t" // broadcast fp16 scale bits
+ "fcvt.ps.f16 f3, f3 \n\t" // -> d in all 8 lanes (fp32)
+
+ "fgb.ps f4, f2(%[qs0]) \n\t" // gather qs[0..7]
+ "fandi.pi f5, f4, 15 \n\t" // low nibble
+ "faddi.pi f5, f5, -8 \n\t"
+ "fcvt.ps.pw f5, f5, rne \n\t"
+ "fmul.ps f5, f5, f3 \n\t"
+ "fscw.ps f5, f1(%[c0]) \n\t" // k=0..7 -> lines 0..7
+ "fsrli.pi f6, f4, 4 \n\t" // high nibble
+ "fandi.pi f6, f6, 15 \n\t"
+ "faddi.pi f6, f6, -8 \n\t"
+ "fcvt.ps.pw f6, f6, rne \n\t"
+ "fmul.ps f6, f6, f3 \n\t"
+ "fscw.ps f6, f1(%[c16]) \n\t" // k=16..23 -> lines 16..23
+
+ "fgb.ps f4, f2(%[qs8]) \n\t" // gather qs[8..15]
+ "fandi.pi f5, f4, 15 \n\t"
+ "faddi.pi f5, f5, -8 \n\t"
+ "fcvt.ps.pw f5, f5, rne \n\t"
+ "fmul.ps f5, f5, f3 \n\t"
+ "fscw.ps f5, f1(%[c8]) \n\t" // k=8..15 -> lines 8..15
+ "fsrli.pi f6, f4, 4 \n\t"
+ "fandi.pi f6, f6, 15 \n\t"
+ "faddi.pi f6, f6, -8 \n\t"
+ "fcvt.ps.pw f6, f6, rne \n\t"
+ "fmul.ps f6, f6, f3 \n\t"
+ "fscw.ps f6, f1(%[c24]) \n\t" // k=24..31 -> lines 24..31
+ :
+ : [sb] "r"(scale_raw), [qs0] "r"(qs), [qs8] "r"(qs + 8), [c0] "r"(col), [c8] "r"(col + 8 * 64),
+ [c16] "r"(col + 16 * 64), [c24] "r"(col + 24 * 64)
+ : "f3", "f4", "f5", "f6", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(old_mask));
+}
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ (void) env;
+
+ uint64_t hart_id = get_hart_id();
+ uint64_t shire_id = get_shire_id();
+
+ if (shire_id >= NUM_COMPUTE_SHIRES) {
+ return 0;
+ }
+
+ const int is_hart1 = hart_id & 1;
+ uint64_t local_minion = (hart_id >> 1) & 0x1F;
+
+ // Dimensions (both harts need these for tile assignment)
+ const int64_t K = params->src0.ne[0];
+ const int64_t M = params->src0.ne[1];
+ const int64_t N = params->src1.ne[1];
+
+ if ((M % TILE_M) != 0) {
+ return 0;
+ }
+ if ((K % BLOCK_K) != 0) {
+ return 0;
+ }
+
+ const int64_t ne2_0 = params->src0.ne[2], ne3_0 = params->src0.ne[3];
+ const int64_t ne2_1 = params->src1.ne[2], ne3_1 = params->src1.ne[3];
+
+ const int64_t nb1_0 = params->src0.nb[1];
+ const int64_t nb2_0 = params->src0.nb[2], nb3_0 = params->src0.nb[3];
+
+ const int64_t nb1_1 = params->src1.nb[1];
+ const int64_t nb2_1 = params->src1.nb[2], nb3_1 = params->src1.nb[3];
+
+ const int64_t nb1_d = params->dst.nb[1];
+ const int64_t nb2_d = params->dst.nb[2], nb3_d = params->dst.nb[3];
+
+ const char * src0_base = (const char *) params->src0.data;
+ const char * src1_base = (const char *) params->src1.data;
+ char * dst_base = (char *) params->dst.data;
+
+ const int64_t m_tiles = M / TILE_M;
+ const int64_t n_tiles = (N + TILE_N - 1) / TILE_N;
+ const int64_t batch_count = ne2_1 * ne3_1;
+ const int64_t base_tiles = m_tiles * n_tiles * batch_count;
+
+ const int64_t r2 = ne2_1 / ne2_0;
+ const int64_t r3 = ne3_1 / ne3_0;
+
+ const int64_t k_steps = K / BLOCK_K; // number of Q4_0 blocks
+
+ // Force a single K-split.
+ const int64_t k_splits = 1;
+
+ const int64_t tiles_per_shire = MINIONS_PER_SHIRE / k_splits;
+ const int64_t k_split = local_minion % k_splits;
+ const int64_t local_tile_idx = local_minion / k_splits;
+ const int64_t tiles_stride = (int64_t) NUM_COMPUTE_SHIRES * tiles_per_shire;
+
+ const int64_t k_steps_per_split = k_steps / k_splits;
+ const int64_t kb_start = k_split * k_steps_per_split; // first block
+ const int64_t kb_end = kb_start + k_steps_per_split; // one past last
+
+ // L2 SCP pointers for this minion's double-buffered panels + sync.
+ uint64_t scp_base = local_minion * SCP_PER_MINION;
+ float * scp_panel[2] = {
+ (float *) et_shire_l2scp_local(scp_base),
+ (float *) et_shire_l2scp_local(scp_base + SCP_PANEL_SIZE),
+ };
+ volatile uint32_t * ready_ctr = (volatile uint32_t *) et_shire_l2scp_local(scp_base + SCP_READY_OFF);
+ volatile uint32_t * consumed_ctr = (volatile uint32_t *) et_shire_l2scp_local(scp_base + SCP_CONSUMED_OFF);
+
+ // ================================================================
+ // Hart 1: Q4_0 weight dequant producer
+ // ================================================================
+ if (is_hart1) {
+ scp_signal(ready_ctr, 0);
+ scp_signal(consumed_ctr, 0);
+
+ uint32_t chunk_id = 0;
+
+ for (int64_t tile = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES; tile < base_tiles;
+ tile += tiles_stride) {
+ const int64_t tiles_per_batch = m_tiles * n_tiles;
+ const int64_t batch_idx = tile / tiles_per_batch;
+ const int64_t tile_in_batch = tile % tiles_per_batch;
+
+ const int64_t mb_idx = tile_in_batch % m_tiles;
+
+ const int64_t i3 = batch_idx / ne2_1;
+ const int64_t i2 = batch_idx % ne2_1;
+ const int64_t i2_0 = i2 / r2;
+ const int64_t i3_0 = i3 / r3;
+
+ const char * src0_batch = src0_base + i3_0 * nb3_0 + i2_0 * nb2_0;
+ const int64_t mb = mb_idx * TILE_M;
+
+ for (int64_t kb = kb_start; kb < kb_end; ++kb) {
+ int buf = chunk_id & 1;
+
+ // Back-pressure: wait for hart 0 to finish with this buffer.
+ if (chunk_id >= 2) {
+ scp_wait(consumed_ctr, chunk_id - 1);
+ }
+
+ dequant_q4_0_panel(scp_panel[buf], src0_batch, mb, kb, nb1_0);
+
+ FENCE;
+ flush_to_l2(scp_panel[buf], BLOCK_K, 64);
+ WAIT_CACHEOPS;
+
+ chunk_id++;
+ scp_signal(ready_ctr, chunk_id);
+ }
+ }
+
+ FENCE;
+ return 0;
+ }
+
+ // ================================================================
+ // Hart 0: tensor engine compute
+ // ================================================================
+ uint64_t my_minion_id = get_minion_id();
+ const uint64_t group_base_global = my_minion_id - k_split;
+
+ setup_cache_scp();
+#if CACHEOP_MAX > 0 || REP_RATE > 0
+ ucache_control(1, REP_RATE, CACHEOP_MAX);
+#endif
+ CLEAR_TENSOR_ERROR;
+
+ evict_to_l2((const void *) ready_ctr, 1, 64);
+ WAIT_CACHEOPS;
+ evict_to_l2((const void *) consumed_ctr, 1, 64);
+ WAIT_CACHEOPS;
+
+ uint32_t chunk_id = 0;
+
+ for (int64_t tile = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES; tile < base_tiles;
+ tile += tiles_stride) {
+ const int64_t tiles_per_batch = m_tiles * n_tiles;
+ const int64_t batch_idx = tile / tiles_per_batch;
+ const int64_t tile_in_batch = tile % tiles_per_batch;
+
+ const int64_t nb_idx = tile_in_batch / m_tiles;
+ const int64_t mb_idx = tile_in_batch % m_tiles;
+
+ const int64_t i3 = batch_idx / ne2_1;
+ const int64_t i2 = batch_idx % ne2_1;
+
+ const char * src1_batch = src1_base + i3 * nb3_1 + i2 * nb2_1;
+ char * dst_batch = dst_base + i3 * nb3_d + i2 * nb2_d;
+
+ const int64_t mb = mb_idx * TILE_M;
+ const int64_t nb = nb_idx * TILE_N;
+ const int64_t n_cur = (nb + TILE_N <= N) ? TILE_N : (N - nb);
+
+ // Partial-N tiles run TensorFMA32 with a_num_rows = n_cur-1.
+ // Errata Type D workaround for n_cur == 4 (AROWS==3): pad A to AROWS==4.
+ const int64_t arows_fma = (n_cur == 4) ? 4 : (n_cur - 1);
+
+ if (n_cur == 4) {
+ // Zero the padded 5th A row (line A_L1_START+4) once; the per-pass A
+ // load only writes lines A_L1_START..+3, so this persists.
+ static const float __attribute__((aligned(64))) zero_line[16] = { 0 };
+ tensor_load(false, false, A_L1_START + 4, TENSOR_LOAD_PLAIN, 0, (uint64_t) zero_line, 0,
+ 0, // 1 line
+ 64, 0);
+ tensor_wait(TENSOR_LOAD_WAIT_0);
+ }
+
+ int first = 1; // first_pass=1 only for the very first FMA of the tile
+
+ for (int64_t kb = kb_start; kb < kb_end; ++kb) {
+ int buf = chunk_id & 1;
+
+ // Wait for hart 1 to finish dequantizing this block.
+ chunk_id++;
+ scp_wait(ready_ctr, chunk_id);
+
+ // Two FMA passes over the 32-wide block (16 K-cols each).
+ for (int half = 0; half < 2; ++half) {
+ const int64_t k_elem = kb * BLOCK_K + half * FMA_K;
+
+ // Load A (activations) for this 16-K sub-tile, PLAIN.
+ tensor_load(false, false, A_L1_START, TENSOR_LOAD_PLAIN, 0,
+ (uint64_t) (src1_batch + nb * nb1_1 + k_elem * (int64_t) sizeof(float)), 0, n_cur - 1,
+ (uint64_t) nb1_1, 0);
+
+ // Load B (dequantized weights) half from L2 SCP panel, PLAIN.
+ tensor_load(false, false, B_L1_START, TENSOR_LOAD_PLAIN, 0,
+ (uint64_t) (scp_panel[buf] + (int64_t) half * FMA_K * TILE_M), 0, FMA_K - 1, 64, 1);
+
+ tensor_wait(TENSOR_LOAD_WAIT_0);
+ tensor_wait(TENSOR_LOAD_WAIT_1);
+
+ tensor_fma(false,
+ 3, // b_num_col: (16/4)-1
+ arows_fma, // a_num_rows (n_cur-1, or 4 for the n_cur==4 errata pad)
+ FMA_K - 1, // a_num_cols
+ 0, false, false, false, false, B_L1_START, A_L1_START, TENSOR_FMA_OP_FP32, first);
+
+ tensor_wait(TENSOR_FMA_WAIT);
+ first = 0;
+ }
+
+ // Signal that this buffer is free for hart 1 to reuse.
+ scp_signal(consumed_ctr, chunk_id);
+ }
+
+ // K-split ring reduce.
+ if (k_splits > 1) {
+ const uint64_t num_regs = (uint64_t) n_cur * 2;
+
+ if (k_split > 0) {
+ tensor_reduce_recv(0, TENSOR_REDUCE_OP_FADD, num_regs, group_base_global + k_split - 1);
+ tensor_wait(TENSOR_REDUCE_WAIT);
+ }
+
+ if (k_split < k_splits - 1) {
+ tensor_reduce_send(0, num_regs, group_base_global + k_split + 1);
+ tensor_wait(TENSOR_REDUCE_WAIT);
+ }
+ }
+
+ // Store FP32 result tile (only the last k-split owns the final sum).
+ if (k_split == k_splits - 1) {
+ tensor_store(0, 0, 3, n_cur - 1, (uint64_t) (dst_batch + nb * nb1_d + mb * (int64_t) sizeof(float)), 0,
+ (uint64_t) nb1_d);
+ tensor_wait(TENSOR_STORE_WAIT);
+ }
+ }
+
+ FENCE;
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// MUL_MAT Kernel
+// Matrix multiplication: C[M,N] = A[M,K] * B[K,N]
+//******************************************************************************
+
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <stdint.h>
+
+#define STRIDE_M 2048 /* 32 shires x 32 minions x 2 harts */
+#define STRIDE_M_KSPLIT 1024 /* 32 shires x 32 minions (both harts share rows) */
+#define KSPLIT_MIN_K_BLOCKS 256 /* K >= 8192 elements */
+#define KSPLIT_SMALL_ROWS_K_BLOCKS 64 /* K >= 2048 elements for very small M */
+#define KSPLIT_MAX_ROWS 8 /* max rows per minion for K-split */
+#define TILE_KB 256 /* K-tile size in Q8_0 blocks (8192 elems, 32KB B data) */
+#define KSPLIT_GROUP_ROWS 4
+#define SIMPLE_X2_ROWS 2
+
+static inline size_t tensor_bytes(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+int entry_point(struct ggml_et_mm_q8_params * params, void * env) {
+ uint64_t hart_id = get_hart_id();
+
+ // Matrix dimensions
+ const int64_t K = params->src0.ne[0];
+ const int64_t M = params->src0.ne[1];
+ const int64_t N = params->src1.ne[1];
+ const int64_t ne02 = params->src0.ne[2];
+ const int64_t ne03 = params->src0.ne[3];
+ const int64_t ne12 = params->src1.ne[2];
+ const int64_t ne13 = params->src1.ne[3];
+
+ // Strides (in bytes)
+ const size_t nb01 = params->src0.nb[1];
+ const size_t nb02 = params->src0.nb[2];
+ const size_t nb03 = params->src0.nb[3];
+
+ const size_t nb11 = params->src1.nb[1];
+ const size_t nb12 = params->src1.nb[2];
+ const size_t nb13 = params->src1.nb[3];
+
+ const size_t nbd1 = params->dst.nb[1];
+ const size_t nbd2 = params->dst.nb[2];
+ const size_t nbd3 = params->dst.nb[3];
+
+ // Optional residual bias
+ const char * bias_base = (const char *) params->bias.data;
+ const size_t nbb1 = params->bias.nb[1];
+ const size_t nbb2 = params->bias.nb[2];
+ const size_t nbb3 = params->bias.nb[3];
+
+ // Q8_0 block size is 32
+ const int64_t K_blocks = K / 32;
+ const int use_simple_x2 = ((nb01 & 31) == 0);
+
+ // Broadcasting ratios
+ const int64_t r2 = ne12 / ne02;
+ const int64_t r3 = ne13 / ne03;
+
+ // K-split decision
+ const int64_t minion_id = hart_id >> 1; /* 0..1023 global */
+ const int64_t local_minion = (hart_id >> 1) & 0x1F; /* 0..31 within shire */
+ const int is_hart1 = hart_id & 1;
+ const int64_t rows_per_minion = (M + STRIDE_M_KSPLIT - 1) / STRIDE_M_KSPLIT;
+ const int64_t k_half = K_blocks / 2;
+ const int use_ksplit_small_rows = (rows_per_minion <= 2) && (K_blocks >= KSPLIT_SMALL_ROWS_K_BLOCKS);
+ /*
+ * K-split when K is large enough to benefit, and either:
+ * - few rows (≤4): always safe, proven working
+ * - more rows (5-8): only if each hart's half fits in one tile,
+ * otherwise L1 thrashing from 2 harts × 8 rows kills performance
+ *
+ * Also allow K-split earlier for the low-M regime (≤2 rows/minion). In
+ * that case the simple row-striped path leaves half the machine idle, so
+ * using both harts on each row pays off even for moderate K.
+ */
+ const int use_ksplit = ((K_blocks >= KSPLIT_MIN_K_BLOCKS) && (rows_per_minion <= KSPLIT_MAX_ROWS) &&
+ (rows_per_minion <= 4 || k_half <= TILE_KB)) ||
+ use_ksplit_small_rows;
+ const int use_ksplit_group = !use_ksplit && (K_blocks >= KSPLIT_MIN_K_BLOCKS) && (rows_per_minion > 4) &&
+ (rows_per_minion <= KSPLIT_MAX_ROWS);
+
+ evict_region_past_l2(params->src1.data, tensor_bytes(¶ms->src1));
+ if (params->bias.data) {
+ evict_region_past_l2(params->bias.data, tensor_bytes(¶ms->bias));
+ }
+
+ if (use_ksplit) {
+ /* Each hart processes half the K dimension */
+ const int64_t k_start = is_hart1 ? k_half : 0;
+ const int64_t k_len = is_hart1 ? (K_blocks - k_half) : k_half;
+
+ /* One cache-line-aligned L2SCP slot per minion for exchange */
+ volatile float * l2scp_slot = (volatile float *) et_shire_l2scp_local(local_minion * 64);
+
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+ const char * bias_ptr3 = bias_base ? bias_base + i3 * nbb3 : (const char *) 0;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+ const char * bias_ptr2 = bias_ptr3 ? bias_ptr3 + i2 * nbb2 : (const char *) 0;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+ const float * bias_n = bias_ptr2 ? (const float *) (bias_ptr2 + n * nbb1) : (const float *) 0;
+
+ for (int64_t m = minion_id; m < M; m += STRIDE_M_KSPLIT) {
+ const block_q8_0 * q_row = (const block_q8_0 *) (src0_ptr2 + m * nb01);
+
+ float partial = compute_row_dot_q8_0(q_row + k_start, b_col_base + k_start * 32, k_len);
+
+ if (is_hart1) {
+ *l2scp_slot = partial;
+ FENCE;
+ flush_to_l2((const void *) l2scp_slot, 1, 64);
+ WAIT_CACHEOPS;
+ et_sem_post(ET_BARRIER_MINION);
+ et_sem_wait(ET_BARRIER_MINION);
+ } else {
+ et_sem_wait(ET_BARRIER_MINION);
+ float other = *l2scp_slot;
+ et_sem_post(ET_BARRIER_MINION);
+
+ float * dst_entry = (float *) (dst_ptr2 + n * nbd1 + m * sizeof(float));
+ float sum = partial + other;
+ if (bias_n) {
+ sum += bias_n[m];
+ }
+ atomic_store_f32((volatile float *) dst_entry, sum);
+ }
+ }
+ }
+ }
+ }
+ } else if (use_ksplit_group) {
+ /*
+ * Grouped K-split for the 5-8 rows/minion regime.
+ *
+ * Both harts process the same 4-row group, each on half of K, and
+ * exchange 4 partial sums once per group instead of once per row.
+ * This keeps the K-split bandwidth benefit while cutting semaphore
+ * traffic by 4x relative to the old per-row exchange.
+ */
+ const int64_t k_start = is_hart1 ? k_half : 0;
+ const int64_t k_len = is_hart1 ? (K_blocks - k_half) : k_half;
+ volatile float * l2scp_slot = (volatile float *) et_shire_l2scp_local(local_minion * 64);
+
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+ const char * bias_ptr3 = bias_base ? bias_base + i3 * nbb3 : (const char *) 0;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+ const char * bias_ptr2 = bias_ptr3 ? bias_ptr3 + i2 * nbb2 : (const char *) 0;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+ const float * bias_n = bias_ptr2 ? (const float *) (bias_ptr2 + n * nbb1) : (const float *) 0;
+
+ for (int64_t m_base = minion_id; m_base < M; m_base += STRIDE_M_KSPLIT * KSPLIT_GROUP_ROWS) {
+ const int64_t m0 = m_base;
+ const int64_t m1 = m0 + STRIDE_M_KSPLIT;
+ const int64_t m2 = m1 + STRIDE_M_KSPLIT;
+ const int64_t m3 = m2 + STRIDE_M_KSPLIT;
+
+ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
+
+ for (int64_t kb = 0; kb < K_blocks; kb += TILE_KB) {
+ int64_t tile_len = k_len - kb;
+ if (tile_len > TILE_KB) {
+ tile_len = TILE_KB;
+ }
+ if (tile_len <= 0) {
+ break;
+ }
+ const float * b_tile = b_col_base + (k_start + kb) * 32;
+ const int64_t row_kb = k_start + kb;
+
+ if (m0 < M) {
+ s0 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m0 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ if (m1 < M) {
+ s1 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m1 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ if (m2 < M) {
+ s2 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m2 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ if (m3 < M) {
+ s3 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m3 * nb01) + row_kb,
+ b_tile, tile_len);
+ }
+ }
+
+ if (is_hart1) {
+ l2scp_slot[0] = s0;
+ l2scp_slot[1] = s1;
+ l2scp_slot[2] = s2;
+ l2scp_slot[3] = s3;
+ FENCE;
+ flush_to_l2((const void *) l2scp_slot, 1, 64);
+ WAIT_CACHEOPS;
+ et_sem_post(ET_BARRIER_MINION);
+ et_sem_wait(ET_BARRIER_MINION);
+ } else {
+ et_sem_wait(ET_BARRIER_MINION);
+ const float p0 = l2scp_slot[0];
+ const float p1 = l2scp_slot[1];
+ const float p2 = l2scp_slot[2];
+ const float p3 = l2scp_slot[3];
+ et_sem_post(ET_BARRIER_MINION);
+
+ float * c_base = (float *) (dst_ptr2 + n * nbd1);
+ const float b0 = bias_n ? bias_n[m0] : 0.0f;
+ const float b1 = (bias_n && m1 < M) ? bias_n[m1] : 0.0f;
+ const float b2 = (bias_n && m2 < M) ? bias_n[m2] : 0.0f;
+ const float b3 = (bias_n && m3 < M) ? bias_n[m3] : 0.0f;
+ if (m0 < M) {
+ atomic_store_f32((volatile float *) (c_base + m0), s0 + p0 + b0);
+ }
+ if (m1 < M) {
+ atomic_store_f32((volatile float *) (c_base + m1), s1 + p1 + b1);
+ }
+ if (m2 < M) {
+ atomic_store_f32((volatile float *) (c_base + m2), s2 + p2 + b2);
+ }
+ if (m3 < M) {
+ atomic_store_f32((volatile float *) (c_base + m3), s3 + p3 + b3);
+ }
+ }
+ }
+ }
+ }
+ }
+ } else if (K_blocks > TILE_KB) {
+ /*
+ * Tile-outer with scalar row groups: process up to 4 rows per
+ * hart sharing each B tile before advancing to the next tile.
+ * Uses scalar float variables (not an array) to accumulate across
+ * tiles — avoids the flw/fadd.s/fsw stack ops that corrupt vector
+ * register state on ET-SoC-1's MMX-style shared FP file.
+ */
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+ const char * bias_ptr3 = bias_base ? bias_base + i3 * nbb3 : (const char *) 0;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+ const char * bias_ptr2 = bias_ptr3 ? bias_ptr3 + i2 * nbb2 : (const char *) 0;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+ const float * bias_n = bias_ptr2 ? (const float *) (bias_ptr2 + n * nbb1) : (const float *) 0;
+
+ for (int64_t m0 = hart_id; m0 < M; m0 += STRIDE_M * 4) {
+ const int64_t m1 = m0 + STRIDE_M;
+ const int64_t m2 = m0 + STRIDE_M * 2;
+ const int64_t m3 = m0 + STRIDE_M * 3;
+
+ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
+
+ for (int64_t kb = 0; kb < K_blocks; kb += TILE_KB) {
+ int64_t tile_len = K_blocks - kb;
+ if (tile_len > TILE_KB) {
+ tile_len = TILE_KB;
+ }
+ const float * b_tile = b_col_base + kb * 32;
+
+ s0 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m0 * nb01) + kb, b_tile,
+ tile_len);
+ if (m1 < M) {
+ s1 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m1 * nb01) + kb, b_tile,
+ tile_len);
+ }
+ if (m2 < M) {
+ s2 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m2 * nb01) + kb, b_tile,
+ tile_len);
+ }
+ if (m3 < M) {
+ s3 += compute_row_dot_q8_0((const block_q8_0 *) (src0_ptr2 + m3 * nb01) + kb, b_tile,
+ tile_len);
+ }
+ }
+
+ float * dst_base = (float *) (dst_ptr2 + n * nbd1);
+ const float b0 = bias_n ? bias_n[m0] : 0.0f;
+ const float b1 = (bias_n && m1 < M) ? bias_n[m1] : 0.0f;
+ const float b2 = (bias_n && m2 < M) ? bias_n[m2] : 0.0f;
+ const float b3 = (bias_n && m3 < M) ? bias_n[m3] : 0.0f;
+ atomic_store_f32((volatile float *) (dst_base + m0), s0 + b0);
+ if (m1 < M) {
+ atomic_store_f32((volatile float *) (dst_base + m1), s1 + b1);
+ }
+ if (m2 < M) {
+ atomic_store_f32((volatile float *) (dst_base + m2), s2 + b2);
+ }
+ if (m3 < M) {
+ atomic_store_f32((volatile float *) (dst_base + m3), s3 + b3);
+ }
+ }
+ }
+ }
+ }
+ } else {
+ /*
+ * Simple path for small K.
+ *
+ * When `nb01` is 32-byte aligned, every row has the same block-alignment
+ * pattern. That lets us compute two rows together and reuse each loaded
+ * B chunk across both rows instead of reloading it in a second dot call.
+ */
+ for (int64_t i3 = 0; i3 < ne13; i3++) {
+ const int64_t i03 = i3 / r3;
+ const char * src0_ptr3 = (const char *) params->src0.data + i03 * nb03;
+ const char * src1_ptr3 = (const char *) params->src1.data + i3 * nb13;
+ char * dst_ptr3 = (char *) params->dst.data + i3 * nbd3;
+ const char * bias_ptr3 = bias_base ? bias_base + i3 * nbb3 : (const char *) 0;
+
+ for (int64_t i2 = 0; i2 < ne12; i2++) {
+ const int64_t i02 = i2 / r2;
+ const char * src0_ptr2 = src0_ptr3 + i02 * nb02;
+ const char * src1_ptr2 = src1_ptr3 + i2 * nb12;
+ char * dst_ptr2 = dst_ptr3 + i2 * nbd2;
+ const char * bias_ptr2 = bias_ptr3 ? bias_ptr3 + i2 * nbb2 : (const char *) 0;
+
+ for (int64_t n = 0; n < N; n++) {
+ const float * b_col_base = (const float *) (src1_ptr2 + n * nb11);
+ const float * bias_n = bias_ptr2 ? (const float *) (bias_ptr2 + n * nbb1) : (const float *) 0;
+ q8_dot_state q8_state;
+ q8_dot_begin(&q8_state);
+
+ if (use_simple_x2) {
+ for (int64_t m0 = hart_id; m0 < M; m0 += STRIDE_M * SIMPLE_X2_ROWS) {
+ const int64_t m1 = m0 + STRIDE_M;
+ const block_q8_0 * q_row0 = (const block_q8_0 *) (src0_ptr2 + m0 * nb01);
+
+ if (m1 < M) {
+ const block_q8_0 * q_row1 = (const block_q8_0 *) (src0_ptr2 + m1 * nb01);
+ float s0, s1;
+ q8_dot_compute_x2_aligned(q_row0, q_row1, b_col_base, K_blocks, &s0, &s1);
+
+ float * dst0 = (float *) (dst_ptr2 + n * nbd1 + m0 * sizeof(float));
+ float * dst1 = (float *) (dst_ptr2 + n * nbd1 + m1 * sizeof(float));
+ if (bias_n) {
+ s0 += bias_n[m0];
+ s1 += bias_n[m1];
+ }
+ atomic_store_f32((volatile float *) dst0, s0);
+ atomic_store_f32((volatile float *) dst1, s1);
+ } else {
+ float sum = q8_dot_compute(q_row0, b_col_base, K_blocks);
+ float * dst = (float *) (dst_ptr2 + n * nbd1 + m0 * sizeof(float));
+ if (bias_n) {
+ sum += bias_n[m0];
+ }
+ atomic_store_f32((volatile float *) dst, sum);
+ }
+ }
+ } else {
+ for (int64_t m = hart_id; m < M; m += STRIDE_M) {
+ const block_q8_0 * q_row = (const block_q8_0 *) (src0_ptr2 + m * nb01);
+
+ float sum = q8_dot_compute(q_row, b_col_base, K_blocks);
+
+ float * dst_entry = (float *) (dst_ptr2 + n * nbd1 + m * sizeof(float));
+ if (bias_n) {
+ sum += bias_n[m];
+ }
+ atomic_store_f32((volatile float *) dst_entry, sum);
+ }
+ }
+
+ q8_dot_end(&q8_state);
+ }
+ }
+ }
+ }
+
+#ifdef ET_UBERKERNEL
+ FENCE;
+ evict_region_past_l2(params->dst.data, tensor_bytes(¶ms->dst));
+ WAIT_CACHEOPS;
+ FENCE;
+#endif
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// MUL_MAT Kernel
+// Matrix multiplication: C[M,N] = A[M,K] * B[K,N]
+//******************************************************************************
+
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <stdint.h>
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env || params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ // Thread coordination
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0 || (thread_id & 1)) {
+ return 0; // Skip odd threads to avoid resource contention
+ }
+
+ int effective_thread_id = thread_id / 2;
+ int effective_num_threads = (num_threads + 1) / 2;
+
+ // Extract tensor references
+ struct ggml_tensor * src0 = ¶ms->src0; // Weight matrix A (F16)
+ struct ggml_tensor * src1 = ¶ms->src1; // Activation matrix B (F16/F32)
+ struct ggml_tensor * dst = ¶ms->dst; // Output matrix C (F32)
+
+ // Generic non-matrix-engine path: F16 x (F16/F32) -> F32
+ if (src0->type != GGML_TYPE_F16 || (src1->type != GGML_TYPE_F16 && src1->type != GGML_TYPE_F32) ||
+ dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const uint16_t * src0_data = (const uint16_t *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ // Dimensions and Strides
+ const int64_t K = src0->ne[0];
+ const int64_t M = src0->ne[1];
+ const int64_t N = src1->ne[1];
+
+ const int64_t ne02 = src0->ne[2], ne03 = src0->ne[3];
+ const int64_t ne12 = src1->ne[2], ne13 = src1->ne[3];
+ const int64_t ne2 = dst->ne[2], ne3 = dst->ne[3];
+
+ const size_t nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
+ const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+
+ // F16 specific block size (Usually QK_F16)
+ const int block_size = QK_F16;
+ const int64_t K_blocks = K / block_size;
+ const int64_t K_remainder = K % block_size;
+
+ // Threading distribution
+ const uint64_t total_elements = M * N * ne2 * ne3;
+ const uint64_t per_thread = 16;
+ const uint64_t threads_stride = per_thread * effective_num_threads;
+
+ if (effective_thread_id * per_thread >= total_elements) {
+ return 0;
+ }
+
+ // Broadcasting support
+ const int64_t r2 = ne12 / ne02;
+ const int64_t r3 = ne13 / ne03;
+
+ for (uint64_t base_idx = effective_thread_id * per_thread; base_idx < total_elements; base_idx += threads_stride) {
+ for (uint64_t j = 0; j < per_thread; j++) {
+ const uint64_t idx = base_idx + j;
+ if (idx >= total_elements) {
+ break;
+ }
+
+ // Index decoding
+ const int64_t i3 = idx / (M * N * ne2);
+ const int64_t rem3 = idx % (M * N * ne2);
+ const int64_t i2 = rem3 / (M * N);
+ const int64_t rem2 = rem3 % (M * N);
+ const int64_t n = rem2 / M;
+ const int64_t m = rem2 % M;
+
+ const int64_t i03 = i3 / r3, i02 = i2 / r2;
+ const int64_t i13 = (ne13 > 1) ? i3 : 0, i12 = (ne12 > 1) ? i2 : 0;
+
+ float sum = 0.0f;
+ const uint16_t * f16_row =
+ (const uint16_t *) ((const char *) src0_data + m * nb01 + i02 * nb02 + i03 * nb03);
+
+ if (src1->type == GGML_TYPE_F32) {
+ const float * src1_data = (const float *) src1->data;
+
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const float * b_col_ptr =
+ (const float *) ((const char *) src1_data + (kb * block_size) * sizeof(float) + n * nb11 +
+ i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f16_naive(&f16_row[kb * block_size], b_col_ptr);
+ }
+
+ if (K_remainder > 0) {
+ const int64_t offset = K_blocks * block_size;
+ const float * b_col_ptr = (const float *) ((const char *) src1_data + offset * sizeof(float) +
+ n * nb11 + i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f16_partial(&f16_row[offset], b_col_ptr, K_remainder);
+ }
+ } else {
+ const uint16_t * src1_data = (const uint16_t *) src1->data;
+
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const uint16_t * b_col_ptr =
+ (const uint16_t *) ((const char *) src1_data + (kb * block_size) * sizeof(uint16_t) + n * nb11 +
+ i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f16_f16_partial(&f16_row[kb * block_size], b_col_ptr, block_size);
+ }
+
+ if (K_remainder > 0) {
+ const int64_t offset = K_blocks * block_size;
+ const uint16_t * b_col_ptr =
+ (const uint16_t *) ((const char *) src1_data + offset * sizeof(uint16_t) + n * nb11 +
+ i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f16_f16_partial(&f16_row[offset], b_col_ptr, K_remainder);
+ }
+ }
+
+ // Atomic store for output
+ volatile float * c_element =
+ (volatile float *) ((char *) dst_data + m * dst->nb[0] + n * nb1 + i2 * nb2 + i3 * nb3);
+ atomic_store_f32(c_element, sum);
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+#include "ggml_tensor.h"
+#include "platform.h"
+#include "tensor.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+
+// FP16 x FP16 -> FP32 MUL_MAT with hart 1 B-panel packing
+//
+// Hart 0: tensor engine (load A, load B from SCP, FMA, reduce, store)
+// Hart 1: pack B into double-buffered L2 SCP panels, flush for tensor_load
+//
+// Sync: monotonic counters in L2 SCP with evict-based coherency.
+// Double-buffered bpanel allows pack/FMA overlap.
+//
+#define NUM_COMPUTE_SHIRES 32
+#define MINIONS_PER_SHIRE 32
+
+#define TILE_M 16
+#define TILE_N 16
+#define TILE_K 32
+
+#define CACHEOP_MAX 0
+#define REP_RATE 0
+
+#define A_L1_START 0 // SCP lines 0..15 for A
+#define B_L1_START 16 // SCP lines 16..31 for B
+
+typedef uint16_t et_fp16_t;
+
+// L2 SCP layout per minion (double-buffered bpanel + sync counters)
+// [0..1023] bpanel buffer 0 (16 lines x 64 bytes)
+// [1024..2047] bpanel buffer 1
+// [2048..2111] ready counter (hart1 -> hart0, own cache line)
+// [2112..2175] consumed counter (hart0 -> hart1, own cache line)
+#define SCP_BPANEL_SIZE (16 * 32 * sizeof(et_fp16_t)) // 1024 bytes
+#define SCP_READY_OFF (2 * SCP_BPANEL_SIZE) // 2048
+#define SCP_CONSUMED_OFF (SCP_READY_OFF + 64) // 2112
+#define SCP_PER_MINION (SCP_CONSUMED_OFF + 64) // 2176
+
+// Signal a counter value to the other hart via L2 SCP.
+static inline void __attribute__((always_inline)) scp_signal(volatile uint32_t * flag, uint32_t value) {
+ *flag = value;
+ FENCE;
+ evict_to_l2((const void *) flag, 1, 64);
+ WAIT_CACHEOPS;
+}
+
+// Wait for a counter in L2 SCP to reach the expected value.
+static inline void __attribute__((always_inline)) scp_wait(volatile uint32_t * flag, uint32_t expected) {
+ while (1) {
+ evict_to_l2((const void *) flag, 1, 64);
+ WAIT_CACHEOPS;
+ if (*flag >= expected) {
+ return;
+ }
+ }
+}
+
+/**
+ * Build the interleaved B panel that TensorFMA16A32 expects (vectorized).
+ *
+ * Output: 16 lines x 32 fp16 = 1024 bytes, 64-byte aligned.
+ * out[l][j*2+0] = src0[mb + j][kb + 2*l]
+ * out[l][j*2+1] = src0[mb + j][kb + 2*l + 1]
+ *
+ * Uses fsch.ps scatter store: load 8 pairs per row, scatter to 8 output lines.
+ */
+static inline void __attribute__((always_inline)) pack_b_interleaved(et_fp16_t * out,
+ const char * src0_batch,
+ int64_t mb,
+ int64_t kb,
+ int64_t nb1_0) {
+ static const int32_t __attribute__((aligned(32))) scatter_idx[8] = { 0, 64, 128, 192, 256, 320, 384, 448 };
+
+ unsigned long old_mask;
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f1, 0(%[idx]) \n\t"
+ : [ms] "=&r"(old_mask)
+ : [idx] "r"(scatter_idx)
+ : "f1");
+
+ for (int j = 0; j < TILE_M; ++j) {
+ const et_fp16_t * row = (const et_fp16_t *) (src0_batch + (mb + j) * nb1_0) + kb;
+ char * dst = (char *) out + j * 4;
+
+ __asm__ volatile(
+ "flw.ps f2, 0(%[src]) \n\t"
+ "flw.ps f3, 32(%[src]) \n\t"
+ "fscw.ps f2, f1(%[d0]) \n\t"
+ "fscw.ps f3, f1(%[d1]) \n\t"
+ :
+ : [src] "r"(row), [d0] "r"(dst), [d1] "r"(dst + 512)
+ : "f2", "f3", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %[ms] \n\t" : : [ms] "r"(old_mask));
+}
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ (void) env;
+
+ uint64_t hart_id = get_hart_id();
+ uint64_t shire_id = get_shire_id();
+
+ if (shire_id >= NUM_COMPUTE_SHIRES) {
+ return 0;
+ }
+
+ const int is_hart1 = hart_id & 1;
+ uint64_t local_minion = (hart_id >> 1) & 0x1F;
+
+ // Dimensions (both harts need these for tile assignment)
+ const int64_t K = params->src0.ne[0];
+ const int64_t M = params->src0.ne[1];
+ const int64_t N = params->src1.ne[1];
+
+ const int64_t ne2_0 = params->src0.ne[2], ne3_0 = params->src0.ne[3];
+ const int64_t ne2_1 = params->src1.ne[2], ne3_1 = params->src1.ne[3];
+
+ const int64_t nb1_0 = params->src0.nb[1];
+ const int64_t nb2_0 = params->src0.nb[2], nb3_0 = params->src0.nb[3];
+
+ const int64_t nb1_1 = params->src1.nb[1];
+ const int64_t nb2_1 = params->src1.nb[2], nb3_1 = params->src1.nb[3];
+
+ const int64_t nb1_d = params->dst.nb[1];
+ const int64_t nb2_d = params->dst.nb[2], nb3_d = params->dst.nb[3];
+
+ const char * src0_base = (const char *) params->src0.data;
+ const char * src1_base = (const char *) params->src1.data;
+ char * dst_base = (char *) params->dst.data;
+
+ if ((M % TILE_M) != 0) {
+ return 0;
+ }
+ if ((K % TILE_K) != 0) {
+ return 0;
+ }
+
+ const int64_t m_tiles = M / TILE_M;
+ const int64_t n_tiles = (N + TILE_N - 1) / TILE_N;
+ const int64_t batch_count = ne2_1 * ne3_1;
+ const int64_t base_tiles = m_tiles * n_tiles * batch_count;
+
+ const int64_t r2 = ne2_1 / ne2_0;
+ const int64_t r3 = ne3_1 / ne3_0;
+
+ const int64_t total_harts = NUM_COMPUTE_SHIRES * MINIONS_PER_SHIRE;
+ const int64_t k_steps = K / TILE_K;
+
+ int64_t k_splits = 1;
+ if (base_tiles < total_harts) {
+ k_splits = (total_harts + base_tiles - 1) / base_tiles;
+ int64_t ks = 1;
+ while (ks * 2 <= k_splits && ks * 2 <= 32 && k_steps % (ks * 2) == 0) {
+ ks *= 2;
+ }
+ k_splits = ks;
+ }
+
+ const int64_t tiles_per_shire = MINIONS_PER_SHIRE / k_splits;
+ const int64_t k_split = local_minion % k_splits;
+ const int64_t local_tile_idx = local_minion / k_splits;
+ const int64_t tiles_stride = (int64_t) NUM_COMPUTE_SHIRES * tiles_per_shire;
+
+ const int64_t k_steps_per_split = k_steps / k_splits;
+ const int64_t k_start = k_split * k_steps_per_split * TILE_K;
+ const int64_t k_end = k_start + k_steps_per_split * TILE_K;
+
+ // L2 SCP pointers for this minion's double-buffered panels + sync
+ uint64_t scp_base = local_minion * SCP_PER_MINION;
+ et_fp16_t * scp_bp[2] = {
+ (et_fp16_t *) et_shire_l2scp_local(scp_base),
+ (et_fp16_t *) et_shire_l2scp_local(scp_base + SCP_BPANEL_SIZE),
+ };
+ volatile uint32_t * ready_ctr = (volatile uint32_t *) et_shire_l2scp_local(scp_base + SCP_READY_OFF);
+ volatile uint32_t * consumed_ctr = (volatile uint32_t *) et_shire_l2scp_local(scp_base + SCP_CONSUMED_OFF);
+
+ // ================================================================
+ // Hart 1: B-panel packer
+ // ================================================================
+ if (is_hart1) {
+ // Initialize sync counters
+ scp_signal(ready_ctr, 0);
+ scp_signal(consumed_ctr, 0);
+
+ uint32_t chunk_id = 0;
+
+ for (int64_t tile = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES; tile < base_tiles;
+ tile += tiles_stride) {
+ const int64_t tiles_per_batch = m_tiles * n_tiles;
+ const int64_t batch_idx = tile / tiles_per_batch;
+ const int64_t tile_in_batch = tile % tiles_per_batch;
+
+ const int64_t mb_idx = tile_in_batch % m_tiles;
+
+ const int64_t i3 = batch_idx / ne2_1;
+ const int64_t i2 = batch_idx % ne2_1;
+ const int64_t i2_0 = i2 / r2;
+ const int64_t i3_0 = i3 / r3;
+
+ const char * src0_batch = src0_base + i3_0 * nb3_0 + i2_0 * nb2_0;
+ const int64_t mb = mb_idx * TILE_M;
+
+ for (int64_t kb = k_start; kb < k_end; kb += TILE_K) {
+ int buf = chunk_id & 1;
+
+ // Back-pressure: wait for hart 0 to finish with this buffer
+ if (chunk_id >= 2) {
+ scp_wait(consumed_ctr, chunk_id - 1);
+ }
+
+ pack_b_interleaved(scp_bp[buf], src0_batch, mb, kb, nb1_0);
+
+ FENCE;
+ flush_to_l2(scp_bp[buf], 16, 64);
+ WAIT_CACHEOPS;
+
+ chunk_id++;
+ scp_signal(ready_ctr, chunk_id);
+ }
+ }
+
+ FENCE;
+ return 0;
+ }
+
+ // ================================================================
+ // Hart 0: tensor engine compute
+ // ================================================================
+ uint64_t my_minion_id = get_minion_id();
+ const uint64_t group_base_global = my_minion_id - k_split;
+
+ setup_cache_scp();
+#if CACHEOP_MAX > 0 || REP_RATE > 0
+ ucache_control(1, REP_RATE, CACHEOP_MAX);
+#endif
+ CLEAR_TENSOR_ERROR;
+
+ // Evict any stale L1D copies of sync counters
+ evict_to_l2((const void *) ready_ctr, 1, 64);
+ WAIT_CACHEOPS;
+ evict_to_l2((const void *) consumed_ctr, 1, 64);
+ WAIT_CACHEOPS;
+
+ uint32_t chunk_id = 0;
+
+ for (int64_t tile = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES; tile < base_tiles;
+ tile += tiles_stride) {
+ const int64_t tiles_per_batch = m_tiles * n_tiles;
+ const int64_t batch_idx = tile / tiles_per_batch;
+ const int64_t tile_in_batch = tile % tiles_per_batch;
+
+ const int64_t nb_idx = tile_in_batch / m_tiles;
+ const int64_t mb_idx = tile_in_batch % m_tiles;
+
+ const int64_t i3 = batch_idx / ne2_1;
+ const int64_t i2 = batch_idx % ne2_1;
+
+ const char * src1_batch = src1_base + i3 * nb3_1 + i2 * nb2_1;
+ char * dst_batch = dst_base + i3 * nb3_d + i2 * nb2_d;
+
+ const int64_t mb = mb_idx * TILE_M;
+ const int64_t nb = nb_idx * TILE_N;
+ const int64_t n_cur = (nb + TILE_N <= N) ? TILE_N : (N - nb);
+
+ // Set tensor_mask for partial N tiles
+ if (n_cur < TILE_N) {
+ uint64_t mask = (1ULL << n_cur) - 1;
+ __asm__ __volatile__("csrw 0x805, %0" : : "r"(mask));
+ }
+
+ for (int64_t kb = k_start; kb < k_end; kb += TILE_K) {
+ int buf = chunk_id & 1;
+
+ // Start loading A from DRAM (overlaps with waiting for hart 1)
+ tensor_load((n_cur < TILE_N), false, A_L1_START, TENSOR_LOAD_PLAIN, 0,
+ (uint64_t) (src1_batch + nb * nb1_1 + kb * (int64_t) sizeof(et_fp16_t)), 0, n_cur - 1,
+ (uint64_t) nb1_1, 0);
+
+ // Wait for hart 1 to finish packing this chunk
+ chunk_id++;
+ scp_wait(ready_ctr, chunk_id);
+
+ // Load B from L2 SCP (hart 1 already flushed it)
+ tensor_load(false, false, B_L1_START, TENSOR_LOAD_PLAIN, 0, (uint64_t) scp_bp[buf], 0, 15, 64, 1);
+
+ tensor_wait(TENSOR_LOAD_WAIT_0);
+ tensor_wait(TENSOR_LOAD_WAIT_1);
+
+ // TensorFMA16A32
+ tensor_fma((n_cur < TILE_N), 3, n_cur - 1, 15, 0, false, false, false, false, B_L1_START, A_L1_START,
+ TENSOR_FMA_OP_FP16, (kb == k_start));
+
+ tensor_wait(TENSOR_FMA_WAIT);
+
+ // Signal that this buffer is free for hart 1 to reuse
+ scp_signal(consumed_ctr, chunk_id);
+ }
+
+ // K-split ring reduce
+ if (k_splits > 1) {
+ const uint64_t num_regs = (uint64_t) n_cur * 2;
+
+ if (k_split > 0) {
+ tensor_reduce_recv(0, TENSOR_REDUCE_OP_FADD, num_regs, group_base_global + k_split - 1);
+ tensor_wait(TENSOR_REDUCE_WAIT);
+ }
+
+ if (k_split < k_splits - 1) {
+ tensor_reduce_send(0, num_regs, group_base_global + k_split + 1);
+ tensor_wait(TENSOR_REDUCE_WAIT);
+ }
+ }
+
+ // Store FP32 result tile
+ if (k_split == k_splits - 1) {
+ tensor_store(0, 0, 3, n_cur - 1, (uint64_t) (dst_batch + nb * nb1_d + mb * (int64_t) sizeof(float)), 0,
+ (uint64_t) nb1_d);
+ tensor_wait(TENSOR_STORE_WAIT);
+ }
+ }
+
+ FENCE;
+ return 0;
+}
--- /dev/null
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+#include <stdio.h>
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env || params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ // Thread coordination
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0 || (thread_id & 1)) {
+ return 0; // Skip odd threads to avoid resource contention
+ }
+
+ int effective_thread_id = thread_id / 2;
+ int effective_num_threads = (num_threads + 1) / 2;
+
+ // Extract tensor references
+ struct ggml_tensor * src0 = ¶ms->src0; // Weight matrix A (F32)
+ struct ggml_tensor * src1 = ¶ms->src1; // Activation matrix B (F16/F32)
+ struct ggml_tensor * dst = ¶ms->dst; // Output matrix C (F32)
+
+ // Generic non-matrix-engine path: F32 x (F16/F32) -> F32
+ if (src0->type != GGML_TYPE_F32 || (src1->type != GGML_TYPE_F16 && src1->type != GGML_TYPE_F32) ||
+ dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * src0_data = (const float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ // Dimensions and Strides
+ const int64_t K = src0->ne[0];
+ const int64_t M = src0->ne[1];
+ const int64_t N = src1->ne[1];
+
+ const int64_t ne02 = src0->ne[2], ne03 = src0->ne[3];
+ const int64_t ne12 = src1->ne[2], ne13 = src1->ne[3];
+ const int64_t ne2 = dst->ne[2], ne3 = dst->ne[3];
+
+ const size_t nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
+ const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+
+ // F32 specific block size and counts
+ const int block_size = QK_F32;
+ const int64_t K_blocks = K / block_size;
+ const int64_t K_remainder = K % block_size;
+
+ // Threading distribution
+ const uint64_t total_elements = M * N * ne2 * ne3;
+ const uint64_t per_thread = 16;
+ const uint64_t threads_stride = per_thread * effective_num_threads;
+
+ if (effective_thread_id * per_thread >= total_elements) {
+ return 0;
+ }
+
+ // Broadcasting support
+ const int64_t r2 = ne12 / ne02;
+ const int64_t r3 = ne13 / ne03;
+
+ for (uint64_t base_idx = effective_thread_id * per_thread; base_idx < total_elements; base_idx += threads_stride) {
+ for (uint64_t j = 0; j < per_thread; j++) {
+ const uint64_t idx = base_idx + j;
+ if (idx >= total_elements) {
+ break;
+ }
+
+ // Index decoding
+ const int64_t i3 = idx / (M * N * ne2);
+ const int64_t rem3 = idx % (M * N * ne2);
+ const int64_t i2 = rem3 / (M * N);
+ const int64_t rem2 = rem3 % (M * N);
+ const int64_t n = rem2 / M;
+ const int64_t m = rem2 % M;
+
+ const int64_t i03 = i3 / r3, i02 = i2 / r2;
+ const int64_t i13 = (ne13 > 1) ? i3 : 0, i12 = (ne12 > 1) ? i2 : 0;
+
+ float sum = 0.0f;
+ const float * f32_row = (const float *) ((const char *) src0_data + m * nb01 + i02 * nb02 + i03 * nb03);
+
+ if (src1->type == GGML_TYPE_F32) {
+ const float * src1_data = (const float *) src1->data;
+
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const float * b_col_ptr =
+ (const float *) ((const char *) src1_data + (kb * block_size) * sizeof(float) + n * nb11 +
+ i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f32(&f32_row[kb * block_size], b_col_ptr);
+ }
+
+ if (K_remainder > 0) {
+ const int64_t offset = K_blocks * block_size;
+ const float * b_col_ptr = (const float *) ((const char *) src1_data + offset * sizeof(float) +
+ n * nb11 + i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f32_partial(&f32_row[offset], b_col_ptr, K_remainder);
+ }
+ } else {
+ const uint16_t * src1_data = (const uint16_t *) src1->data;
+
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const uint16_t * b_col_ptr =
+ (const uint16_t *) ((const char *) src1_data + (kb * block_size) * sizeof(uint16_t) + n * nb11 +
+ i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f32_f16_partial(&f32_row[kb * block_size], b_col_ptr, block_size);
+ }
+
+ if (K_remainder > 0) {
+ const int64_t offset = K_blocks * block_size;
+ const uint16_t * b_col_ptr =
+ (const uint16_t *) ((const char *) src1_data + offset * sizeof(uint16_t) + n * nb11 +
+ i12 * nb12 + i13 * nb13);
+ sum += compute_block_dot_product_f32_f16_partial(&f32_row[offset], b_col_ptr, K_remainder);
+ }
+ }
+
+ // Atomic store for output
+ volatile float * c_element =
+ (volatile float *) ((char *) dst_data + m * dst->nb[0] + n * nb1 + i2 * nb2 + i3 * nb3);
+ atomic_store_f32(c_element, sum);
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+#include "ggml_tensor.h"
+#include "platform.h"
+#include "tensor.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+
+/*
+ * F32 Matrix Multiply for ET-SoC-1 — TensorFMA32.
+ *
+ * K-parallel + interleaved tiles + ring reduce.
+ * No batched-K yet (needs investigation on hang).
+ * This is the last known working version.
+ */
+
+#define NUM_COMPUTE_SHIRES 32
+#define MINIONS_PER_SHIRE 32
+#define TILE_K 16
+#define TILE_M 16
+
+/* ── Tuning knobs ───────────────────────────────────────────────────── */
+#define TILE_N 16
+#define CACHEOP_MAX 0
+#define REP_RATE 0
+
+/* ─────────────────────────────────────────────────────────────────── */
+
+int entry_point(struct ggml_et_binary_params * params, void * env) {
+ uint64_t hart_id = get_hart_id();
+ uint64_t shire_id = get_shire_id();
+
+ if (shire_id >= NUM_COMPUTE_SHIRES) {
+ return 0;
+ }
+ if (hart_id & 1) {
+ return 0;
+ }
+
+ uint64_t local_minion = (hart_id >> 1) & 0x1F;
+ uint64_t my_minion_id = get_minion_id();
+
+ const int64_t K = params->src0.ne[0];
+ const int64_t M = params->src0.ne[1];
+ const int64_t N = params->src1.ne[1];
+
+ const int64_t ne2_0 = params->src0.ne[2], ne3_0 = params->src0.ne[3];
+ const int64_t ne2_1 = params->src1.ne[2], ne3_1 = params->src1.ne[3];
+
+ const int64_t nb1_0 = params->src0.nb[1];
+ const int64_t nb2_0 = params->src0.nb[2], nb3_0 = params->src0.nb[3];
+ const int64_t nb1_1 = params->src1.nb[1];
+ const int64_t nb2_1 = params->src1.nb[2], nb3_1 = params->src1.nb[3];
+ const int64_t nb1_d = params->dst.nb[1];
+ const int64_t nb2_d = params->dst.nb[2], nb3_d = params->dst.nb[3];
+
+ const char * src0_base = (const char *) params->src0.data;
+ const char * src1_base = (const char *) params->src1.data;
+ char * dst_base = (char *) params->dst.data;
+
+ setup_cache_scp();
+#if CACHEOP_MAX > 0 || REP_RATE > 0
+ ucache_control(1, REP_RATE, CACHEOP_MAX);
+#endif
+ CLEAR_TENSOR_ERROR;
+
+ const int64_t m_tiles = M / TILE_M;
+ const int64_t n_tiles = (N + TILE_N - 1) / TILE_N;
+ const int64_t batch_count = ne2_1 * ne3_1;
+ const int64_t base_tiles = m_tiles * n_tiles * batch_count;
+
+ const int64_t r2 = ne2_1 / ne2_0;
+ const int64_t r3 = ne3_1 / ne3_0;
+
+ const int64_t total_harts = NUM_COMPUTE_SHIRES * MINIONS_PER_SHIRE;
+ const int64_t k_steps = K / TILE_K;
+ int64_t k_splits = 1;
+ if (base_tiles < total_harts) {
+ k_splits = (total_harts + base_tiles - 1) / base_tiles;
+ int64_t ks = 1;
+ while (ks * 2 <= k_splits && ks * 2 <= 32 && k_steps % (ks * 2) == 0) {
+ ks *= 2;
+ }
+ k_splits = ks;
+ }
+
+ const int64_t tiles_per_shire = MINIONS_PER_SHIRE / k_splits;
+ const int64_t k_split = local_minion % k_splits;
+ const int64_t local_tile_idx = local_minion / k_splits;
+ const int64_t tiles_stride = (int64_t) NUM_COMPUTE_SHIRES * tiles_per_shire;
+
+ const int64_t k_steps_per_split = k_steps / k_splits;
+ const int64_t k_start = k_split * k_steps_per_split * TILE_K;
+ const int64_t k_end = k_start + k_steps_per_split * TILE_K;
+
+ const uint64_t group_base_global = my_minion_id - k_split;
+
+ for (int64_t tile = (int64_t) shire_id + local_tile_idx * NUM_COMPUTE_SHIRES; tile < base_tiles;
+ tile += tiles_stride) {
+ const int64_t tiles_per_batch = m_tiles * n_tiles;
+ const int64_t batch_idx = tile / tiles_per_batch;
+ const int64_t tile_in_batch = tile % tiles_per_batch;
+ const int64_t nb_idx = tile_in_batch / m_tiles;
+ const int64_t mb_idx = tile_in_batch % m_tiles;
+
+ const int64_t i3 = batch_idx / ne2_1;
+ const int64_t i2 = batch_idx % ne2_1;
+ const int64_t i2_0 = i2 / r2;
+ const int64_t i3_0 = i3 / r3;
+
+ const char * src0_batch = src0_base + i3_0 * nb3_0 + i2_0 * nb2_0;
+ const char * src1_batch = src1_base + i3 * nb3_1 + i2 * nb2_1;
+ char * dst_batch = dst_base + i3 * nb3_d + i2 * nb2_d;
+
+ const int64_t mb = mb_idx * TILE_M;
+ const int64_t nb = nb_idx * TILE_N;
+ const int64_t n_cur = (nb + TILE_N <= N) ? TILE_N : (N - nb);
+
+ for (int64_t kb = k_start; kb < k_end; kb += TILE_K) {
+ tensor_load(false, false, 0, 0, 0, (uint64_t) (src1_batch + nb * nb1_1 + kb * sizeof(float)), 0, n_cur - 1,
+ (uint64_t) nb1_1, 0);
+
+ tensor_load(false, false, TILE_K, 7, 0, (uint64_t) (src0_batch + mb * nb1_0 + kb * sizeof(float)), 0,
+ TILE_K - 1, (uint64_t) nb1_0, 1);
+
+ tensor_wait(TENSOR_LOAD_WAIT_0);
+ tensor_wait(TENSOR_LOAD_WAIT_1);
+
+ tensor_fma(false, 3, n_cur - 1, TILE_K - 1, 0, false, false, false, false, TILE_K, 0, 0, (kb == k_start));
+
+ tensor_wait(TENSOR_FMA_WAIT);
+ }
+
+ if (k_splits > 1) {
+ const uint64_t num_regs = (uint64_t) n_cur * 2;
+
+ if (k_split > 0) {
+ tensor_reduce_recv(0, TENSOR_REDUCE_OP_FADD, num_regs, group_base_global + k_split - 1);
+ tensor_wait(TENSOR_REDUCE_WAIT);
+ }
+ if (k_split < k_splits - 1) {
+ tensor_reduce_send(0, num_regs, group_base_global + k_split + 1);
+ tensor_wait(TENSOR_REDUCE_WAIT);
+ }
+ }
+
+ if (k_split == k_splits - 1) {
+ tensor_store(0, 0, 3, n_cur - 1, (uint64_t) (dst_batch + nb * nb1_d + mb * sizeof(float)), 0,
+ (uint64_t) nb1_d);
+ tensor_wait(TENSOR_STORE_WAIT);
+ }
+ }
+
+ FENCE;
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// MUL_MAT_ID kernel specialized for Q4_0 weights (Mixture of Experts).
+//
+// C[m, s, b] = Sum(k=0..K-1) A[k, m, ids[s,b]] * B[k, s % ne11, b]
+// A: Q4_0 [K, M, n_expert] weights
+// B: F32 [K, n_cols, batch] activations
+// ids: I32 [n_expert_used, batch]
+// C: F32 [M, n_expert_used, batch]
+//
+// Strategy: All harts active. Flat m-major output partition allows amortized
+// expert lookups and 2-row x2 dot products.
+//******************************************************************************
+
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <stdint.h>
+
+int entry_point(struct ggml_et_mul_mat_id_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+ if (!kernel_env || !params) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * src2 = ¶ms->src2;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_Q4_0 || src1->type != GGML_TYPE_F32 || src2->type != GGML_TYPE_I32 ||
+ dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const void * src0_data = src0->data;
+ const float * src1_data = (const float *) src1->data;
+ const int32_t * src2_data = (const int32_t *) src2->data;
+ float * dst_data = (float *) dst->data;
+ if (!src0_data || !src1_data || !src2_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t K = src0->ne[0];
+ const int64_t M = src0->ne[1];
+ const int64_t n_expert = src0->ne[2];
+ const int64_t n_expert_used = src2->ne[0];
+ const int64_t batch = src2->ne[1];
+ const int64_t ne11 = src1->ne[1];
+
+ if (K % QK4_0 != 0) {
+ return -1;
+ }
+
+ const size_t nb01 = src0->nb[1]; // Q4_0 row stride
+ const size_t nb02 = src0->nb[2]; // expert stride
+ const size_t nb11 = src1->nb[1]; // activation column stride
+ const size_t nb12 = src1->nb[2]; // activation batch stride
+ const size_t nb20 = src2->nb[0];
+ const size_t nb21 = src2->nb[1];
+ const size_t nbd0 = dst->nb[0];
+ const size_t nbd1 = dst->nb[1];
+ const size_t nbd2 = dst->nb[2];
+
+ if (src0->nb[0] != sizeof(block_q4_0) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(int32_t) ||
+ nbd0 != sizeof(float)) {
+ return -1;
+ }
+
+ const int64_t K_blocks = K / QK4_0;
+ const int use_x2 = ((nb01 & 31) == 0);
+
+ const uint64_t total_outputs = (uint64_t) M * (uint64_t) n_expert_used * (uint64_t) batch;
+ if (total_outputs == 0) {
+ return 0;
+ }
+
+ // Even partition: hart h owns outputs [h*chunk, (h+1)*chunk).
+ const uint64_t chunk = (total_outputs + (uint64_t) num_threads - 1) / (uint64_t) num_threads;
+ const uint64_t my_start = (uint64_t) thread_id * chunk;
+ if (my_start >= total_outputs) {
+ return 0;
+ }
+ uint64_t my_end = my_start + chunk;
+ if (my_end > total_outputs) {
+ my_end = total_outputs;
+ }
+
+ // Save mask register once; full lanes for vector dot.
+ q4_dot_state q4_state;
+ q4_dot_begin(&q4_state);
+
+ const uint64_t per_batch = (uint64_t) M * (uint64_t) n_expert_used;
+
+ uint64_t idx = my_start;
+ while (idx < my_end) {
+ // Decode (m, slot, batch) from the m-major linear index.
+ const int64_t batch_idx = (int64_t) (idx / per_batch);
+ const uint64_t rem = idx - (uint64_t) batch_idx * per_batch;
+ const int64_t slot_idx = (int64_t) (rem / (uint64_t) M);
+ const int64_t m0 = (int64_t) (rem - (uint64_t) slot_idx * (uint64_t) M);
+
+ // How many outputs left in this (slot, batch) run AND in my range.
+ const uint64_t run_end_global =
+ (uint64_t) batch_idx * per_batch + (uint64_t) slot_idx * (uint64_t) M + (uint64_t) M;
+ const uint64_t end_in_my = (run_end_global < my_end) ? run_end_global : my_end;
+ int64_t run_len = (int64_t) (end_in_my - idx);
+
+ // Resolve expert + B column + dst slot for this run.
+ const int32_t expert_id =
+ *(const int32_t *) ((const char *) src2_data + slot_idx * (int64_t) nb20 + batch_idx * (int64_t) nb21);
+
+ char * dst_slot = (char *) dst_data + slot_idx * (int64_t) nbd1 + batch_idx * (int64_t) nbd2;
+
+ if (expert_id < 0 || expert_id >= n_expert) {
+ // Invalid expert id — zero out this run's outputs.
+ int64_t m = m0;
+ for (int64_t i = 0; i < run_len; i++, m++) {
+ atomic_store_f32((volatile float *) (dst_slot + m * (int64_t) nbd0), 0.0f);
+ }
+ idx += (uint64_t) run_len;
+ continue;
+ }
+
+ const int64_t col_idx = slot_idx % ne11;
+ const float * b_col_base =
+ (const float *) ((const char *) src1_data + col_idx * (int64_t) nb11 + batch_idx * (int64_t) nb12);
+ const char * expert_base = (const char *) src0_data + expert_id * (int64_t) nb02;
+
+ int64_t m = m0;
+ int64_t left = run_len;
+
+ // Paired-row dots: halves B bandwidth for runs >= 2.
+ if (use_x2) {
+ while (left >= 2) {
+ const block_q4_0 * row0 = (const block_q4_0 *) (expert_base + m * (int64_t) nb01);
+ const block_q4_0 * row1 = (const block_q4_0 *) (expert_base + (m + 1) * (int64_t) nb01);
+ float s0, s1;
+ q4_dot_compute_x2_aligned(row0, row1, b_col_base, K_blocks, &s0, &s1);
+ atomic_store_f32((volatile float *) (dst_slot + m * (int64_t) nbd0), s0);
+ atomic_store_f32((volatile float *) (dst_slot + (m + 1) * (int64_t) nbd0), s1);
+ m += 2;
+ left -= 2;
+ }
+ }
+
+ // Tail / non-aligned fallback: single-row dots.
+ while (left > 0) {
+ const block_q4_0 * row = (const block_q4_0 *) (expert_base + m * (int64_t) nb01);
+ float s = q4_dot_compute(row, b_col_base, K_blocks);
+ atomic_store_f32((volatile float *) (dst_slot + m * (int64_t) nbd0), s);
+ m++;
+ left--;
+ }
+
+ idx += (uint64_t) run_len;
+ }
+
+ q4_dot_end(&q4_state);
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// MUL_MAT_ID kernel specialized for Q8_0 weights (Mixture of Experts).
+//
+// C[m, s, b] = Sum(k=0..K-1) A[k, m, ids[s,b]] * B[k, s % ne11, b]
+// A: Q8_0 [K, M, n_expert] weights
+// B: F32 [K, n_cols, batch] activations
+// ids: I32 [n_expert_used, batch]
+// C: F32 [M, n_expert_used, batch]
+//
+// Strategy mirrors mul_mat_id_Q4_0.c.
+//******************************************************************************
+
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <stdint.h>
+
+int entry_point(struct ggml_et_mul_mat_id_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+ if (!kernel_env || !params) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * src2 = ¶ms->src2;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_Q8_0 || src1->type != GGML_TYPE_F32 || src2->type != GGML_TYPE_I32 ||
+ dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const void * src0_data = src0->data;
+ const float * src1_data = (const float *) src1->data;
+ const int32_t * src2_data = (const int32_t *) src2->data;
+ float * dst_data = (float *) dst->data;
+ if (!src0_data || !src1_data || !src2_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t K = src0->ne[0];
+ const int64_t M = src0->ne[1];
+ const int64_t n_expert = src0->ne[2];
+ const int64_t n_expert_used = src2->ne[0];
+ const int64_t batch = src2->ne[1];
+ const int64_t ne11 = src1->ne[1];
+
+ if (K % QK8_0 != 0) {
+ return -1;
+ }
+
+ const size_t nb01 = src0->nb[1];
+ const size_t nb02 = src0->nb[2];
+ const size_t nb11 = src1->nb[1];
+ const size_t nb12 = src1->nb[2];
+ const size_t nb20 = src2->nb[0];
+ const size_t nb21 = src2->nb[1];
+ const size_t nbd0 = dst->nb[0];
+ const size_t nbd1 = dst->nb[1];
+ const size_t nbd2 = dst->nb[2];
+
+ if (src0->nb[0] != sizeof(block_q8_0) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(int32_t) ||
+ nbd0 != sizeof(float)) {
+ return -1;
+ }
+
+ const int64_t K_blocks = K / QK8_0;
+ const int use_x2 = ((nb01 & 31) == 0);
+
+ const uint64_t total_outputs = (uint64_t) M * (uint64_t) n_expert_used * (uint64_t) batch;
+ if (total_outputs == 0) {
+ return 0;
+ }
+
+ const uint64_t chunk = (total_outputs + (uint64_t) num_threads - 1) / (uint64_t) num_threads;
+ const uint64_t my_start = (uint64_t) thread_id * chunk;
+ if (my_start >= total_outputs) {
+ return 0;
+ }
+ uint64_t my_end = my_start + chunk;
+ if (my_end > total_outputs) {
+ my_end = total_outputs;
+ }
+
+ q8_dot_state q8_state;
+ q8_dot_begin(&q8_state);
+
+ const uint64_t per_batch = (uint64_t) M * (uint64_t) n_expert_used;
+
+ uint64_t idx = my_start;
+ while (idx < my_end) {
+ const int64_t batch_idx = (int64_t) (idx / per_batch);
+ const uint64_t rem = idx - (uint64_t) batch_idx * per_batch;
+ const int64_t slot_idx = (int64_t) (rem / (uint64_t) M);
+ const int64_t m0 = (int64_t) (rem - (uint64_t) slot_idx * (uint64_t) M);
+
+ const uint64_t run_end_global =
+ (uint64_t) batch_idx * per_batch + (uint64_t) slot_idx * (uint64_t) M + (uint64_t) M;
+ const uint64_t end_in_my = (run_end_global < my_end) ? run_end_global : my_end;
+ int64_t run_len = (int64_t) (end_in_my - idx);
+
+ const int32_t expert_id =
+ *(const int32_t *) ((const char *) src2_data + slot_idx * (int64_t) nb20 + batch_idx * (int64_t) nb21);
+
+ char * dst_slot = (char *) dst_data + slot_idx * (int64_t) nbd1 + batch_idx * (int64_t) nbd2;
+
+ if (expert_id < 0 || expert_id >= n_expert) {
+ int64_t m = m0;
+ for (int64_t i = 0; i < run_len; i++, m++) {
+ atomic_store_f32((volatile float *) (dst_slot + m * (int64_t) nbd0), 0.0f);
+ }
+ idx += (uint64_t) run_len;
+ continue;
+ }
+
+ const int64_t col_idx = slot_idx % ne11;
+ const float * b_col_base =
+ (const float *) ((const char *) src1_data + col_idx * (int64_t) nb11 + batch_idx * (int64_t) nb12);
+ const char * expert_base = (const char *) src0_data + expert_id * (int64_t) nb02;
+
+ int64_t m = m0;
+ int64_t left = run_len;
+
+ if (use_x2) {
+ while (left >= 2) {
+ const block_q8_0 * row0 = (const block_q8_0 *) (expert_base + m * (int64_t) nb01);
+ const block_q8_0 * row1 = (const block_q8_0 *) (expert_base + (m + 1) * (int64_t) nb01);
+ float s0, s1;
+ q8_dot_compute_x2_aligned(row0, row1, b_col_base, K_blocks, &s0, &s1);
+ atomic_store_f32((volatile float *) (dst_slot + m * (int64_t) nbd0), s0);
+ atomic_store_f32((volatile float *) (dst_slot + (m + 1) * (int64_t) nbd0), s1);
+ m += 2;
+ left -= 2;
+ }
+ }
+
+ while (left > 0) {
+ const block_q8_0 * row = (const block_q8_0 *) (expert_base + m * (int64_t) nb01);
+ float s = q8_dot_compute(row, b_col_base, K_blocks);
+ atomic_store_f32((volatile float *) (dst_slot + m * (int64_t) nbd0), s);
+ m++;
+ left--;
+ }
+
+ idx += (uint64_t) run_len;
+ }
+
+ q8_dot_end(&q8_state);
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal MUL_MAT_ID Kernel (Mixture of Experts)
+//
+// ALGORITHM:
+// MUL_MAT_ID performs batched matrix multiplication with expert routing.
+// Each output element selects which expert matrix to use based on an index tensor.
+//
+// INPUTS:
+// src0 (as): Expert weight matrices [K, M, n_expert]
+// - Stack of n_expert matrices, each of size [K, M]
+// src1 (b): Activation vectors [K, n_cols, batch]
+// - n_cols can be 1 (broadcast) or n_expert_used (per-expert inputs)
+// src2 (ids): Expert selection indices [n_expert_used, batch] (int32)
+// - For each (slot, batch), specifies which expert from src0 to use
+//
+// OUTPUT:
+// dst: Result [M, n_expert_used, batch, 1]
+//
+// COMPUTATION:
+// For each output position (m, slot, batch):
+// expert_id = ids[slot, batch] // Which expert to use (0..n_expert-1)
+// col_idx = slot % src1.ne[1] // Which column in src1 (handles broadcasting)
+// dst[m, slot, batch] = dot_product(
+// src0[0:K, m, expert_id], // Row m from selected expert matrix
+// src1[0:K, col_idx, batch] // Column from activations (may broadcast)
+// )
+//
+// BROADCASTING:
+// - When src1.ne[1] == 1: All expert slots use the same activation column
+// - When src1.ne[1] == n_expert_used: Each slot has its own activation column
+// - General case: col_idx = slot % src1.ne[1] (modulo handles both cases)
+//
+// MATH NOTATION:
+// C[m, s, b] = Sum(k=0 to K-1) A[k, m, ids[s,b]] x B[k, s % ne11, b]
+// where:
+// m: [0, M) - output feature index
+// s: [0, n_expert_used) - expert slot index
+// b: [0, batch) - batch index
+// k: [0, K) - hidden dimension
+// ne11 = src1->ne[1] - number of columns in src1
+//******************************************************************************
+
+#include "block_ops.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+#include "quants.h"
+
+#include <stdint.h>
+
+// Main entry point for MUL_MAT_ID kernel (Mixture of Experts)
+int entry_point(struct ggml_et_mul_mat_id_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ // Get thread coordination info
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return -1;
+ }
+
+ // Use even threads only to avoid resource contention
+ // Each minion has 2 threads sharing instruction/data cache, NOC to RAM, and FPU
+ // Odd threads return immediately to avoid fighting for shared resources
+ if (thread_id & 1) {
+ return 0; // Odd thread - skip work
+ }
+
+ // Adjust thread count and ID for even-only threading
+ int effective_thread_id = thread_id / 2;
+ int effective_num_threads = (num_threads + 1) / 2; // Ceiling division
+
+ // Validate params
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ // Extract tensor references
+ struct ggml_tensor * src0 = ¶ms->src0; // Expert weight matrices [K, M, n_expert]
+ struct ggml_tensor * src1 = ¶ms->src1; // Activations [K, n_expert_used, batch]
+ struct ggml_tensor * src2 = ¶ms->src2; // Expert indices [n_expert_used, batch] (I32)
+ struct ggml_tensor * dst = ¶ms->dst; // Output [M, n_expert_used, batch, 1]
+
+ // Validate tensor types
+ if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || src2->type != GGML_TYPE_I32) {
+ return -1;
+ }
+
+ // Get data pointers
+ const void * src0_data = src0->data; // Expert matrices (Q8_0/F16/F32)
+ const float * src1_data = (const float *) src1->data; // Activations (F32)
+ const int32_t * src2_data = (const int32_t *) src2->data; // Expert IDs (I32)
+ float * dst_data = (float *) dst->data; // Output (F32)
+
+ if (!src0_data || !src1_data || !src2_data || !dst_data) {
+ return -1;
+ }
+
+ // Determine block size based on src0 type
+ int block_size;
+ switch (src0->type) {
+ case GGML_TYPE_Q8_0:
+ block_size = QK8_0;
+ break;
+ case GGML_TYPE_Q4_0:
+ block_size = QK4_0;
+ break;
+ case GGML_TYPE_F16:
+ block_size = QK_F16;
+ break;
+ case GGML_TYPE_F32:
+ block_size = QK_F32;
+ break;
+ default:
+ return -1;
+ }
+
+ // Get dimensions
+ // src0: [K, M, n_expert] - expert weight matrices
+ // src1: [K, n_expert_used, batch] - activations
+ // src2: [n_expert_used, batch] - expert indices
+ // dst: [M, n_expert_used, batch, 1] - output
+ const int64_t K = src0->ne[0]; // Hidden dimension
+ const int64_t M = src0->ne[1]; // Output features
+ const int64_t n_expert = src0->ne[2]; // Number of experts
+ const int64_t n_expert_used = src2->ne[0]; // Experts used per token
+ const int64_t batch = src2->ne[1]; // Batch size
+
+ // Strides (in bytes)
+ const size_t nb01 = src0->nb[1]; // src0 row stride
+ const size_t nb02 = src0->nb[2]; // src0 expert stride
+ const size_t nb11 = src1->nb[1]; // src1 column stride
+ const size_t nb12 = src1->nb[2]; // src1 batch stride
+ const size_t nb20 = src2->nb[0]; // src2 element stride
+ const size_t nb21 = src2->nb[1]; // src2 batch stride
+ const size_t nb1 = dst->nb[1]; // dst column stride
+ const size_t nb2 = dst->nb[2]; // dst batch stride
+
+ // Verify K dimension alignment for quantization
+ // Q8_0 requires strict alignment (quantized data must be block-aligned)
+ // F32 and F16 can handle partial blocks with scalar remainders
+ if ((src0->type == GGML_TYPE_Q8_0 || src0->type == GGML_TYPE_Q4_0) && K % block_size != 0) {
+ return -1; // Q8_0 requires K to be multiple of block_size
+ }
+
+ // Verify first dimension is contiguous
+ size_t expected_element_size_src0;
+ if (src0->type == GGML_TYPE_Q8_0) {
+ expected_element_size_src0 = sizeof(block_q8_0);
+ } else if (src0->type == GGML_TYPE_Q4_0) {
+ expected_element_size_src0 = sizeof(block_q4_0);
+ } else if (src0->type == GGML_TYPE_F16) {
+ expected_element_size_src0 = sizeof(uint16_t);
+ } else if (src0->type == GGML_TYPE_F32) {
+ expected_element_size_src0 = sizeof(float);
+ } else {
+ return -1;
+ }
+
+ if (src0->nb[0] != expected_element_size_src0 || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(int32_t) ||
+ dst->nb[0] != sizeof(float)) {
+ return -1;
+ }
+
+ const int64_t K_blocks = K / block_size;
+
+ // Threading: distribute output elements across threads
+ // Total output elements = M * n_expert_used * batch
+ const uint64_t total_elements = M * n_expert_used * batch;
+
+ const uint64_t per_thread = 16;
+ const uint64_t threads_stride = per_thread * effective_num_threads;
+
+ if (effective_thread_id * per_thread >= total_elements) {
+ return 0;
+ }
+
+ // Process elements assigned to this thread
+ for (uint64_t base_idx = effective_thread_id * per_thread; base_idx < total_elements; base_idx += threads_stride) {
+ for (uint64_t j = 0; j < per_thread; j++) {
+ const uint64_t idx = base_idx + j;
+
+ if (idx >= total_elements) {
+ break;
+ }
+
+ // Decode linear index to (m, n_idx, batch_idx)
+ // Layout: m + M * (n_idx + n_expert_used * batch_idx)
+ const int64_t batch_idx = idx / (M * n_expert_used);
+ const int64_t rem = idx % (M * n_expert_used);
+ const int64_t n_idx = rem / M;
+ const int64_t m = rem % M;
+
+ // Get expert ID from src2[n_idx, batch_idx]
+ const int32_t expert_id = *(const int32_t *) ((const char *) src2_data + n_idx * nb20 + batch_idx * nb21);
+
+ // Validate expert ID
+ if (expert_id < 0 || expert_id >= n_expert) {
+ // Invalid expert ID - write zero and continue
+ volatile float * dst_element =
+ (volatile float *) ((char *) dst_data + m * dst->nb[0] + n_idx * nb1 + batch_idx * nb2);
+ atomic_store_f32(dst_element, 0.0f);
+ continue;
+ }
+
+ // Compute dot product: expert_matrix[m, :] x activations[:, col_idx, batch_idx]
+ // Use modulo to handle broadcasting: when src1 has fewer columns than expert slots,
+ // multiple slots share the same activation column (col_idx = n_idx % src1->ne[1])
+ const int64_t col_idx = n_idx % src1->ne[1];
+ float sum = 0.0f;
+
+ // Type switch hoisted outside block loop: one branch per element, not per block
+ const char * expert_row_base = (const char *) src0_data + m * nb01 + expert_id * nb02;
+
+ switch (src0->type) {
+ case GGML_TYPE_Q8_0:
+ {
+ const block_q8_0 * q8_row = (const block_q8_0 *) expert_row_base;
+ const float * b_col_base =
+ (const float *) ((const char *) src1_data + col_idx * nb11 + batch_idx * nb12);
+ sum += compute_row_dot_q8_0(q8_row, b_col_base, K_blocks);
+ break;
+ }
+ case GGML_TYPE_Q4_0:
+ {
+ const block_q4_0 * q4_row = (const block_q4_0 *) expert_row_base;
+ const float * b_col_base =
+ (const float *) ((const char *) src1_data + col_idx * nb11 + batch_idx * nb12);
+ sum += compute_row_dot_q4_0(q4_row, b_col_base, K_blocks);
+ break;
+ }
+ case GGML_TYPE_F16:
+ {
+ const uint16_t * f16_row = (const uint16_t *) expert_row_base;
+ const int64_t K_remainder = K % block_size;
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const float * b_col_ptr =
+ (const float *) ((const char *) src1_data + (kb * block_size) * sizeof(float) +
+ col_idx * nb11 + batch_idx * nb12);
+ sum += compute_block_dot_product_f16_naive(&f16_row[kb * block_size], b_col_ptr);
+ }
+ if (K_remainder > 0) {
+ const int64_t offset = K_blocks * block_size;
+ const float * b_col_ptr =
+ (const float *) ((const char *) src1_data + offset * sizeof(float) + col_idx * nb11 +
+ batch_idx * nb12);
+ sum += compute_block_dot_product_f16_partial(&f16_row[offset], b_col_ptr, K_remainder);
+ }
+ break;
+ }
+ case GGML_TYPE_F32:
+ {
+ const float * f32_row = (const float *) expert_row_base;
+ const int64_t K_remainder = K % block_size;
+ for (int64_t kb = 0; kb < K_blocks; kb++) {
+ const float * b_col_ptr =
+ (const float *) ((const char *) src1_data + (kb * block_size) * sizeof(float) +
+ col_idx * nb11 + batch_idx * nb12);
+ sum += compute_block_dot_product_f32(&f32_row[kb * block_size], b_col_ptr);
+ }
+ if (K_remainder > 0) {
+ const int64_t offset = K_blocks * block_size;
+ const float * b_col_ptr =
+ (const float *) ((const char *) src1_data + offset * sizeof(float) + col_idx * nb11 +
+ batch_idx * nb12);
+ sum += compute_block_dot_product_f32_partial(&f32_row[offset], b_col_ptr, K_remainder);
+ }
+ break;
+ }
+ default:
+ return -1;
+ }
+
+ // Store result using atomic store to avoid cache coherency issues
+ // when multiple threads write to the same cache line (64 bytes = 16 floats)
+ volatile float * dst_element =
+ (volatile float *) ((char *) dst_data + m * dst->nb[0] + n_idx * nb1 + batch_idx * nb2);
+ atomic_store_f32(dst_element, sum);
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Norm F32 Kernel (Layer Normalization)
+// y[i] = (x[i] - mean) / sqrt(variance + eps)
+// where mean = sum(x) / N, variance = sum((x - mean)^2) / N
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <stdint.h>
+#include <string.h>
+
+// Norm kernel parameters structure
+struct ggml_et_norm_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon parameter for numerical stability
+};
+
+int entry_point(struct ggml_et_norm_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+ float eps = params->eps;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+
+ if (eps < 0.0f) {
+ return -1; // Invalid epsilon
+ }
+
+ const int64_t ne0 = dst->ne[0];
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t ne3 = dst->ne[3];
+
+ const size_t nb0 = dst->nb[0], nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+ const size_t nb00 = src0->nb[0], nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+
+ if (src0->ne[0] != ne0 || src0->ne[1] != ne1 || src0->ne[2] != ne2 || src0->ne[3] != ne3) {
+ return -1; // Shape mismatch
+ }
+
+ const int32_t total_rows = (int32_t) (ne1 * ne2 * ne3);
+ const int shire_threads = SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+
+ if (total_rows >= shire_threads) {
+ // Row-parallel: each thread processes whole rows
+ for (int64_t i3 = 0; i3 < ne3; i3++) {
+ for (int64_t i2 = 0; i2 < ne2; i2++) {
+ for (int64_t i1 = thread_id; i1 < ne1; i1 += num_threads) {
+ const float * src_ptr =
+ (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ // Step 1: sum for mean
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fadd.ps f10, f10, f11\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ float sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ const float mean = et_fdiv(sum, (float) (int32_t) ne0);
+
+ // Step 2: compute (x - mean) → dst, accumulate variance
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fbc.ps f12, %[mean_ptr]\n"
+ "fsub.ps f13, f11, f12\n"
+ "fsw.ps f13, %[result]\n"
+ "fmadd.ps f10, f13, f13, f10\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0]), [mean_ptr] "m"(mean)
+ : "f10", "f11", "f12", "f13");
+ }
+
+ float var_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(var_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ const float variance = et_fdiv(var_sum, (float) (int32_t) ne0);
+ const float scale = et_powf(variance + eps, -0.5f);
+
+ if (!(scale > 0.0f)) {
+ return -1;
+ }
+
+ // Step 3: apply scale to centered values in dst
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[y_vec]\n"
+ "fbc.ps f13, %[scale_ptr]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [y_vec] "m"(*(const float (*)[8]) & dst_ptr[i0]), [scale_ptr] "m"(scale)
+ : "f12", "f13", "f14");
+ }
+ }
+ }
+ }
+ } else {
+ // Intra-row: threads within each shire cooperate via L2 SCP.
+ // Two reductions needed: sum (for mean), then variance sum.
+ int shire_tid = thread_id % shire_threads;
+ int threads_per_row = shire_threads / total_rows;
+ int my_row = shire_tid / threads_per_row;
+ int local_tid = shire_tid % threads_per_row;
+ int group_base = my_row * threads_per_row;
+
+ if (my_row >= total_rows) {
+ FENCE;
+ et_barrier(ET_BARRIER_SHIRE);
+ // Second barrier for variance exchange
+ FENCE;
+ et_barrier(ET_BARRIER_SHIRE);
+ return 0;
+ }
+
+ int64_t i1 = my_row % ne1;
+ int64_t i2 = (my_row / ne1) % ne2;
+ int64_t i3 = my_row / (ne1 * ne2);
+
+ const float * src_ptr = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ const int32_t elems_per_cl = 16;
+ int32_t total_cls = ((int32_t) ne0 + elems_per_cl - 1) / elems_per_cl;
+ int32_t cls_per_thread = (total_cls + threads_per_row - 1) / threads_per_row;
+ int32_t my_start = local_tid * cls_per_thread * elems_per_cl;
+ int32_t my_end = my_start + cls_per_thread * elems_per_cl;
+ if (my_end > (int32_t) ne0) {
+ my_end = (int32_t) ne0;
+ }
+ if (my_start >= (int32_t) ne0) {
+ my_start = 0;
+ my_end = 0;
+ }
+
+ int workers = threads_per_row < total_cls ? threads_per_row : total_cls;
+
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // ---- Reduction 1: partial sum for mean ----
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fadd.ps f10, f10, f11\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ float partial_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(partial_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ // L2SCP exchange for sum
+ volatile float * my_slot = (volatile float *) et_shire_l2scp_local((uint64_t) shire_tid * 64);
+ *my_slot = partial_sum;
+ FENCE;
+ evict_to_l2((const void *) my_slot, 1, 64);
+ WAIT_CACHEOPS;
+
+ et_barrier(ET_BARRIER_SHIRE);
+
+ // All threads read sum, compute mean
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ evict_to_l2((const void *) slot, 1, 64);
+ }
+ WAIT_CACHEOPS;
+
+ float total_sum = 0.0f;
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ total_sum += *slot;
+ }
+
+ const float mean = et_fdiv(total_sum, (float) (int32_t) ne0);
+
+ // ---- Reduction 2: compute (x - mean) → dst chunk, partial variance ----
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+
+ if (my_start < my_end) {
+ uint32_t mean_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(mean_bits) : "f"(mean));
+ __asm__ volatile("fbcx.ps f15, %[mb]\n" : : [mb] "r"(mean_bits) : "f15");
+
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fsub.ps f13, f11, f15\n"
+ "fsw.ps f13, %[result]\n"
+ "fmadd.ps f10, f13, f13, f10\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11", "f13");
+ }
+ }
+
+ float partial_var;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(partial_var)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ // L2SCP exchange for variance (reuse same slots)
+ *my_slot = partial_var;
+ FENCE;
+ evict_to_l2((const void *) my_slot, 1, 64);
+ WAIT_CACHEOPS;
+
+ et_barrier(ET_BARRIER_SHIRE);
+
+ // All threads read variance, compute scale, apply to own chunk
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ evict_to_l2((const void *) slot, 1, 64);
+ }
+ WAIT_CACHEOPS;
+
+ float total_var = 0.0f;
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ total_var += *slot;
+ }
+
+ const float variance = et_fdiv(total_var, (float) (int32_t) ne0);
+ const float scale = et_powf(variance + eps, -0.5f);
+
+ if (!(scale > 0.0f)) {
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return -1;
+ }
+
+ // Apply scale to centered values (already in dst from reduction 2)
+ if (my_start < my_end) {
+ uint32_t scale_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(scale_bits) : "f"(scale));
+ __asm__ volatile("fbcx.ps f13, %[sb]\n" : : [sb] "r"(scale_bits) : "f13");
+
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[y_vec]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [y_vec] "m"(*(const float (*)[8]) & dst_ptr[i0])
+ : "f12", "f14");
+ }
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal PAD F32 Kernel
+// Zero-pads an F32 tensor along dimensions 1-3.
+//
+// Constraints:
+// - No dim0 padding (lp[0]==0, rp[0]==0)
+// - dst contiguous
+// - src nb[0] == 4 (dim0 contiguous for vectorized reads)
+// - Zero-pad only (no circular mode)
+//
+// Two paths:
+// Aligned (ne0 % 16 == 0): rows distributed across harts, vectorized.
+// Small (16 % ne0 == 0): cache-line distributed, scalar per-element.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_pad_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+ int32_t lp[4];
+ int32_t rp[4];
+};
+
+// Vectorized copy with scalar tail
+static inline void vec_copy_f32(float * dst, const float * src, int32_t n) {
+ int32_t i = 0;
+ const int32_t vec_end = (n / 8) * 8;
+ for (; i < vec_end; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[s]\n"
+ "fsw.ps f10, %[d]\n"
+ : [d] "=m"(*(float (*)[8]) & dst[i])
+ : [s] "m"(*(const float (*)[8]) & src[i])
+ : "f10");
+ }
+ for (; i < n; i++) {
+ dst[i] = src[i];
+ }
+}
+
+int entry_point(struct ggml_et_pad_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * src0_data = (const float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ // Dst dimensions
+ const int64_t ne0 = dst->ne[0];
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t ne3 = dst->ne[3];
+
+ // Src strides (byte offsets)
+ const int64_t nb1_src = src0->nb[1];
+ const int64_t nb2_src = src0->nb[2];
+ const int64_t nb3_src = src0->nb[3];
+
+ // Padding values
+ const int32_t lp1 = params->lp[1];
+ const int32_t rp1 = params->rp[1];
+ const int32_t lp2 = params->lp[2];
+ const int32_t rp2 = params->rp[2];
+ const int32_t lp3 = params->lp[3];
+ const int32_t rp3 = params->rp[3];
+
+ const int64_t total_rows = ne1 * ne2 * ne3;
+ const int64_t total_elements = ne0 * total_rows;
+
+ if (total_elements == 0) {
+ return 0;
+ }
+
+ // Broadcast 0.0f to SIMD register for vectorized zero-fill
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f12, %[v]\n" : : [v] "m"(zero) : "f12");
+
+ // Aligned: ne0 % 16 == 0 -> row-based distribution, vectorized
+ if (ne0 % 16 == 0) {
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ const int64_t i3 = row / (ne1 * ne2);
+ const int64_t i2 = (row / ne1) % ne2;
+ const int64_t i1 = row % ne1;
+
+ float * dst_row = dst_data + row * ne0;
+
+ if (i1 >= lp1 && i1 < ne1 - rp1 && i2 >= lp2 && i2 < ne2 - rp2 && i3 >= lp3 && i3 < ne3 - rp3) {
+ const float * src_row = (const float *) ((const char *) src0_data + (i1 - lp1) * nb1_src +
+ (i2 - lp2) * nb2_src + (i3 - lp3) * nb3_src);
+ vec_copy_f32(dst_row, src_row, (int32_t) ne0);
+ } else {
+ int64_t i = 0;
+ const int64_t vec_end = (ne0 / 8) * 8;
+ for (; i < vec_end; i += 8) {
+ __asm__ volatile("fsw.ps f12, %[d]\n" : [d] "=m"(*(float (*)[8]) & dst_row[i])::"f12");
+ }
+ }
+ }
+ return 0;
+ }
+
+ // Small-ne0 path: 16 % ne0 == 0 -> cache-line distributed, scalar
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t ne1_data_end = ne1 - rp1;
+ const int64_t ne2_data_end = ne2 - rp2;
+ const int64_t ne3_data_end = ne3 - rp3;
+
+ for (int64_t cl = thread_id; cl < total_cl; cl += num_threads) {
+ const int64_t elem_start = cl * elems_per_cl;
+ int64_t elem_end = elem_start + elems_per_cl;
+ if (elem_end > total_elements) {
+ elem_end = total_elements;
+ }
+
+ for (int64_t idx = elem_start; idx < elem_end; idx++) {
+ const int64_t i0 = idx % ne0;
+ const int64_t rem = idx / ne0;
+ const int64_t i1 = rem % ne1;
+ const int64_t rem2 = rem / ne1;
+ const int64_t i2 = rem2 % ne2;
+ const int64_t i3 = rem2 / ne2;
+
+ if (i1 >= lp1 && i1 < ne1_data_end && i2 >= lp2 && i2 < ne2_data_end && i3 >= lp3 && i3 < ne3_data_end) {
+ const float * sp = (const float *) ((const char *) src0_data + i0 * 4 + (i1 - lp1) * nb1_src +
+ (i2 - lp2) * nb2_src + (i3 - lp3) * nb3_src);
+ dst_data[idx] = *sp;
+ } else {
+ dst_data[idx] = 0.0f;
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// ET Platform Hardware Abstraction Layer
+// Provides thread coordination, kernel infrastructure, and platform primitives
+// for bare metal ET kernels
+//******************************************************************************
+
+#ifndef PLATFORM_H
+#define PLATFORM_H
+
+#include "etsoc/common/utils.h"
+#include "etsoc/isa/barriers.h"
+#include "etsoc/isa/cacheops-umode.h"
+#include "etsoc/isa/hart.h"
+
+#include <stdint.h>
+
+#define SOC_MINIONS_PER_SHIRE 32
+#define NUM_HARTS_PER_MINION 2
+#define ET_CACHE_LINE_SIZE_BYTES 64
+
+// Environment structure definition
+typedef struct {
+ uint32_t version; // Version of the ABI (offset 0)
+ uint32_t padding1; // Padding to align shire_mask to offset 8
+ uint64_t shire_mask; // Bitmask of active compute shires (offset 8)
+ uint32_t frequency; // Frequency of Minion cores in MHz (offset 16)
+ uint32_t padding2; // Padding to maintain alignment
+} __attribute__((packed, aligned(64))) kernel_environment_t;
+
+// Manual implementation of count trailing zeros for bare metal environment
+// NOTE: This simple loop-based implementation is used for portability.
+// Production implementations (like libgcc's __ctzdi2) use optimized bit manipulation
+// algorithms with lookup tables and parallel bit operations for O(log n) performance.
+static inline int manual_ctzll(uint64_t x) {
+ if (x == 0) return 64;
+ int count = 0;
+ while ((x & 1) == 0) {
+ x >>= 1;
+ count++;
+ }
+ return count;
+}
+
+// Manual implementation of population count for bare metal environment
+// NOTE: This simple loop-based implementation is used for portability.
+// Production implementations (like libgcc's __popcountdi2) use optimized bit-parallel
+// algorithms with magic constants and bit manipulation tricks for O(1) performance.
+static inline int manual_popcountll(uint64_t x) {
+ int count = 0;
+ while (x) {
+ count += x & 1;
+ x >>= 1;
+ }
+ return count;
+}
+
+// Binary GCD (Stein's algorithm) — avoids expensive 64-bit division/remainder.
+// Uses only shifts, subtraction, and comparison (all single-cycle on ET cores).
+static inline int64_t et_gcd_i64(int64_t a, int64_t b) {
+ while (b) {
+ const int64_t t = b;
+ b = a % b;
+ a = t;
+ }
+ return a;
+}
+
+// Return the number of consecutive rows of width row_elems needed so the
+// combined write footprint spans an integer number of cache lines.
+static inline int64_t et_rows_per_cacheline_group(int64_t row_elems, int64_t elem_size_bytes) {
+ if (row_elems <= 0 || elem_size_bytes <= 0) {
+ return 1;
+ }
+
+ const int64_t row_bytes = row_elems * elem_size_bytes;
+ const int64_t gcd = et_gcd_i64(ET_CACHE_LINE_SIZE_BYTES, row_bytes);
+ return ET_CACHE_LINE_SIZE_BYTES / gcd;
+}
+
+// Calculate relative thread ID from absolute hart ID using shire mask
+// Returns -1 if this hart is not active (not in shire mask)
+static inline int get_relative_thread_id(uint64_t shire_mask) {
+ int hart_id = (int) get_hart_id();
+
+ // Find starting hart offset from lowest active shire
+ int starting_hart = manual_ctzll(shire_mask) * SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+
+ // Return -1 if not an active thread
+ if (hart_id < starting_hart) {
+ return -1;
+ }
+
+ // Calculate relative thread ID
+ int thread_id = hart_id - starting_hart;
+ return thread_id;
+}
+
+// Calculate total number of threads from shire mask
+static inline int get_num_threads(uint64_t shire_mask) {
+ // Count active shires using popcount, multiply by minions per shire and harts per minion
+ return manual_popcountll(shire_mask) * SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+}
+
+//******************************************************************************
+// Synchronization Primitives
+//******************************************************************************
+
+#define NOP __asm__ __volatile__("nop\n");
+#define FENCE __asm__ __volatile__("fence\n" ::: "memory");
+#define WFI __asm__ __volatile__("wfi\n");
+
+//******************************************************************************
+// Atomic Operations
+//******************************************************************************
+
+// Global AMO primitives — ET custom 'g' suffix instructions that go through
+// the NoC coherence fabric for chip-wide atomicity.
+
+// Atomic swap (word), returns previous value.
+static inline uint32_t __attribute__((always_inline)) et_global_swap_w(volatile void * addr, uint32_t val) {
+ uint32_t ret;
+ __asm__ __volatile__("amoswapg.w %0, %1, (%2)" : "=r"(ret) : "r"(val), "r"(addr) : "memory");
+ return ret;
+}
+
+// Atomic add (word), returns previous value.
+static inline uint32_t __attribute__((always_inline)) et_global_add_w(volatile void * addr, uint32_t val) {
+ uint32_t ret;
+ __asm__ __volatile__("amoaddg.w %0, %1, (%2)" : "=r"(ret) : "r"(val), "r"(addr) : "memory");
+ return ret;
+}
+
+// Atomic store (halfword, global). Address must be 16-bit aligned.
+static inline void __attribute__((always_inline)) et_global_store_hw(volatile void * addr, uint16_t val) {
+ __asm__ __volatile__("shg %0, (%1)" : : "r"(val), "r"(addr) : "memory");
+}
+
+// Convenience wrappers — float types, fire-and-forget (old value discarded).
+static inline void atomic_store_f32(volatile float * addr, float value) {
+ et_global_swap_w(addr, *(uint32_t *) &value);
+}
+
+static inline void atomic_add_f32(volatile float * addr, float value) {
+ et_global_add_w(addr, *(uint32_t *) &value);
+}
+
+static inline void atomic_store_f16(volatile uint16_t * addr, uint16_t value) {
+ et_global_store_hw(addr, value);
+}
+
+//******************************************************************************
+// Barrier Primitives
+//
+// Hardware resources used (per shire):
+// - 32 FLBs: 8-bit atomic counters, non-blocking (CSR 0x820)
+// - 2 FCCs per hart: credit counters, hardware-stall on consume (CSR 0x821)
+//
+// Convention:
+// MINION barriers: FLB = local_minion_id (0-31), FCC 0
+// SHIRE barriers: FLB 0, FCC 1
+//
+// MINION and SHIRE barriers MUST NOT be concurrent. All minion barriers
+// must complete before a shire barrier, and vice versa. FLB 0 is shared
+// between minion 0's barrier and the shire barrier — safe only because
+// the FLB counter auto-resets on match.
+//
+// FCC 0 is safe for all 32 concurrent minion barriers because each
+// barrier's fcc_send targets only its own minion (per-hart private
+// counters, scoped by CREDINC mask). FCC 1 is reserved for shire-wide
+// broadcast.
+//******************************************************************************
+
+#define ET_DEFAULT_SHIRE_MASK 0xFFFFFFFFULL
+
+typedef enum {
+ ET_BARRIER_MINION, // sync both harts within each minion (FLB=minion_id, FCC 0)
+ ET_BARRIER_SHIRE, // sync all harts across the shire (FLB=0, FCC 1)
+ ET_BARRIER_GLOBAL, // sync all harts across all active shires (FLB+global AMO+FCC)
+} et_barrier_scope_t;
+
+//******************************************************************************
+// Global Barrier (cross-shire)
+//
+// Synchronizes all harts across multiple shires on the chip.
+// Algorithm:
+// 1. FLB within each shire to elect one representative hart
+// 2. Elected hart does a global atomic increment on a shared counter
+// 3. The last shire to arrive resets the counter and sends FCC credits
+// to all active shires to release them
+// 4. All harts wait on FCC to complete the barrier
+//
+// Uses FLB 0, FCC 1 (same as ET_BARRIER_SHIRE, these must not overlap).
+// The counter lives in a cache-line-aligned global to avoid coherency problems
+//******************************************************************************
+
+// Barrier counter cache-line aligned to avoid coherency problems
+// Must be zero-initialized (BSS).
+static uint32_t __attribute__((aligned(64))) et_global_barrier_count[64 / sizeof(uint32_t)] = { 0 };
+
+// Cross-shire barrier: all harts in num_active_shires shires synchronize.
+// Returns 1 if this hart was the globally-last to arrive, 0 otherwise.
+//
+// num_active_shires - number of shires participating
+// (typically popcount(shire_mask) from kernel_environment_t)
+static inline uint64_t __attribute__((always_inline)) et_barrier_global(uint64_t num_active_shires) {
+ uint64_t last_global = 0;
+
+ // FLB within this shire. Elect one hart per shire.
+ // Master shire has only 16 minions (32 harts), others have 32 (64 harts).
+ uint64_t shire_id = get_shire_id();
+ uint32_t harts_in_shire = (shire_id == SHIRE_MASTER) ? (SOC_MINIONS_PER_SHIRE / 2) * NUM_HARTS_PER_MINION :
+ SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+ uint64_t last_in_shire = flbarrier(0, harts_in_shire - 1);
+
+ if (last_in_shire) {
+ // Global atomic increment. Count arriving shires
+ uint32_t prev = et_global_add_w(et_global_barrier_count, 1);
+
+ if (prev == num_active_shires - 1) {
+ // Last shire. reset counter and fan out FCC to all shires
+ last_global = 1;
+ et_global_swap_w(et_global_barrier_count, 0);
+
+ for (uint64_t sid = 0; sid < 33; sid++) {
+ // Send FCC 1 credit to all harts (both threads) in each shire
+ fcc_send(sid, THREAD_0, FCC_1, 0xFFFFFFFF);
+ fcc_send(sid, THREAD_1, FCC_1, 0xFFFFFFFF);
+ }
+ }
+ }
+
+ // All harts wait for the FCC credit from the last shire
+ fcc_consume(FCC_1);
+ return last_global;
+}
+
+// Barrier with scope-derived parameters.
+// Returns 1 if this hart was the last to arrive, 0 otherwise.
+//
+// ET_BARRIER_GLOBAL uses ET_DEFAULT_SHIRE_MASK (32 shires). For a different
+// shire count, use et_barrier_global(n) directly.
+static inline uint64_t __attribute__((always_inline)) et_barrier(et_barrier_scope_t scope) {
+ if (scope == ET_BARRIER_MINION) {
+ uint32_t local_minion = (get_hart_id() >> 1) & 0x1F;
+ uint32_t mask = 1u << local_minion;
+ return shire_barrier(local_minion, 0, 2, mask, mask);
+ } else if (scope == ET_BARRIER_SHIRE) {
+ uint64_t shire_id = get_shire_id();
+ uint32_t thread_count = (shire_id == SHIRE_MASTER) ? 32 : 64;
+ uint32_t mask = (shire_id == SHIRE_MASTER) ? 0xFFFF0000U : 0xFFFFFFFFU;
+ return shire_barrier(0, 1, thread_count, mask, mask);
+ } else { /* ET_BARRIER_GLOBAL */
+ return et_barrier_global(manual_popcountll(ET_DEFAULT_SHIRE_MASK));
+ }
+}
+
+// Raw barrier — caller manages FLB/FCC allocation.
+// Use when et_barrier() doesn't fit (custom thread counts, subgroups,
+// only even harts active, etc).
+//
+// flb - which FLB counter (0-31)
+// fcc - which FCC counter (0 or 1)
+// thread_count - number of harts that will call this barrier
+// mask_t0 - CREDINC bitmask: which minions' hart 0 gets a credit
+// mask_t1 - CREDINC bitmask: which minions' hart 1 gets a credit
+static inline uint64_t __attribute__((always_inline)) et_barrier_raw(uint32_t flb,
+ uint32_t fcc,
+ uint32_t thread_count,
+ uint32_t mask_t0,
+ uint32_t mask_t1) {
+ return shire_barrier(flb, fcc, thread_count, mask_t0, mask_t1);
+}
+
+// One-way semaphore between harts (non-blocking post, blocking wait).
+//
+// et_sem_post(): increment the partner hart's semaphore. Non-blocking.
+// the caller continues immediately. Multiple posts accumulate.
+//
+// et_sem_wait(): block until the semaphore is non-zero, then decrement it.
+//
+// Backed by hardware FCC (Flow Control Credit) counters. Uses FCC 0 for
+// ET_BARRIER_MINION scope. Counters are per-hart private, so both harts
+// can post/wait on the same scope independently.
+//
+// Must not be mixed with et_barrier() of the same scope in the
+// same kernel (shared FCC channel).
+static inline void __attribute__((always_inline)) et_sem_post(et_barrier_scope_t scope) {
+ if (scope == ET_BARRIER_MINION) {
+ uint64_t hart_id = get_hart_id();
+ uint32_t local_minion = (hart_id >> 1) & 0x1F;
+ uint32_t mask = 1u << local_minion;
+ uint64_t shire_id = get_shire_id();
+
+ if (hart_id & 1) {
+ // Hart 1 → hart 0
+ fcc_send(shire_id, THREAD_0, FCC_0, mask);
+ } else {
+ // Hart 0 → hart 1
+ fcc_send(shire_id, THREAD_1, FCC_0, mask);
+ }
+ }
+}
+
+// Block until a post from et_sem_post() is available, then consume it.
+static inline void __attribute__((always_inline)) et_sem_wait(et_barrier_scope_t scope) {
+ if (scope == ET_BARRIER_MINION) {
+ fcc_consume(FCC_0);
+ }
+}
+
+//******************************************************************************
+// Tensor Engine Wait & Error Macros
+//
+// These write to CSR 0x830 (tensor_wait) to stall the hart until the specified
+// tensor unit completes its current operation. The immediate encodes which
+// unit to wait on.
+//******************************************************************************
+
+#define WAIT_TENSOR_LOAD_0 __asm__ __volatile__("csrwi 0x830, 0\n" : :);
+#define WAIT_TENSOR_LOAD_1 __asm__ __volatile__("csrwi 0x830, 1\n" : :);
+#define WAIT_TENSOR_LOAD_L2_0 __asm__ __volatile__("csrwi 0x830, 2\n" : :);
+#define WAIT_TENSOR_LOAD_L2_1 __asm__ __volatile__("csrwi 0x830, 3\n" : :);
+#define WAIT_PREFETCH_0 __asm__ __volatile__("csrwi 0x830, 4\n" : :);
+#define WAIT_PREFETCH_1 __asm__ __volatile__("csrwi 0x830, 5\n" : :);
+#define WAIT_CACHEOPS __asm__ __volatile__("csrwi 0x830, 6\n" : :);
+#define WAIT_TENSOR_FMA __asm__ __volatile__("csrwi 0x830, 7\n" : :);
+#define WAIT_TENSOR_STORE __asm__ __volatile__("csrwi 0x830, 8\n" : :);
+#define WAIT_TENSOR_REDUCE __asm__ __volatile__("csrwi 0x830, 9\n" : :);
+#define WAIT_TENSOR_QUANT __asm__ __volatile__("csrwi 0x830, 10\n" : :);
+#define STALL __asm__ __volatile__("csrw stall, x0\n" : :);
+
+// Write 0 to CSR 0x808 (tensor_error) to clear any latched tensor error bits.
+// Must be issued before the first tensor operation in a kernel to avoid stale
+// errors from a previous invocation causing spurious faults.
+#define CLEAR_TENSOR_ERROR __asm__ __volatile__("csrwi 0x808, 0" : :);
+
+//******************************************************************************
+// L1 Data Cache / Scratchpad (SCP) Configuration
+//
+// The ET-SoC-1 L1 data cache can be split so that half its ways operate as a
+// software-managed scratchpad (SCP). Tensor load/store/FMA instructions
+// require SCP mode to be active.
+//
+// CSR 0x810 — ucache_control:
+//
+// Bit(s) Field Description
+// ────── ──────────── ──────────────────────────────────────────────────
+// [0] D1Split 1 = L1 is split (half cache, half SCP).
+// Read-only from U-mode; set by M-mode firmware
+// before kernel launch. Writing ScpEnable while
+// D1Split=0 is silently ignored.
+// [1] ScpEnable 1 = scratchpad is active and zeroed.
+// [4:2] RepRate Cache-op replay rate (0 = no delay between ops).
+// [10:6] CacheOpMax Max outstanding cache ops (0 = unlimited).
+//
+// Typical kernel prologue for tensor operations:
+// setup_cache_scp(); // enables SCP, waits for zeroing
+// CLEAR_TENSOR_ERROR; // clear stale error bits
+//******************************************************************************
+
+// Write the ucache_control CSR (0x810).
+//
+// scp_en — 1 to enable SCP mode (requires D1Split already set)
+// cacheop_rate — cache-op replay rate (0–7; 0 = no delay)
+// cacheop_max — max outstanding cache ops (0–31; 0 = unlimited)
+static inline void __attribute__((always_inline)) ucache_control(uint64_t scp_en,
+ uint64_t cacheop_rate,
+ uint64_t cacheop_max) {
+ uint64_t csr_enc = ((cacheop_max & 0x1F) << 6) | ((cacheop_rate & 0x7) << 2) | ((scp_en & 0x1) << 1);
+
+ __asm__ __volatile__("csrw 0x810, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) : "x31");
+}
+
+// Enable L1 scratchpad mode and wait for the transition to complete.
+// After this call the SCP lines are zeroed and ready for tensor operations.
+//
+// Prerequisites:
+// - D1Split must already be 1 (set by M-mode firmware at boot).
+// - Only even harts (hart 0 per minion) should call this, as only they
+// can issue tensor instructions.
+static inline void setup_cache_scp(void) {
+ FENCE; // drain pending stores before reconfiguring cache
+ ucache_control(1, 0, 0); // ScpEnable=1
+ WAIT_CACHEOPS; // wait for SCP mode transition + zeroing
+}
+
+//******************************************************************************
+// L2 Scratchpad (L2 SCP) Address Computation
+//
+// Each shire has 4 MB of SRAM that can be split across L2 cache, L3 cache,
+// and scratchpad. The scratchpad region occupies 0x00_8000_0000~0x00_FFFF_FFFF
+// and is accessible via regular load/store from any minion core.
+//
+// Two addressing formats (differentiated by address bit 30):
+//
+// Format 0 (bit[30]=0): Direct shire addressing
+// [29:23] = shire ID (0–33, or 0x7F for local shire)
+// [22:0] = byte offset within shire's scratchpad
+//
+// Format 1 (bit[30]=1): Striped (round-robin) addressing
+// [29:28] = shire ID[6:5]
+// [27:11] = offset[22:6] (cache-line-aligned upper bits)
+// [10:6] = shire ID[4:0]
+// [5:0] = offset[5:0] (byte within cache line)
+// Consecutive 64-byte cache lines cycle through different shires,
+// distributing bandwidth across the mesh.
+//
+// Shire ID 0x7F always targets the local shire (instead of figureing out which
+// shire you are on).
+//******************************************************************************
+
+#define L2SCP_BASE 0x0080000000ULL
+#define L2SCP_SHIRE_LOCAL 0x7FULL
+
+// Format 0: direct address into a specific shire's L2 SCP.
+// shire: 0–33 for explicit shire, L2SCP_SHIRE_LOCAL (0x7F) for local
+// offset: byte offset within the shire's scratchpad
+static inline void * __attribute__((always_inline)) et_shire_l2scp(uint64_t shire, uint64_t offset) {
+ return (void *) (L2SCP_BASE | ((shire & 0x7F) << 23) | (offset & 0x7FFFFF));
+}
+
+// Format 0: local shire shorthand — no cross-shire traffic.
+static inline void * __attribute__((always_inline)) et_shire_l2scp_local(uint64_t offset) {
+ return (void *) (L2SCP_BASE | (L2SCP_SHIRE_LOCAL << 23) | (offset & 0x7FFFFF));
+}
+
+// Format 1: flat offset into a hardware-striped global address space.
+// Consecutive 64-byte cache lines automatically land on different shires,
+// distributing bandwidth across the mesh. No shire parameter — the
+// hardware derives the target shire from the address bits.
+static inline void * __attribute__((always_inline)) et_global_l2scp(uint64_t offset) {
+ return (void *) (L2SCP_BASE | (1ULL << 30) | (offset & 0x3FFFFFFF));
+}
+
+//******************************************************************************
+// Cache Operatons
+//******************************************************************************
+
+// Prefetch nlines cache lines into L2 starting at addr, with stride bytes
+// between each line. Uses PrefetchVA (CSR 0x81F) with dest=L2 (bits 59:58=01).
+//
+// The hardware fetches nlines consecutive cache-line-sized (64B) blocks from
+// DRAM/L3 into L2, starting at addr and advancing by stride bytes per line.
+// This is asynchronous — use WAIT_PREFETCH_0 or WAIT_PREFETCH_1 if the hart
+// must stall until the prefetch completes.
+//
+// NOTE: nlines is encoded in a 4-bit field (max 16). Passing nlines > 16
+// silently truncates. DO NOT pass nlines > 16.
+static inline void __attribute__((always_inline)) l2_prefetch(const void * addr, uint64_t nlines, uint64_t stride) {
+ uint64_t csr_val = (0x1ULL << 58) | ((uint64_t) addr & 0xFFFFFFFFFFC0ULL) | ((nlines - 1) & 0xF);
+
+ __asm__ __volatile__(
+ "mv x31, %[stride]\n"
+ "csrw 0x81f, %[val]\n"
+ :
+ : [stride] "r"(stride & 0xFFFFFFFFFFC0ULL), [val] "r"(csr_val)
+ : "x31", "memory");
+}
+
+// Flush nlines cache lines at stride apart starting at addr from L1 to L2.
+// Uses FlushVA (CSR 0x8BF). Caller must FENCE before (to drain stores to L1)
+// and WAIT_CACHEOPS after (to ensure flush completes before tensor loads).
+//
+// NOTE: nlines is encoded in a 4-bit field (max 16). Passing nlines > 16
+// silently truncates. DO NOT pass nlines > 16.
+static inline void __attribute__((always_inline)) flush_to_l2(const void * addr, uint64_t nlines, uint64_t stride) {
+ // dest=01 (L2) in bits 59:58, VA in bits 47:6, numlines-1 in bits 3:0
+ uint64_t csr_val = (0x1ULL << 58) | ((uint64_t) addr & 0xFFFFFFFFFFC0ULL) | ((nlines - 1) & 0xF);
+ uint64_t x31_val = stride & 0xFFFFFFFFFFC0ULL;
+
+ __asm__ __volatile__(
+ "mv x31, %[x31]\n"
+ "csrw 0x8BF, %[val]\n"
+ :
+ : [x31] "r"(x31_val), [val] "r"(csr_val)
+ : "x31", "memory");
+}
+
+// Evict nlines cache lines at stride apart starting at addr from L1 to L2.
+// Uses EvictVA (CSR 0x89F). Unlike flush_to_l2, this guarantees the line is
+// NOT present in L1 after the operation - subsequent loads will miss and go
+// to L2/SCP. Caller must FENCE before and WAIT_CACHEOPS after.
+//
+// NOTE: nlines is encoded in a 4-bit field (max 16). DO NOT pass nlines > 16.
+static inline void __attribute__((always_inline)) evict_to_l2(const void * addr, uint64_t nlines, uint64_t stride) {
+ // dest=01 (L2) in bits 59:58, VA in bits 47:6, numlines-1 in bits 3:0
+ uint64_t csr_val = (0x1ULL << 58) | ((uint64_t) addr & 0xFFFFFFFFFFC0ULL) | ((nlines - 1) & 0xF);
+ uint64_t x31_val = stride & 0xFFFFFFFFFFC0ULL;
+
+ __asm__ __volatile__(
+ "mv x31, %[x31]\n"
+ "csrw 0x89F, %[val]\n"
+ :
+ : [x31] "r"(x31_val), [val] "r"(csr_val)
+ : "x31", "memory");
+}
+
+// Evict nlines cache lines at stride apart starting at addr from BOTH L1
+// and L2. Uses EvictVA (CSR 0x89F) with dest=10 (L3/DRAM). Guarantees the
+// line is NOT present in L1 or L2 after the operation — subsequent loads
+// will fetch from L3 or DRAM. Needed because both L1 and L2 are incoherent
+// on ET-SoC-1 (L2 is per-shire).
+// Caller must FENCE before and WAIT_CACHEOPS after.
+//
+// NOTE: nlines is encoded in a 4-bit field (max 16). DO NOT pass nlines > 16.
+static inline void __attribute__((always_inline)) evict_past_l2(const void * addr, uint64_t nlines, uint64_t stride) {
+ // dest=10 in bits 59:58, VA in bits 47:6, numlines-1 in bits 3:0
+ uint64_t csr_val = (0x2ULL << 58) | ((uint64_t) addr & 0xFFFFFFFFFFC0ULL) | ((nlines - 1) & 0xF);
+ uint64_t x31_val = stride & 0xFFFFFFFFFFC0ULL;
+
+ __asm__ __volatile__(
+ "mv x31, %[x31]\n"
+ "csrw 0x89F, %[val]\n"
+ :
+ : [x31] "r"(x31_val), [val] "r"(csr_val)
+ : "x31", "memory");
+}
+
+// Evict a contiguous region from both L1 and L2 so subsequent loads fetch
+// from L3/DRAM. Both L1 and L2 are incoherent on ET-SoC-1 (L2 is per-shire),
+// so every op must evict its inputs before reading if a prior op in the same
+// uberkernel batch may have written to them via fsw.ps or tensor_store.
+//
+// Handles regions larger than the 16-line hardware limit by issuing multiple
+// evict_past_l2 calls.
+static void evict_region_past_l2(const void * addr, size_t bytes) {
+ if (!addr || bytes == 0) {
+ return;
+ }
+
+ const uint64_t CL = 64;
+ uint64_t base = (uint64_t) addr & ~(CL - 1);
+ uint64_t end = ((uint64_t) addr + bytes + CL - 1) & ~(CL - 1);
+ uint64_t nlines = (end - base) / CL;
+ // FENCE;
+ for (uint64_t off = 0; off < nlines; off += 16) {
+ uint64_t batch = nlines - off;
+ if (batch > 16) {
+ batch = 16;
+ }
+ evict_past_l2((const void *) (base + off * CL), batch, CL);
+ }
+}
+
+#endif // PLATFORM_H
--- /dev/null
+// Scalar dequantization helpers and ET-side block-size aliases.
+
+#ifndef QUANTS_H
+#define QUANTS_H
+
+#include "math_fp.h"
+
+#include <stdint.h>
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+
+// 64-byte (one cache line) F16 / F32 block sizes.
+#define QK_F16 32
+#define QK_F32 16
+
+static inline void dequantize_q8_0_block(const block_q8_0 * block, float * dst) {
+ const float scale = fp16_to_fp32(block->d);
+
+ for (int i = 0; i < QK8_0; i++) {
+ dst[i] = scale * (float) block->qs[i];
+ }
+}
+
+// Low nibbles -> dst[0..15], high nibbles -> dst[16..31].
+static inline void dequantize_q4_0_block(const block_q4_0 * block, float * dst) {
+ const float scale = fp16_to_fp32(block->d);
+
+ for (int i = 0; i < QK4_0 / 2; i++) {
+ const uint8_t byte = block->qs[i];
+ dst[i] = scale * (float) ((int) (byte & 0xF) - 8);
+ dst[i + QK4_0 / 2] = scale * (float) ((int) (byte >> 4) - 8);
+ }
+}
+
+// Unpack the 6-bit scale/min pair for Q4_K group j (groups 4-7 split their high bits).
+static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t * d, uint8_t * m) {
+ if (j < 4) {
+ *d = q[j] & 63;
+ *m = q[j + 4] & 63;
+ } else {
+ *d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
+ *m = (q[j + 4] >> 4) | ((q[j] >> 6) << 4);
+ }
+}
+
+static inline void dequantize_q4_K_block(const block_q4_K * block, float * dst) {
+ const uint8_t * q = block->qs;
+ const float d = fp16_to_fp32(block->d);
+ const float min = fp16_to_fp32(block->dmin);
+
+ int is = 0;
+ uint8_t sc, m;
+ for (int j = 0; j < QK_K; j += 64) {
+ get_scale_min_k4(is + 0, block->scales, &sc, &m);
+ const float d1 = d * sc;
+ const float m1 = min * m;
+ get_scale_min_k4(is + 1, block->scales, &sc, &m);
+ const float d2 = d * sc;
+ const float m2 = min * m;
+ for (int l = 0; l < 32; ++l) {
+ *dst++ = d1 * (q[l] & 0xF) - m1;
+ }
+ for (int l = 0; l < 32; ++l) {
+ *dst++ = d2 * (q[l] >> 4) - m2;
+ }
+ q += 32;
+ is += 2;
+ }
+}
+
+#endif // QUANTS_H
--- /dev/null
+//******************************************************************************
+// Repeat F32 Kernel
+// Tiles src0 into dst: dst.ne[i] = src0.ne[i] * nr[i] for each dimension.
+// All copies are cacheline-aligned (ne00 % 16 == 0).
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+#include <string.h>
+
+struct ggml_et_repeat_params {
+ struct ggml_tensor src0; // F32 input tensor (tile)
+ struct ggml_tensor dst; // F32 output tensor (tiled result)
+};
+
+// Copy n floats from src to dst using 8-wide vector loads/stores.
+// n must be a multiple of 16 (cacheline-aligned).
+static inline void copy_row_aligned(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[src_vec]\n"
+ "fsw.ps f11, %[dst_vec]\n"
+ : [dst_vec] "=m"(*(float (*)[8]) & dst[i])
+ : [src_vec] "m"(*(const float (*)[8]) & src[i])
+ : "f11");
+ }
+}
+
+// Broadcast a single scalar to n floats using fbc.ps (broadcast to all lanes).
+// n must be a multiple of 16 (cacheline-aligned).
+static inline void broadcast_scalar_aligned(float * dst, float val, int32_t n) {
+ __asm__ volatile("fbc.ps f11, %[v]\n" : : [v] "m"(val) : "f11");
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile("fsw.ps f11, %[dst_vec]\n" : [dst_vec] "=m"(*(float (*)[8]) & dst[i])::"f11");
+ }
+}
+
+int entry_point(struct ggml_et_repeat_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne00 = src0->ne[0], ne01 = src0->ne[1], ne02 = src0->ne[2], ne03 = src0->ne[3];
+ const int64_t ne0 = dst->ne[0], ne1 = dst->ne[1], ne2 = dst->ne[2], ne3 = dst->ne[3];
+
+ // src0 strides in bytes
+ const size_t nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ // dst strides in bytes
+ const size_t dnb0 = dst->nb[0], dnb1 = dst->nb[1], dnb2 = dst->nb[2], dnb3 = dst->nb[3];
+
+ // Repeat counts per dimension
+ const int32_t nr0 = (int32_t) (ne0 / ne00);
+ const int32_t nr1 = (int32_t) (ne1 / ne01);
+ const int32_t nr2 = (int32_t) (ne2 / ne02);
+ const int32_t nr3 = (int32_t) (ne3 / ne03);
+
+ // Total output rows across all dimensions (excluding dim 0 tiling)
+ const int64_t total_rows = ne1 * ne2 * ne3;
+
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ // Decompose linear row index into dst (i1, i2, i3)
+ int64_t i1 = row % ne1;
+ int64_t i2 = (row / ne1) % ne2;
+ int64_t i3 = row / (ne1 * ne2);
+
+ // Map dst indices back to src0 indices (modular wrap)
+ int64_t k1 = i1 % ne01;
+ int64_t k2 = i2 % ne02;
+ int64_t k3 = i3 % ne03;
+
+ const float * src_row = (const float *) ((const char *) src0_data + k1 * nb01 + k2 * nb02 + k3 * nb03);
+ float * dst_row = (float *) ((char *) dst_data + i1 * dnb1 + i2 * dnb2 + i3 * dnb3);
+
+ if (ne00 == 1) {
+ // Scalar broadcast: splat single value across entire dst row
+ broadcast_scalar_aligned(dst_row, *src_row, (int32_t) ne0);
+ } else if (nr0 == 1) {
+ // No tiling along dim 0 - single cacheline-aligned row copy
+ copy_row_aligned(dst_row, src_row, (int32_t) ne00);
+ } else {
+ // Tile ne00-sized chunks across dim 0
+ for (int32_t i0 = 0; i0 < nr0; i0++) {
+ copy_row_aligned(dst_row + i0 * ne00, src_row, (int32_t) ne00);
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// RMS Norm F32 Kernel
+// Root Mean Square normalization: y[i] = x[i] / sqrt(mean(x^2) + eps)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <stdint.h>
+#include <string.h>
+
+// RMS norm kernel parameters structure
+struct ggml_et_rms_norm_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon parameter for numerical stability
+};
+
+int entry_point(struct ggml_et_rms_norm_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+ float eps = params->eps;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+
+ if (eps < 0.0f) {
+ return -1; // Invalid epsilon
+ }
+
+ const int64_t ne0 = dst->ne[0]; // Inner dimension (row size)
+ const int64_t ne1 = dst->ne[1]; // Dimension 1
+ const int64_t ne2 = dst->ne[2]; // Dimension 2
+ const int64_t ne3 = dst->ne[3]; // Dimension 3
+
+ // Get dst strides (in bytes)
+ const size_t nb0 = dst->nb[0], nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+
+ // Get src0 strides (in bytes)
+ const size_t nb00 = src0->nb[0], nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+
+ // Verify that src0 and dst have same shape (required for RMS norm)
+ if (src0->ne[0] != ne0 || src0->ne[1] != ne1 || src0->ne[2] != ne2 || src0->ne[3] != ne3) {
+ return -1; // Shape mismatch
+ }
+
+ // RMS norm processes rows independently
+ // Parallelize across rows using simple striding
+ // TODO: ensure lines don't cross cache lines
+ // Precompute reciprocal of row length (constant across all rows)
+ const float inv_ne0 = et_fdiv(1.0f, (float) (int32_t) ne0);
+ const int32_t total_rows = (int32_t) (ne1 * ne2 * ne3);
+
+ // Intra-row cooperation only works within a single shire (barrier + L2SCP
+ // are shire-local). Use per-shire thread count for the threshold.
+ const int shire_threads = SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION; // 64
+
+ if (total_rows >= shire_threads) {
+ // Row-parallel: each thread processes whole rows
+ for (int64_t i3 = 0; i3 < ne3; i3++) {
+ for (int64_t i2 = 0; i2 < ne2; i2++) {
+ for (int64_t i1 = thread_id; i1 < ne1; i1 += num_threads) {
+ const float * src_ptr =
+ (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ // Set mask to enable all 8 vector lanes
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Step 1: Compute sum of squares using 8-wide vectors
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fmadd.ps f10, f11, f11, f10\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ // Horizontal reduce
+ float sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ // Step 2: scale = rsqrt(mean + eps)
+ const float scale = et_powf(sum * inv_ne0 + eps, -0.5f);
+
+ if (!(scale > 0.0f)) {
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return -1;
+ }
+
+ // Step 3: Apply scaling: broadcast scale once, reuse across loop
+ uint32_t scale_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(scale_bits) : "f"(scale));
+ __asm__ volatile("fbcx.ps f13, %[sb]\n" : : [sb] "r"(scale_bits) : "f13");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[x_vec]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f12", "f14");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+ }
+ }
+ } else {
+ // Intra-row: threads within each shire cooperate on rows via L2 SCP.
+ // L2 SCP + barrier are shire-local, so use shire-local thread index.
+ int shire_tid = thread_id % shire_threads; // 0..63 within this shire
+ int threads_per_row = shire_threads / total_rows;
+ int my_row = shire_tid / threads_per_row;
+ int local_tid = shire_tid % threads_per_row;
+ int group_base = my_row * threads_per_row; // shire-local group base
+
+ // Excess threads within this shire, barrier and leave
+ if (my_row >= total_rows) {
+ FENCE;
+ et_barrier(ET_BARRIER_SHIRE);
+ return 0;
+ }
+
+ // Unflatten row index
+ int64_t i1 = my_row % ne1;
+ int64_t i2 = (my_row / ne1) % ne2;
+ int64_t i3 = my_row / (ne1 * ne2);
+
+ const float * src_ptr = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ // Chunk boundaries aligned to 16 floats (64-byte cache line)
+ const int32_t elems_per_cl = 16;
+ int32_t total_cls = ((int32_t) ne0 + elems_per_cl - 1) / elems_per_cl;
+ int32_t cls_per_thread = (total_cls + threads_per_row - 1) / threads_per_row;
+ int32_t my_start = local_tid * cls_per_thread * elems_per_cl;
+ int32_t my_end = my_start + cls_per_thread * elems_per_cl;
+ if (my_end > (int32_t) ne0) {
+ my_end = (int32_t) ne0;
+ }
+ if (my_start >= (int32_t) ne0) {
+ my_start = 0;
+ my_end = 0;
+ }
+
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Phase 1: each thread computes partial sum of squares on its chunk
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fmadd.ps f10, f11, f11, f10\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ // Horizontal reduce to scalar
+ float partial_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(partial_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ // Phase 2: write partial sum to L2 SCP, evict from L1D
+ volatile float * my_slot = (volatile float *) et_shire_l2scp_local((uint64_t) shire_tid * 64);
+ *my_slot = partial_sum;
+ FENCE;
+ evict_to_l2((const void *) my_slot, 1, 64);
+ WAIT_CACHEOPS;
+
+ et_barrier(ET_BARRIER_SHIRE);
+
+ // Phase 3: ALL threads read partial sums, compute scale, apply to own chunk.
+ // Each thread independently reduces to avoid a second barrier.
+ int workers = threads_per_row < total_cls ? threads_per_row : total_cls;
+
+ // Evict stale L1D entries for worker slots
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ evict_to_l2((const void *) slot, 1, 64);
+ }
+ WAIT_CACHEOPS;
+
+ // Every thread reduces the same partial sums -> same scale
+ float total_sum = 0.0f;
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ total_sum += *slot;
+ }
+
+ const float scale = et_powf(total_sum * inv_ne0 + eps, -0.5f);
+ if (!(scale > 0.0f)) {
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return -1;
+ }
+
+ // Each thread applies scale to its own chunk only
+ if (my_start < my_end) {
+ uint32_t scale_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(scale_bits) : "f"(scale));
+ __asm__ volatile("fbcx.ps f13, %[sb]\n" : : [sb] "r"(scale_bits) : "f13");
+
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[x_vec]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f12", "f14");
+ }
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+
+ return 0;
+}
--- /dev/null
+
+// Fused RMS Norm + MUL F32 Kernel
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <stdint.h>
+#include <string.h>
+
+// Fused RMS norm + MUL kernel parameters structure
+struct ggml_et_rms_norm_mul_params {
+ struct ggml_tensor src0; // F32 input tensor (to be normalized)
+ struct ggml_tensor src1; // F32 weights tensor (element-wise multiply)
+ struct ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon for numerical stability
+};
+
+static inline size_t tensor_bytes(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+int entry_point(struct ggml_et_rms_norm_mul_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ float eps = params->eps;
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * src1_data = (float *) src1->data;
+ float * dst_data = (float *) dst->data;
+ // #ifdef ET_UBERKERNEL
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+ // evict_region_past_l2(src1_data, tensor_bytes(src1));
+ // // WAIT_CACHEOPS;
+ // FENCE;
+ // // et_barrier(ET_BARRIER_GLOBAL);
+ // #endif
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ if (eps < 0.0f) {
+ return -1; // Invalid epsilon
+ }
+
+ const int64_t ne0 = dst->ne[0]; // Inner dimension (row size)
+ const int64_t ne1 = dst->ne[1]; // Dimension 1
+ const int64_t ne2 = dst->ne[2]; // Dimension 2
+ const int64_t ne3 = dst->ne[3]; // Dimension 3
+
+ // Get dst strides (in bytes)
+ const size_t nb0 = dst->nb[0], nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+
+ // Get src0 strides (in bytes)
+ const size_t nb00 = src0->nb[0], nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+
+ // Get src1 (weights) strides (in bytes), supports broadcasting in dims 1,2,3
+ const size_t nb10 = src1->nb[0], nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
+
+ // Verify that src0 and dst have same shape (required for RMS norm)
+ if (src0->ne[0] != ne0 || src0->ne[1] != ne1 || src0->ne[2] != ne2 || src0->ne[3] != ne3) {
+ return -1; // Shape mismatch
+ }
+ // et_barrier(ET_BARRIER_GLOBAL);
+
+ const float inv_ne0 = et_fdiv(1.0f, (float) (int32_t) ne0);
+ const int32_t total_rows = (int32_t) (ne1 * ne2 * ne3);
+ const int shire_threads = SOC_MINIONS_PER_SHIRE * NUM_HARTS_PER_MINION;
+
+ if (total_rows >= shire_threads) {
+ // Row-parallel: each thread processes whole rows
+ for (int64_t i3 = 0; i3 < ne3; i3++) {
+ for (int64_t i2 = 0; i2 < ne2; i2++) {
+ for (int64_t i1 = thread_id; i1 < ne1; i1 += num_threads) {
+ const float * src_ptr =
+ (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ const float * wgt_ptr = (const float *) ((const char *) src1_data + (i3 % src1->ne[3]) * nb13 +
+ (i2 % src1->ne[2]) * nb12 + (i1 % src1->ne[1]) * nb11);
+
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Sum of squares
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fmadd.ps f10, f11, f11, f10\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ float sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ const float scale = et_powf(sum * inv_ne0 + eps, -0.5f);
+ if (!(scale > 0.0f)) {
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return -1;
+ }
+
+ uint32_t scale_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(scale_bits) : "f"(scale));
+ __asm__ volatile("fbcx.ps f13, %[sb]\n" : : [sb] "r"(scale_bits) : "f13");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne0; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[x_vec]\n"
+ "flw.ps f15, %[w_vec]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fmul.ps f14, f14, f15\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0]), [w_vec] "m"(*(const float (*)[8]) &
+ wgt_ptr[i0])
+ : "f12", "f14", "f15");
+ }
+ // #ifdef ET_UBERKERNEL
+ // FENCE;
+ // evict_region_past_l2(dst_ptr, (size_t)ne0 * sizeof(float));
+ // WAIT_CACHEOPS;
+ // FENCE;
+ // #endif
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+ }
+ }
+ } else {
+ // Intra-row: threads within each shire cooperate on rows via L2 SCP.
+ // L2 SCP + barrier are shire-local, so use shire-local thread index.
+ int shire_tid = thread_id % shire_threads;
+ int threads_per_row = shire_threads / total_rows;
+ int my_row = shire_tid / threads_per_row;
+ int local_tid = shire_tid % threads_per_row;
+ int group_base = my_row * threads_per_row;
+
+ // Excess threads within this shire
+ if (my_row >= total_rows) {
+ __asm__ __volatile__("fence\n" ::: "memory");
+ et_barrier(ET_BARRIER_SHIRE);
+ return 0;
+ }
+
+ // Unflatten row index
+ int64_t i1 = my_row % ne1;
+ int64_t i2 = (my_row / ne1) % ne2;
+ int64_t i3 = my_row / (ne1 * ne2);
+
+ const float * src_ptr = (const float *) ((const char *) src0_data + i3 * nb03 + i2 * nb02 + i1 * nb01);
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
+
+ const float * wgt_ptr = (const float *) ((const char *) src1_data + (i3 % src1->ne[3]) * nb13 +
+ (i2 % src1->ne[2]) * nb12 + (i1 % src1->ne[1]) * nb11);
+
+ // Chunk boundaries aligned to 16 floats (64-byte cache line)
+ const int32_t elems_per_cl = 16;
+ int32_t total_cls = ((int32_t) ne0 + elems_per_cl - 1) / elems_per_cl;
+ int32_t cls_per_thread = (total_cls + threads_per_row - 1) / threads_per_row;
+ int32_t my_start = local_tid * cls_per_thread * elems_per_cl;
+ int32_t my_end = my_start + cls_per_thread * elems_per_cl;
+ if (my_end > (int32_t) ne0) {
+ my_end = (int32_t) ne0;
+ }
+ if (my_start >= (int32_t) ne0) {
+ my_start = 0;
+ my_end = 0;
+ }
+
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Phase 1: partial sum of squares on own chunk
+ __asm__ volatile("fbci.pi f10, 0" ::: "f10");
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fmadd.ps f10, f11, f11, f10\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ float partial_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(partial_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ // Phase 2: write partial sum to L2 SCP, evict from L1D
+ volatile float * my_slot = (volatile float *) et_shire_l2scp_local((uint64_t) shire_tid * 64);
+ *my_slot = partial_sum;
+ __asm__ __volatile__("fence\n" ::: "memory");
+ evict_to_l2((const void *) my_slot, 1, 64);
+ WAIT_CACHEOPS;
+
+ et_barrier(ET_BARRIER_SHIRE);
+
+ // Phase 3: all threads read partial sums, compute scale, apply to own chunk
+ int workers = threads_per_row < total_cls ? threads_per_row : total_cls;
+
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ evict_to_l2((const void *) slot, 1, 64);
+ }
+ WAIT_CACHEOPS;
+
+ float total_sum = 0.0f;
+ for (int t = 0; t < workers; t++) {
+ volatile float * slot = (volatile float *) et_shire_l2scp_local((uint64_t) (group_base + t) * 64);
+ total_sum += *slot;
+ }
+
+ const float scale = et_powf(total_sum * inv_ne0 + eps, -0.5f);
+ if (!(scale > 0.0f)) {
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ return -1;
+ }
+
+ // Apply scale * weights to own chunk
+ if (my_start < my_end) {
+ uint32_t scale_bits;
+ __asm__ volatile("fmv.x.s %0, %1" : "=r"(scale_bits) : "f"(scale));
+ __asm__ volatile("fbcx.ps f13, %[sb]\n" : : [sb] "r"(scale_bits) : "f13");
+
+ for (int32_t i0 = my_start; i0 < my_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f12, %[x_vec]\n"
+ "flw.ps f15, %[w_vec]\n"
+ "fmul.ps f14, f12, f13\n"
+ "fmul.ps f14, f14, f15\n"
+ "fsw.ps f14, %[result]\n"
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0]), [w_vec] "m"(*(const float (*)[8]) & wgt_ptr[i0])
+ : "f12", "f14", "f15");
+ }
+ // #ifdef ET_UBERKERNEL
+ // FENCE;
+ // evict_region_past_l2(dst_ptr + my_start, (size_t)(my_end - my_start) * sizeof(float));
+ // WAIT_CACHEOPS;
+ // FENCE;
+ // #endif
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// ROPE (Rotary Position Encoding) Kernel
+// Experiment 1:
+// - Keep old scheduling and rotate logic
+// - ONLY SIMD-ize sin/cos approximation inside compute_rope_cache()
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <etsoc/common/utils.h>
+#include <stdint.h>
+
+// ROPE constants (matching GGML definitions)
+#define GGML_ROPE_TYPE_NEOX 2
+#define GGML_ROPE_TYPE_MROPE 8
+#define GGML_ROPE_TYPE_IMROPE 40
+#define MAX_ROPE_HALF_DIMS 256 // supports up to n_dims=512
+
+#define ROPE_VEC_WIDTH 8
+
+#define ROPE_PI 3.14159265358979323846f
+#define ROPE_TWO_PI 6.28318530717958647693f
+#define ROPE_PI_OVER_2 1.57079632679489661923f
+#define ROPE_INV_TWO_PI 0.15915494309189533577f
+
+// ROPE operation parameters structure (matches ggml-et-ops.h)
+typedef struct {
+ int32_t n_past;
+ int32_t n_dims; // Number of dimensions to apply ROPE to (must be even)
+ int32_t mode; // ROPE mode (0=normal, 2=neox)
+ int32_t n_ctx;
+ int32_t n_ctx_orig;
+ float freq_base; // Base frequency (usually 10000.0f)
+ float freq_scale; // Frequency scaling factor
+ float ext_factor; // Extension factor for YaRN
+ float attn_factor; // Attention factor for YaRN
+ float beta_fast; // Fast beta for YaRN
+ float beta_slow; // Slow beta for YaRN
+ int32_t sections[4]; // Sections for multi-modal ROPE
+} rope_params_t;
+
+// ROPE kernel parameters structure (matches ggml_et_rope_params)
+struct ggml_et_rope_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor src1; // I32 position tensor
+ struct ggml_tensor src2; // F32 frequency factors (optional)
+ struct ggml_tensor dst; // F32 output tensor
+ rope_params_t rope_params;
+};
+
+//------------------------------------------------------------------------------
+// Existing scalar helpers
+//------------------------------------------------------------------------------
+
+// floor/ceil with ±inf and NaN passthrough.
+static inline float rope_floorf(float x) {
+ union {
+ float f;
+ uint32_t u;
+ } v = { .f = x };
+
+ const uint32_t expo = (v.u >> 23) & 0xFF;
+ if (expo == 0xFF) {
+ return x; // inf or NaN
+ }
+ if (expo >= 23 + 127) {
+ return x; // already integer-valued
+ }
+ int i = (int) x;
+ return (x < 0.0f && (float) i != x) ? (float) (i - 1) : (float) i;
+}
+
+static inline float rope_ceilf(float x) {
+ union {
+ float f;
+ uint32_t u;
+ } v = { .f = x };
+
+ const uint32_t expo = (v.u >> 23) & 0xFF;
+ if (expo == 0xFF) {
+ return x; // inf or NaN
+ }
+ if (expo >= 23 + 127) {
+ return x; // already integer-valued
+ }
+ int i = (int) x;
+ return (x > 0.0f && (float) i != x) ? (float) (i + 1) : (float) i;
+}
+
+static inline float rope_yarn_ramp(const float low, const float high, const int i0) {
+ float denom = high - low;
+ if (denom < 0.001f) {
+ denom = 0.001f;
+ }
+
+ const float y = et_fdiv((float) (i0 / 2) - low, denom);
+ const float clamped = y < 0.0f ? 0.0f : (y > 1.0f ? 1.0f : y);
+ return 1.0f - clamped;
+}
+
+// Matches CPU reference (ggml_rope_yarn_corr_dim).
+static inline float rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float beta, float freq_base) {
+ return (float) n_dims *
+ et_fdiv(et_logf(et_fdiv((float) n_ctx_orig, beta * ROPE_TWO_PI)), 2.0f * et_logf(freq_base));
+}
+
+static inline void rope_yarn_corr_dims(int n_dims,
+ int n_ctx_orig,
+ float freq_base,
+ float beta_fast,
+ float beta_slow,
+ float dims[2]) {
+ // Match CPU: floor on start, ceil on end, then clamp to [0, n_dims-1].
+ float start = rope_floorf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base));
+ float end = rope_ceilf(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base));
+
+ dims[0] = start > 0.0f ? start : 0.0f;
+ dims[1] = end < (float) (n_dims - 1) ? end : (float) (n_dims - 1);
+}
+
+//------------------------------------------------------------------------------
+// SIMD sin/cos approximation
+//------------------------------------------------------------------------------
+
+static const float rope_ps_one[ROPE_VEC_WIDTH]
+ __attribute__((aligned(32))) = { 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f };
+static const float rope_ps_c3[ROPE_VEC_WIDTH]
+ __attribute__((aligned(32))) = { 1.0f / 6.0f, 1.0f / 6.0f, 1.0f / 6.0f, 1.0f / 6.0f,
+ 1.0f / 6.0f, 1.0f / 6.0f, 1.0f / 6.0f, 1.0f / 6.0f };
+static const float rope_ps_c5[ROPE_VEC_WIDTH]
+ __attribute__((aligned(32))) = { 1.0f / 120.0f, 1.0f / 120.0f, 1.0f / 120.0f, 1.0f / 120.0f,
+ 1.0f / 120.0f, 1.0f / 120.0f, 1.0f / 120.0f, 1.0f / 120.0f };
+static const float rope_ps_c7[ROPE_VEC_WIDTH]
+ __attribute__((aligned(32))) = { 1.0f / 5040.0f, 1.0f / 5040.0f, 1.0f / 5040.0f, 1.0f / 5040.0f,
+ 1.0f / 5040.0f, 1.0f / 5040.0f, 1.0f / 5040.0f, 1.0f / 5040.0f };
+static const float rope_ps_c9[ROPE_VEC_WIDTH]
+ __attribute__((aligned(32))) = { 1.0f / 362880.0f, 1.0f / 362880.0f, 1.0f / 362880.0f, 1.0f / 362880.0f,
+ 1.0f / 362880.0f, 1.0f / 362880.0f, 1.0f / 362880.0f, 1.0f / 362880.0f };
+static const float rope_ps_c11[ROPE_VEC_WIDTH]
+ __attribute__((aligned(32))) = { 1.0f / 39916800.0f, 1.0f / 39916800.0f, 1.0f / 39916800.0f, 1.0f / 39916800.0f,
+ 1.0f / 39916800.0f, 1.0f / 39916800.0f, 1.0f / 39916800.0f, 1.0f / 39916800.0f };
+
+static inline uint64_t rope_ps_enter_fullmask(void) {
+ uint64_t old_mask;
+ __asm__ volatile(
+ "mova.x.m %0 \n\t"
+ "li t0, -1 \n\t"
+ "mova.m.x t0 \n\t"
+ : "=r"(old_mask)
+ :
+ : "t0", "memory");
+ return old_mask;
+}
+
+static inline void rope_ps_leave_fullmask(uint64_t old_mask) {
+ __asm__ volatile("mova.m.x %0 \n\t" : : "r"(old_mask) : "memory");
+}
+
+static inline void rope_poly_sin_block8(float * out, const float * x) {
+ __asm__ volatile(
+ "flw.ps f0, %[x] \n\t"
+ "fmul.ps f1, f0, f0 \n\t"
+
+ "flw.ps f2, %[c11] \n\t"
+ "flw.ps f3, %[c9] \n\t"
+ "fnmsub.ps f2, f1, f2, f3 \n\t"
+
+ "flw.ps f3, %[c7] \n\t"
+ "fnmsub.ps f2, f1, f2, f3 \n\t"
+
+ "flw.ps f3, %[c5] \n\t"
+ "fnmsub.ps f2, f1, f2, f3 \n\t"
+
+ "flw.ps f3, %[c3] \n\t"
+ "fnmsub.ps f2, f1, f2, f3 \n\t"
+
+ "flw.ps f3, %[one] \n\t"
+ "fnmsub.ps f2, f1, f2, f3 \n\t"
+
+ "fmul.ps f4, f0, f2 \n\t"
+ "fsw.ps f4, %[out] \n\t"
+ : [out] "=m"(*(float (*)[ROPE_VEC_WIDTH]) out)
+ : [x] "m"(*(const float (*)[ROPE_VEC_WIDTH]) x), [one] "m"(*(const float (*)[ROPE_VEC_WIDTH]) rope_ps_one),
+ [c3] "m"(*(const float (*)[ROPE_VEC_WIDTH]) rope_ps_c3),
+ [c5] "m"(*(const float (*)[ROPE_VEC_WIDTH]) rope_ps_c5),
+ [c7] "m"(*(const float (*)[ROPE_VEC_WIDTH]) rope_ps_c7),
+ [c9] "m"(*(const float (*)[ROPE_VEC_WIDTH]) rope_ps_c9),
+ [c11] "m"(*(const float (*)[ROPE_VEC_WIDTH]) rope_ps_c11)
+ : "f0", "f1", "f2", "f3", "f4", "memory");
+}
+
+static inline void rope_sincos_block8(float * sin8, float * cos8, const float * theta8) {
+ float sin_fold[ROPE_VEC_WIDTH] __attribute__((aligned(32)));
+ float cos_fold[ROPE_VEC_WIDTH] __attribute__((aligned(32)));
+ float sin_sign[ROPE_VEC_WIDTH] __attribute__((aligned(32)));
+ float cos_sign[ROPE_VEC_WIDTH] __attribute__((aligned(32)));
+
+ for (int i = 0; i < ROPE_VEC_WIDTH; ++i) {
+ float x = theta8[i];
+
+ if (x > ROPE_PI || x < -ROPE_PI) {
+ float cycles = x * ROPE_INV_TWO_PI;
+ int n = (int) cycles;
+ if (x < 0.0f) {
+ n--;
+ }
+ x = x - (float) n * ROPE_TWO_PI;
+ }
+
+ {
+ float y = x;
+ float s = 1.0f;
+ if (y > ROPE_PI_OVER_2) {
+ y = ROPE_PI - y;
+ } else if (y < -ROPE_PI_OVER_2) {
+ y = -ROPE_PI - y;
+ s = -1.0f;
+ }
+ sin_fold[i] = y;
+ sin_sign[i] = s;
+ }
+
+ {
+ float y = x + ROPE_PI_OVER_2;
+ if (y > ROPE_PI || y < -ROPE_PI) {
+ float cycles = y * ROPE_INV_TWO_PI;
+ int n = (int) cycles;
+ if (y < 0.0f) {
+ n--;
+ }
+ y = y - (float) n * ROPE_TWO_PI;
+ }
+
+ float s = 1.0f;
+ if (y > ROPE_PI_OVER_2) {
+ y = ROPE_PI - y;
+ } else if (y < -ROPE_PI_OVER_2) {
+ y = -ROPE_PI - y;
+ s = -1.0f;
+ }
+ cos_fold[i] = y;
+ cos_sign[i] = s;
+ }
+ }
+
+ {
+ const uint64_t saved_mask = rope_ps_enter_fullmask();
+
+ rope_poly_sin_block8(sin8, sin_fold);
+ rope_poly_sin_block8(cos8, cos_fold);
+
+ __asm__ volatile(
+ "flw.ps f0, %[sinv] \n\t"
+ "flw.ps f1, %[sinsgn] \n\t"
+ "fmul.ps f2, f0, f1 \n\t"
+ "fsw.ps f2, %[sout] \n\t"
+
+ "flw.ps f3, %[cosv] \n\t"
+ "flw.ps f4, %[cossgn] \n\t"
+ "fmul.ps f5, f3, f4 \n\t"
+ "fsw.ps f5, %[cout] \n\t"
+ : [sout] "=m"(*(float (*)[ROPE_VEC_WIDTH]) sin8), [cout] "=m"(*(float (*)[ROPE_VEC_WIDTH]) cos8)
+ : [sinv] "m"(*(const float (*)[ROPE_VEC_WIDTH]) sin8),
+ [sinsgn] "m"(*(const float (*)[ROPE_VEC_WIDTH]) sin_sign),
+ [cosv] "m"(*(const float (*)[ROPE_VEC_WIDTH]) cos8),
+ [cossgn] "m"(*(const float (*)[ROPE_VEC_WIDTH]) cos_sign)
+ : "f0", "f1", "f2", "f3", "f4", "f5", "memory");
+
+ rope_ps_leave_fullmask(saved_mask);
+ }
+}
+
+//------------------------------------------------------------------------------
+// Cache build
+//------------------------------------------------------------------------------
+
+// scalar fallback for tail / tiny sizes
+static inline void rope_yarn_scalar(float theta_extrap,
+ float freq_scale,
+ const float corr_dims[2],
+ int64_t i0,
+ float ext_factor,
+ float mscale,
+ float * cos_theta,
+ float * sin_theta) {
+ float theta_interp = freq_scale * theta_extrap;
+ float theta = theta_interp;
+
+ if (ext_factor != 0.0f) {
+ float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], (int) i0) * ext_factor;
+ theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
+ mscale *= 1.0f + 0.1f * et_logf(et_fdiv(1.0f, freq_scale));
+ }
+
+ *cos_theta = et_cosf(theta) * mscale;
+ *sin_theta = et_sinf(theta) * mscale;
+}
+
+// Populate cos/sin cache for a given position using running theta product
+// Experiment 1:
+// - theta construction and YaRN mixing stay scalar
+// - actual sin/cos approximation is done in vec8 blocks
+static inline void compute_rope_cache(float * cos_cache,
+ float * sin_cache,
+ int32_t n_dims,
+ float theta_scale,
+ int32_t pos,
+ const float * freq_factors,
+ float freq_scale,
+ const float corr_dims[2],
+ float ext_factor,
+ float attn_factor) {
+ const int32_t half_dims = n_dims / 2;
+ float theta = 1.0f;
+
+ int32_t dim_idx = 0;
+
+ for (; dim_idx + ROPE_VEC_WIDTH <= half_dims; dim_idx += ROPE_VEC_WIDTH) {
+ float theta_block[ROPE_VEC_WIDTH] __attribute__((aligned(32)));
+ float theta_local = theta;
+ float mscale = attn_factor;
+
+ if (ext_factor != 0.0f) {
+ mscale *= 1.0f + 0.1f * et_logf(et_fdiv(1.0f, freq_scale));
+ }
+
+ for (int i = 0; i < ROPE_VEC_WIDTH; ++i) {
+ const int32_t pair_idx = dim_idx + i;
+ const float ff = freq_factors ? freq_factors[pair_idx] : 1.0f;
+ const float theta_base = (float) pos * theta_local;
+ const float theta_extrap = et_fdiv(theta_base, ff);
+
+ float theta_interp = freq_scale * theta_extrap;
+ float theta_mix = theta_interp;
+
+ if (ext_factor != 0.0f) {
+ float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], pair_idx * 2) * ext_factor;
+ theta_mix = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
+ }
+
+ theta_block[i] = theta_mix;
+ theta_local *= theta_scale;
+ }
+
+ rope_sincos_block8(&sin_cache[dim_idx], &cos_cache[dim_idx], theta_block);
+
+ for (int i = 0; i < ROPE_VEC_WIDTH; ++i) {
+ sin_cache[dim_idx + i] *= mscale;
+ cos_cache[dim_idx + i] *= mscale;
+ }
+
+ theta = theta_local;
+ }
+
+ // tail fallback
+ for (; dim_idx < half_dims; ++dim_idx) {
+ const float ff = freq_factors ? freq_factors[dim_idx] : 1.0f;
+ const float theta_base = (float) pos * theta;
+
+ rope_yarn_scalar(et_fdiv(theta_base, ff), freq_scale, corr_dims, dim_idx * 2, ext_factor, attn_factor,
+ &cos_cache[dim_idx], &sin_cache[dim_idx]);
+
+ theta *= theta_scale;
+ }
+}
+
+//------------------------------------------------------------------------------
+// IMROPE cache build (interleaved multi-modal RoPE for Qwen3VL)
+//------------------------------------------------------------------------------
+
+// Builds cos/sin cache with 4 interleaved position channels.
+// Each dimension pair selects from {theta_t, theta_h, theta_w, theta_e}
+// using a mod-3 sector pattern, matching the CPU reference exactly.
+static inline void compute_imrope_cache(float * cos_cache,
+ float * sin_cache,
+ int32_t n_dims,
+ float theta_scale,
+ int32_t pos_t,
+ int32_t pos_h,
+ int32_t pos_w,
+ int32_t pos_e,
+ const int32_t sections[4],
+ const float * freq_factors,
+ float freq_scale,
+ const float corr_dims[2],
+ float ext_factor,
+ float attn_factor) {
+ const int32_t half_dims = n_dims / 2;
+ const int32_t sect_dims = sections[0] + sections[1] + sections[2] + sections[3];
+
+ float theta_t = (float) pos_t;
+ float theta_h = (float) pos_h;
+ float theta_w = (float) pos_w;
+ float theta_e = (float) pos_e;
+
+ int32_t dim_idx = 0;
+
+ for (; dim_idx + ROPE_VEC_WIDTH <= half_dims; dim_idx += ROPE_VEC_WIDTH) {
+ float theta_block[ROPE_VEC_WIDTH] __attribute__((aligned(32)));
+ float mscale = attn_factor;
+
+ if (ext_factor != 0.0f) {
+ mscale *= 1.0f + 0.1f * et_logf(et_fdiv(1.0f, freq_scale));
+ }
+
+ for (int i = 0; i < ROPE_VEC_WIDTH; ++i) {
+ const int32_t pair_idx = dim_idx + i;
+ const int32_t sector = pair_idx % sect_dims;
+ const float ff = freq_factors ? freq_factors[pair_idx] : 1.0f;
+
+ // Interleaved sector assignment (mod-3 pattern)
+ float theta;
+ if (sector % 3 == 1 && sector < 3 * sections[1]) {
+ theta = theta_h;
+ } else if (sector % 3 == 2 && sector < 3 * sections[2]) {
+ theta = theta_w;
+ } else if (sector % 3 == 0 && sector < 3 * sections[0]) {
+ theta = theta_t;
+ } else {
+ theta = theta_e;
+ }
+
+ const float theta_extrap = et_fdiv(theta, ff);
+ float theta_interp = freq_scale * theta_extrap;
+ float theta_mix = theta_interp;
+
+ if (ext_factor != 0.0f) {
+ float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], pair_idx * 2) * ext_factor;
+ theta_mix = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
+ }
+
+ theta_block[i] = theta_mix;
+
+ // All 4 thetas advance every iteration
+ theta_t *= theta_scale;
+ theta_h *= theta_scale;
+ theta_w *= theta_scale;
+ theta_e *= theta_scale;
+ }
+
+ rope_sincos_block8(&sin_cache[dim_idx], &cos_cache[dim_idx], theta_block);
+
+ for (int i = 0; i < ROPE_VEC_WIDTH; ++i) {
+ sin_cache[dim_idx + i] *= mscale;
+ cos_cache[dim_idx + i] *= mscale;
+ }
+ }
+
+ // Scalar tail
+ for (; dim_idx < half_dims; ++dim_idx) {
+ const int32_t sector = dim_idx % sect_dims;
+ const float ff = freq_factors ? freq_factors[dim_idx] : 1.0f;
+
+ float theta;
+ if (sector % 3 == 1 && sector < 3 * sections[1]) {
+ theta = theta_h;
+ } else if (sector % 3 == 2 && sector < 3 * sections[2]) {
+ theta = theta_w;
+ } else if (sector % 3 == 0 && sector < 3 * sections[0]) {
+ theta = theta_t;
+ } else {
+ theta = theta_e;
+ }
+
+ rope_yarn_scalar(et_fdiv(theta, ff), freq_scale, corr_dims, dim_idx * 2, ext_factor, attn_factor,
+ &cos_cache[dim_idx], &sin_cache[dim_idx]);
+
+ theta_t *= theta_scale;
+ theta_h *= theta_scale;
+ theta_w *= theta_scale;
+ theta_e *= theta_scale;
+ }
+}
+
+//------------------------------------------------------------------------------
+// Entry point
+//------------------------------------------------------------------------------
+
+int entry_point(struct ggml_et_rope_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return -1;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * src2 = ¶ms->src2;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_I32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * src0_data = (const float *) src0->data;
+ const int32_t * src1_data = (const int32_t *) src1->data;
+ const float * freq_factors = (src2 && src2->data) ? (const float *) src2->data : NULL;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1;
+ }
+#ifdef ET_UBERKERNEL
+ const size_t src0_bytes = (size_t) src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3] * src0->nb[0];
+ const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * src1->nb[0];
+ evict_region_past_l2(src0_data, src0_bytes);
+ evict_region_past_l2(src1_data, src1_bytes);
+ WAIT_CACHEOPS;
+ FENCE;
+ et_barrier(ET_BARRIER_GLOBAL);
+#endif
+ const int64_t head_dim = src0->ne[0];
+ const int64_t heads = src0->ne[1];
+ const int64_t seq_len = src0->ne[2];
+ const int64_t batch = src0->ne[3];
+
+ const rope_params_t * rope_params = ¶ms->rope_params;
+ const int32_t n_dims = rope_params->n_dims;
+ const float freq_base = rope_params->freq_base;
+ const float freq_scale = rope_params->freq_scale;
+ const int32_t mode = rope_params->mode;
+
+ if (n_dims <= 0 || n_dims > head_dim || (n_dims & 1) != 0) {
+ return -1;
+ }
+
+ if (n_dims / 2 > MAX_ROPE_HALF_DIMS) {
+ return -1;
+ }
+
+ float cos_cache[MAX_ROPE_HALF_DIMS];
+ float sin_cache[MAX_ROPE_HALF_DIMS];
+
+ float corr_dims[2];
+ rope_yarn_corr_dims(n_dims, rope_params->n_ctx_orig, freq_base, rope_params->beta_fast, rope_params->beta_slow,
+ corr_dims);
+ et_barrier(ET_BARRIER_GLOBAL);
+
+ // Distribute by individual heads: total = batch * seq_len * heads.
+ const int64_t total_heads = batch * seq_len * heads;
+ const int64_t start_wu = (total_heads * thread_id) / num_threads;
+ const int64_t end_wu = (total_heads * (thread_id + 1)) / num_threads;
+
+ if (start_wu >= end_wu) {
+ return 0;
+ }
+
+ const float theta_scale = et_powf(freq_base, et_fdiv(-2.0f, (float) n_dims));
+ const int32_t half_dims = n_dims / 2;
+ const int is_neox = (mode & GGML_ROPE_TYPE_NEOX) != 0;
+ const int is_imrope = (mode == GGML_ROPE_TYPE_IMROPE);
+ const int use_neox_rotation = is_neox || is_imrope;
+
+ // For IMROPE position cache invalidation: track all 4 channels
+ int32_t last_pos = -1;
+ int32_t last_pos_h = -1;
+ int32_t last_pos_w = -1;
+ int32_t last_pos_e = -1;
+
+ for (int64_t wu = start_wu; wu < end_wu; ++wu) {
+ const int64_t h = wu % heads;
+ const int64_t s = (wu / heads) % seq_len;
+ const int64_t b = wu / (heads * seq_len);
+
+ if (is_imrope) {
+ // IMROPE: src1 layout is [p_t(0..S-1), p_h(0..S-1), p_w(0..S-1), p_e(0..S-1)]
+ const int32_t pt = src1_data[s] + rope_params->n_past;
+ const int32_t ph = src1_data[s + seq_len] + rope_params->n_past;
+ const int32_t pw = src1_data[s + seq_len * 2] + rope_params->n_past;
+ const int32_t pe = src1_data[s + seq_len * 3] + rope_params->n_past;
+
+ if (pt != last_pos || ph != last_pos_h || pw != last_pos_w || pe != last_pos_e) {
+ compute_imrope_cache(cos_cache, sin_cache, n_dims, theta_scale, pt, ph, pw, pe, rope_params->sections,
+ freq_factors, freq_scale, corr_dims, rope_params->ext_factor,
+ rope_params->attn_factor);
+ last_pos = pt;
+ last_pos_h = ph;
+ last_pos_w = pw;
+ last_pos_e = pe;
+ }
+ } else {
+ const int32_t pos = src1_data[s] + rope_params->n_past;
+
+ if (pos != last_pos) {
+ compute_rope_cache(cos_cache, sin_cache, n_dims, theta_scale, pos, freq_factors, freq_scale, corr_dims,
+ rope_params->ext_factor, rope_params->attn_factor);
+ last_pos = pos;
+ }
+ }
+
+ const float * head_src =
+ (const float *) ((const char *) src0_data + b * src0->nb[3] + s * src0->nb[2] + h * src0->nb[1]);
+
+ float * head_dst = (float *) ((char *) dst_data + b * dst->nb[3] + s * dst->nb[2] + h * dst->nb[1]);
+
+ // Copy dimensions beyond n_dims unchanged
+ for (int64_t d = n_dims; d < head_dim; ++d) {
+ head_dst[d] = head_src[d];
+ }
+
+ if (use_neox_rotation) {
+ // NEOX/IMROPE: pairs at (i, i+half_dims)
+ uint64_t temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ for (int32_t dim_idx = 0; dim_idx < half_dims; dim_idx += 8) {
+ __asm__ volatile(
+ "flw.ps f0, %[x0_src] \n\t"
+ "flw.ps f1, %[x1_src] \n\t"
+ "flw.ps f2, %[sin_cache] \n\t"
+ "flw.ps f3, %[cos_cache] \n\t"
+ "fmul.ps f4, f0, f3 \n\t"
+ "fmul.ps f5, f0, f2 \n\t"
+ "fnmsub.ps f4, f1, f2, f4 \n\t"
+ "fmadd.ps f5, f1, f3, f5 \n\t"
+ "fsw.ps f4, %[x0_dst] \n\t"
+ "fsw.ps f5, %[x1_dst] \n\t"
+ : [x0_dst] "=m"(*(float (*)[8]) & head_dst[dim_idx]), [x1_dst] "=m"(*(float (*)[8]) &
+ head_dst[dim_idx + half_dims])
+ : [x0_src] "m"(*(const float (*)[8]) & head_src[dim_idx]),
+ [x1_src] "m"(*(const float (*)[8]) & head_src[dim_idx + half_dims]),
+ [sin_cache] "m"(*(const float (*)[8]) & sin_cache[dim_idx]),
+ [cos_cache] "m"(*(const float (*)[8]) & cos_cache[dim_idx])
+ : "f0", "f1", "f2", "f3", "f4", "f5", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+ } else {
+ // Standard: adjacent pairs (2i, 2i+1)
+ for (int32_t pair_idx = 0; pair_idx < half_dims; ++pair_idx) {
+ const int32_t dim_in_head = pair_idx * 2;
+ const float x0 = head_src[dim_in_head];
+ const float x1 = head_src[dim_in_head + 1];
+
+ head_dst[dim_in_head] = x0 * cos_cache[pair_idx] - x1 * sin_cache[pair_idx];
+ head_dst[dim_in_head + 1] = x0 * sin_cache[pair_idx] + x1 * cos_cache[pair_idx];
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// RWKV WKV6 F32 Kernel
+//
+// Implements the RWKV-6 linear attention recurrence:
+// dst = r @ (time_faaaa * (k @ v) + state)
+// state = time_decay * state + (k @ v)
+//
+// For each head h, timestep t, row i:
+// kv[j] = v[j] * k[i]
+// temp[j] = kv[j] * tf[i] + state[i][j]
+// dst[j] += temp[j] * r[i] (accumulated across all i)
+// state[i][j] = state[i][j] * td[i] + kv[j]
+//
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_rwkv_wkv6_params {
+ float * k; // src[0]: [S, H, T] key
+ float * v; // src[1]: [S, H, T] value
+ float * r; // src[2]: [S, H, T] receptance
+ float * tf; // src[3]: [S, H] time_faaaa (per-head, not per-token)
+ float * td; // src[4]: [S, H, T] time_decay
+ float * state_in; // src[5]: [S*S*H, n_seqs] initial state
+ float * dst; // [C, T + S*n_seqs] output + state_out
+ int32_t C; // total channels (S * H)
+ int32_t H; // number of heads
+ int32_t S; // head size
+ int32_t T; // number of tokens
+ int32_t n_seqs; // number of sequences
+};
+
+int entry_point(struct ggml_et_rwkv_wkv6_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ const float * k = params->k;
+ const float * v = params->v;
+ const float * r = params->r;
+ const float * tf = params->tf;
+ const float * td = params->td;
+ const float * state_in = params->state_in;
+ float * dst_data = params->dst;
+
+ const int32_t C = params->C;
+ const int32_t H = params->H;
+ const int32_t S = params->S;
+ const int32_t T = params->T;
+ const int32_t n_seqs = params->n_seqs;
+
+ if (!k || !v || !r || !tf || !td || !state_in || !dst_data) {
+ return -1;
+ }
+
+ const int32_t tps = T / n_seqs; // tokens per sequence
+ float * state_out = dst_data + C * T;
+ float zero = 0.0f;
+
+ // Tile j by one cache line so each hart's dst/state writes never share
+ // a 64-B line with another hart's writes (the chip is non-coherent).
+ // Tiling on j (not i) is required for WKV6 because dst[j] is accumulated
+ // across i — splitting i across harts would race on dst writes.
+ // For S=64 this gives 4 tiles per head; for S<16 or odd S we fall back
+ // to one-hart-per-head (= the original parallelism).
+ const int32_t j_tile = (S % 16 == 0) ? 16 : S;
+ const int32_t tiles_per_head = S / j_tile;
+ const int32_t total_units = H * tiles_per_head;
+
+ // Parallelize across (head, j-tile) pairs. The t loop stays inside this
+ // unit loop so the same hart owns the same column slice of state across
+ // all timesteps — required for the recurrence to read back its own
+ // writes without going through L2.
+ for (int32_t u = thread_id; u < total_units; u += num_threads) {
+ const int32_t h = u / tiles_per_head;
+ const int32_t tile = u % tiles_per_head;
+ const int32_t j_start = tile * j_tile;
+ const int32_t j_end = j_start + j_tile;
+
+ const int32_t h_off = h * S; // offset within C for this head
+ const int32_t s2d = h * S * S; // offset within state for this head
+
+ for (int32_t t = 0; t < T; t++) {
+ const int32_t seq = t / tps;
+ const int32_t t_in_seq = t % tps;
+ const int32_t seq_state = seq * S * C;
+
+ const float * s_prev;
+ float * s_cur = state_out + seq_state + s2d;
+
+ if (t_in_seq == 0) {
+ s_prev = state_in + seq_state + s2d;
+ } else {
+ s_prev = s_cur;
+ }
+
+ const int32_t th = t * C + h_off;
+
+ // Pointers for this timestep/head
+ const float * k_ptr = k + th;
+ const float * v_ptr = v + th;
+ const float * r_ptr = r + th;
+ const float * tf_ptr = tf + h_off; // tf is per-head, no t offset
+ const float * td_ptr = td + th;
+
+ // Zero this hart's slice of dst: dst[th + j_start..th + j_end-1]
+ // WKV6 accumulates dst[j] across all i, so must start from zero
+ float * dst_row = dst_data + th;
+ for (int32_t j = j_start; j < j_end; j += 8) {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "fsw.ps f10, %[dst_vec]\n"
+ : [dst_vec] "=m"(*(float (*)[8]) & dst_row[j])
+ : [z] "m"(zero)
+ : "f10");
+ }
+
+ for (int32_t i = 0; i < S; i++) {
+ const float * sp_row = s_prev + i * S; // state_prev row i
+ float * sc_row = s_cur + i * S; // state_cur row i
+
+ float k_val = k_ptr[i];
+ float r_val = r_ptr[i];
+ float tf_val = tf_ptr[i];
+ float td_val = td_ptr[i];
+
+ // Broadcast k[i], r[i], tf[i], td[i] to vector registers
+ __asm__ volatile(
+ "fbc.ps f20, %[kv]\n" // f20 = k[i] broadcast
+ "fbc.ps f21, %[rv]\n" // f21 = r[i] broadcast
+ "fbc.ps f22, %[tfv]\n" // f22 = tf[i] broadcast
+ "fbc.ps f23, %[tdv]\n" // f23 = td[i] broadcast
+ :
+ : [kv] "m"(k_val), [rv] "m"(r_val), [tfv] "m"(tf_val), [tdv] "m"(td_val)
+ : "f20", "f21", "f22", "f23");
+
+ for (int32_t j = j_start; j < j_end; j += 8) {
+ __asm__ volatile(
+ // Load v[j], state_prev[i][j], dst[j]
+ "flw.ps f10, %[v_vec]\n" // v[j..j+7]
+ "flw.ps f11, %[s_vec]\n" // state_prev[i][j..j+7]
+ "flw.ps f12, %[d_vec]\n" // dst[j..j+7] (accumulated)
+
+ // kv = v * k_broadcast
+ "fmul.ps f13, f10, f20\n" // kv = v * k
+
+ // temp = kv * tf_broadcast + state_prev
+ "fmadd.ps f14, f13, f22, f11\n" // temp = kv * tf + state
+
+ // dst[j] += temp * r_broadcast
+ "fmadd.ps f12, f14, f21, f12\n" // dst += temp * r
+ "fsw.ps f12, %[d_out]\n" // store updated dst
+
+ // state_cur[i][j] = state_prev * td_broadcast + kv
+ "fmadd.ps f11, f11, f23, f13\n" // state = state * td + kv
+ "fsw.ps f11, %[s_out]\n" // store new state
+
+ : [d_out] "=m"(*(float (*)[8]) & dst_row[j]), [s_out] "=m"(*(float (*)[8]) & sc_row[j])
+ : [v_vec] "m"(*(const float (*)[8]) & v_ptr[j]), [s_vec] "m"(*(const float (*)[8]) & sp_row[j]),
+ [d_vec] "m"(*(const float (*)[8]) & dst_row[j])
+ : "f10", "f11", "f12", "f13", "f14");
+ }
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// RWKV WKV7 F32 Kernel
+//
+// Implements the RWKV-7 linear attention recurrence:
+// For each head h, timestep t, row i:
+// sa = dot(a, state[i])
+// state[i] = state[i] * w + v[i]*k + sa * b
+// output[i]= dot(state[i], r)
+//
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_rwkv_wkv7_params {
+ float * r; // [S, H, T] receptance
+ float * w; // [S, H, T] decay
+ float * k; // [S, H, T] key
+ float * v; // [S, H, T] value
+ float * a; // [S, H, T] bonus gate
+ float * b; // [S, H, T] bonus key
+ float * state_in; // [S*S*H, n_seqs] initial state
+ float * dst; // [C, T + S*n_seqs] output + state_out
+ int32_t C; // total channels (S * H)
+ int32_t H; // number of heads
+ int32_t S; // head size
+ int32_t T; // number of tokens
+ int32_t n_seqs; // number of sequences
+};
+
+// Horizontal sum of 8-wide vector register f10 -> scalar float
+static inline float hsum_f10(void) {
+ float result;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(result)::"t0", "f1", "f2", "f3", "f4", "f5");
+ return result;
+}
+
+int entry_point(struct ggml_et_rwkv_wkv7_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ const float * r = params->r;
+ const float * w = params->w;
+ const float * k = params->k;
+ const float * v = params->v;
+ const float * a = params->a;
+ const float * b = params->b;
+ const float * state_in = params->state_in;
+ float * dst_data = params->dst;
+
+ const int32_t C = params->C;
+ const int32_t H = params->H;
+ const int32_t S = params->S;
+ const int32_t T = params->T;
+ const int32_t n_seqs = params->n_seqs;
+
+ if (!r || !w || !k || !v || !a || !b || !state_in || !dst_data) {
+ return -1;
+ }
+
+ const int32_t tps = T / n_seqs; // tokens per sequence
+ float * state_out = dst_data + C * T;
+
+ // Fix #2: hoist w[0..S-1] across the i loop. In the inner j-loop of pass
+ // 2, w/k/b/r are loop-invariant w.r.t. i but were being reloaded for every
+ // i value (16 times redundantly after Fix #1). Pinning all four arrays
+ // would need 32 vector regs (won't fit), so we hoist just w — it's used
+ // in the critical fmadd chain and lives cleanly in f24-f31, which the
+ // existing kernel never touches. Saves ~20% of pass-2 load issues.
+ //
+ // GCC local register variables: declared as `float` but the underlying
+ // f-reg holds the wide vector loaded by flw.ps. GCC reserves f24-f31 for
+ // these variables for the whole function and never generates code that
+ // touches them on its own, so the upper 7 lanes survive between asm
+ // blocks. Only used when S == 64 (the RWKV-7 case); other head sizes
+ // fall through to the original unhoisted path.
+ register float w_h0 __asm__("f24");
+ register float w_h1 __asm__("f25");
+ register float w_h2 __asm__("f26");
+ register float w_h3 __asm__("f27");
+ register float w_h4 __asm__("f28");
+ register float w_h5 __asm__("f29");
+ register float w_h6 __asm__("f30");
+ register float w_h7 __asm__("f31");
+ const int wkv7_fast = (S == 64);
+
+ // Tile i by one cache line so each hart's output writes never share a
+ // 64-B line with another hart's writes (the chip is non-coherent).
+ // For S=64 this gives 4 tiles per head; for S<16 or odd S we fall back
+ // to one-hart-per-head (= the original parallelism).
+ const int32_t i_tile = (S % 16 == 0) ? 16 : S;
+ const int32_t tiles_per_head = S / i_tile;
+ const int32_t total_units = H * tiles_per_head;
+
+ // Parallelize across (head, i-tile) pairs. The t loop stays inside this
+ // unit loop so the same hart owns the same state rows across all
+ // timesteps — required for the recurrence to read back its own writes
+ // without going through L2.
+ for (int32_t u = thread_id; u < total_units; u += num_threads) {
+ const int32_t h = u / tiles_per_head;
+ const int32_t tile = u % tiles_per_head;
+ const int32_t i_start = tile * i_tile;
+ const int32_t i_end = i_start + i_tile;
+
+ const int32_t h_off = h * S; // offset within C for this head
+ const int32_t s2d = h * S * S; // offset within state for this head
+
+ for (int32_t t = 0; t < T; t++) {
+ const int32_t seq = t / tps;
+ const int32_t t_in_seq = t % tps;
+ const int32_t seq_state = seq * S * C; // state offset for this sequence
+
+ const float * s_prev;
+ float * s_cur = state_out + seq_state + s2d;
+
+ if (t_in_seq == 0) {
+ s_prev = state_in + seq_state + s2d;
+ } else {
+ s_prev = s_cur;
+ }
+
+ // Pointers for this timestep/head
+ const int32_t th = t * C + h_off;
+ const float * r_ptr = r + th;
+ const float * w_ptr = w + th;
+ const float * k_ptr = k + th;
+ const float * v_ptr = v + th;
+ const float * a_ptr = a + th;
+ const float * b_ptr = b + th;
+
+ // Hoist w[0..63] into f24-f31 once per (h, t). These values are
+ // invariant across the i loop below, so the inner j-unroll can
+ // reference them by register name and skip the per-i reload.
+ if (wkv7_fast) {
+ __asm__ volatile(
+ "flw.ps f24, 0(%[wp])\n"
+ "flw.ps f25, 32(%[wp])\n"
+ "flw.ps f26, 64(%[wp])\n"
+ "flw.ps f27, 96(%[wp])\n"
+ "flw.ps f28, 128(%[wp])\n"
+ "flw.ps f29, 160(%[wp])\n"
+ "flw.ps f30, 192(%[wp])\n"
+ "flw.ps f31, 224(%[wp])\n"
+ : "=f"(w_h0), "=f"(w_h1), "=f"(w_h2), "=f"(w_h3), "=f"(w_h4), "=f"(w_h5), "=f"(w_h6), "=f"(w_h7)
+ : [wp] "r"(w_ptr));
+ }
+
+ for (int32_t i = i_start; i < i_end; i++) {
+ const float * sp_row = s_prev + i * S; // state_prev row i
+ float * sc_row = s_cur + i * S; // state_cur row i
+
+ // ----------------------------------------------------------
+ // Step 1: sa = dot(a, state_prev[i])
+ // Accumulate in f10
+ // ----------------------------------------------------------
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int32_t j = 0; j < S; j += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[a_vec]\n"
+ "flw.ps f12, %[s_vec]\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ :
+ : [a_vec] "m"(*(const float (*)[8]) & a_ptr[j]), [s_vec] "m"(*(const float (*)[8]) & sp_row[j])
+ : "f10", "f11", "f12");
+ }
+
+ float sa = hsum_f10();
+
+ // ----------------------------------------------------------
+ // Step 2: state update + result accumulation
+ // kv = v[i] * k[j]
+ // state[j] = state[j] * w[j] + kv + sa * b[j]
+ // result += state[j] * r[j]
+ // ----------------------------------------------------------
+ float v_val = v_ptr[i];
+
+ // Broadcast v_val and sa, zero result accumulator (f10)
+ __asm__ volatile(
+ "fbc.ps f20, %[vv]\n"
+ "fbc.ps f21, %[sv]\n"
+ "fbc.ps f10, %[z]\n"
+ :
+ : [vv] "m"(v_val), [sv] "m"(sa), [z] "m"(zero)
+ : "f10", "f20", "f21");
+
+ if (wkv7_fast) {
+// Fast path: 8 chunks unrolled, w hoisted to f24-f31.
+// Saves one flw per chunk vs the original loop.
+#define WKV7_PASS2_CHUNK(j_off, w_var) \
+ __asm__ volatile( \
+ "flw.ps f11, %[s_vec]\n" \
+ "flw.ps f13, %[k_vec]\n" \
+ "flw.ps f14, %[b_vec]\n" \
+ "flw.ps f15, %[r_vec]\n" \
+ "fmul.ps f16, f20, f13\n" \
+ "fmadd.ps f11, f11, %[w_h], f16\n" \
+ "fmadd.ps f11, f21, f14, f11\n" \
+ "fsw.ps f11, %[sc_vec]\n" \
+ "fmadd.ps f10, f11, f15, f10\n" \
+ : [sc_vec] "=m"(*(float (*)[8]) & sc_row[j_off]) \
+ : [s_vec] "m"(*(const float (*)[8]) & sp_row[j_off]), [k_vec] "m"(*(const float (*)[8]) & k_ptr[j_off]), \
+ [b_vec] "m"(*(const float (*)[8]) & b_ptr[j_off]), [r_vec] "m"(*(const float (*)[8]) & r_ptr[j_off]), \
+ [w_h] "f"(w_var) \
+ : "f10", "f11", "f13", "f14", "f15", "f16")
+
+ WKV7_PASS2_CHUNK(0, w_h0);
+ WKV7_PASS2_CHUNK(8, w_h1);
+ WKV7_PASS2_CHUNK(16, w_h2);
+ WKV7_PASS2_CHUNK(24, w_h3);
+ WKV7_PASS2_CHUNK(32, w_h4);
+ WKV7_PASS2_CHUNK(40, w_h5);
+ WKV7_PASS2_CHUNK(48, w_h6);
+ WKV7_PASS2_CHUNK(56, w_h7);
+
+#undef WKV7_PASS2_CHUNK
+ } else {
+ for (int32_t j = 0; j < S; j += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[s_vec]\n" // state_prev[j..j+7]
+ "flw.ps f12, %[w_vec]\n" // w[j..j+7]
+ "flw.ps f13, %[k_vec]\n" // k[j..j+7]
+ "flw.ps f14, %[b_vec]\n" // b[j..j+7]
+ "flw.ps f15, %[r_vec]\n" // r[j..j+7]
+ "fmul.ps f16, f20, f13\n" // kv = v_broadcast * k
+ "fmadd.ps f11, f11, f12, f16\n" // state*w + kv
+ "fmadd.ps f11, f21, f14, f11\n" // + sa*b
+ "fsw.ps f11, %[sc_vec]\n" // store new state
+ "fmadd.ps f10, f11, f15, f10\n" // result += new_state * r
+
+ : [sc_vec] "=m"(*(float (*)[8]) & sc_row[j])
+ : [s_vec] "m"(*(const float (*)[8]) & sp_row[j]),
+ [w_vec] "m"(*(const float (*)[8]) & w_ptr[j]),
+ [k_vec] "m"(*(const float (*)[8]) & k_ptr[j]),
+ [b_vec] "m"(*(const float (*)[8]) & b_ptr[j]),
+ [r_vec] "m"(*(const float (*)[8]) & r_ptr[j])
+ : "f10", "f11", "f12", "f13", "f14", "f15", "f16");
+ }
+ }
+
+ dst_data[th + i] = hsum_f10();
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Scale F32 Kernel
+// dst[i] = src0[i] * scale + bias
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_scale_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor dst; // F32 output tensor
+ float scale; // Scale factor
+ float bias; // Bias (additive offset)
+};
+
+int entry_point(struct ggml_et_scale_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+
+ float scale = params->scale;
+ float bias = params->bias;
+
+ // Total elements across all dimensions
+ const int64_t total_elements = src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3];
+
+ // Cache line = 64 bytes = 16 floats, but vector width = 8 floats
+ // Parallelize at cache line granularity (16 floats)
+ const int64_t elements_per_cacheline = 16;
+ const int64_t total_cachelines = (total_elements + elements_per_cacheline - 1) / elements_per_cacheline;
+
+ int64_t cachelines_per_thread = (total_cachelines + num_threads - 1) / num_threads;
+ int64_t start_cacheline = thread_id * cachelines_per_thread;
+ int64_t end_cacheline = start_cacheline + cachelines_per_thread;
+
+ if (end_cacheline > total_cachelines) {
+ end_cacheline = total_cachelines;
+ }
+
+ if (start_cacheline >= total_cachelines) {
+ return 0;
+ }
+
+ int64_t start_elem = start_cacheline * elements_per_cacheline;
+ int64_t end_elem = end_cacheline * elements_per_cacheline;
+ if (end_elem > total_elements) {
+ end_elem = total_elements;
+ }
+
+ unsigned long temp_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(temp_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+ __asm__ volatile("fbc.ps f20, %[scale_ptr]\n" : : [scale_ptr] "m"(scale) : "f20");
+ __asm__ volatile("fbc.ps f21, %[bias_ptr]\n" : : [bias_ptr] "m"(bias) : "f21");
+
+ for (int64_t i = start_elem; i < end_elem; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[src]\n"
+ "fmadd.ps f10, f10, f20, f21\n" // dst = src*scale + bias
+ "fsw.ps f10, %[dst_out]\n"
+ : [dst_out] "=m"(*(float (*)[8]) & dst_data[i])
+ : [src] "m"(*(const float (*)[8]) & src0_data[i])
+ : "f10", "f20", "f21");
+ }
+ __asm__ volatile("mova.m.x %0" ::"r"(temp_mask));
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// SET F32 Kernel
+// Minimal ET implementation for inplace F32 SET into a contiguous destination
+// using a contiguous F32 source view and explicit destination view strides.
+//
+// Supported shape family:
+// - dst/base is contiguous F32
+// - src1 is contiguous F32
+// - src1.ne[0] is cacheline-aligned (multiple of 16 floats)
+// - destination view strides/offset are cacheline-aligned
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_set_params {
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+ int32_t nb1;
+ int32_t nb2;
+ int32_t nb3;
+ int32_t offset;
+};
+
+static inline void copy_row_aligned(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[src_vec]\n"
+ "fsw.ps f11, %[dst_vec]\n"
+ : [dst_vec] "=m"(*(float (*)[8]) & dst[i])
+ : [src_vec] "m"(*(const float (*)[8]) & src[i])
+ : "f11");
+ }
+}
+
+int entry_point(struct ggml_et_set_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * src1_data = (const float *) src1->data;
+ float * dst_data = (float *) dst->data;
+ if (!src1_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne10 = src1->ne[0];
+ const int64_t ne11 = src1->ne[1];
+ const int64_t ne12 = src1->ne[2];
+ const int64_t ne13 = src1->ne[3];
+
+ if (src1->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float) || ne10 % 16 != 0) {
+ return -1;
+ }
+
+ const int64_t nb11 = src1->nb[1];
+ const int64_t nb12 = src1->nb[2];
+ const int64_t nb13 = src1->nb[3];
+
+ const int64_t dnb1 = params->nb1;
+ const int64_t dnb2 = params->nb2;
+ const int64_t dnb3 = params->nb3;
+ const int64_t offset = params->offset;
+
+ const int64_t total_rows = ne11 * ne12 * ne13;
+
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ const int64_t i1 = row % ne11;
+ const int64_t i2 = (row / ne11) % ne12;
+ const int64_t i3 = row / (ne11 * ne12);
+
+ const float * src_row = (const float *) ((const char *) src1_data + i1 * nb11 + i2 * nb12 + i3 * nb13);
+ float * dst_row = (float *) ((char *) dst_data + offset + i1 * dnb1 + i2 * dnb2 + i3 * dnb3);
+
+ copy_row_aligned(dst_row, src_row, (int32_t) ne10);
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal SET_ROWS F32 Kernel
+// Writes source data rows to specific indices in destination tensor
+//
+// Algorithm:
+// 1. Read row indices from src1 (int64 tensor)
+// 2. For each source row, write it to destination at the specified index
+// 3. Handle type conversion: F32 source -> F32/F16 destination
+// 4. Support multi-dimensional tensor operations
+//
+// Operation: dst[indices[i]] = src[i] for i = 0..num_source_rows
+// This is the inverse of GET_ROWS operation
+//
+// As ET is not a cache coherent processor yet SET_ROWS often are setting
+// small mount of large rows (KV cache). There's several strategies to
+// optimize this operation, including cacheline-based parallelization.
+//
+// - distribute work at cacheline granularity
+// - if previous does not work, find the LCM of cacheline size
+//
+// Features supported:
+// - F32 source data (always F32 input)
+// - F32 and F16 destination data (with transcoding)
+// - Int64 row indices (vs Int32 in GET_ROWS)
+// - Multi-dimensional tensor support
+// - Sequential source reads, scattered destination writes
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+#define CACHE_LINE_SIZE_BYTES 64
+#define CACHE_LINE_F32_ELEMS 16 // 64 / 4
+#define CACHE_LINE_F16_ELEMS 32 // 64 / 2
+
+static int64_t gcd64(int64_t a, int64_t b) {
+ while (b) {
+ int64_t t = b;
+ b = a % b;
+ a = t;
+ }
+ return a;
+}
+
+struct ggml_et_set_rows_params {
+ struct ggml_tensor src0; // F32 source data tensor
+ struct ggml_tensor src1; // I64 row indices tensor
+ struct ggml_tensor dst; // F32/F16 destination tensor
+};
+
+// Copy exactly one cache line (64 bytes = 16 F32 elements) using wide loads/stores
+static void copy_cache_aligned_f32(float * dst, const float * src) {
+ __asm__ volatile(
+ "flq2 f0, 0(%[src]) \n\t" // Load 32 bytes
+ "flq2 f1, 32(%[src]) \n\t" // Load next 32 bytes
+ "fsq2 f0, 0(%[dst]) \n\t" // Store 32 bytes
+ "fsq2 f1, 32(%[dst]) \n\t" // Store next 32 bytes
+ :
+ : [src] "r"(src), [dst] "r"(dst)
+ : "f0", "f1", "memory");
+}
+
+// Convert and copy one dst cache line worth of F32->F16 (32 elements src -> 64 bytes dst)
+static void copy_cache_aligned_f16(uint16_t * dst, const float * src) {
+ unsigned long mask_temp;
+
+ // Build offset vector for consecutive 16-bit stores: [0, 2, 4, 6, 8, 10, 12, 14]
+ float offset_vec_storage[8];
+ uint32_t * offsets = (uint32_t *) offset_vec_storage;
+ for (int j = 0; j < 8; j++) {
+ offsets[j] = j * 2;
+ }
+
+ __asm__ volatile(
+ "mova.x.m %[mask_temp] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "flw.ps f1, 0(%[offsets]) \n\t"
+ : [mask_temp] "=&r"(mask_temp)
+ : [offsets] "r"(offset_vec_storage)
+ : "f1");
+
+ // 4 iterations of 8 elements = 32 F16 elements = 64 bytes = 1 cache line
+ for (int i = 0; i < 32; i += 8) {
+ __asm__ volatile(
+ "flw.ps f2, 0(%[src_ptr]) \n\t"
+ "fcvt.f16.ps f3, f2 \n\t"
+ "fsch.ps f3, f1(%[dst_ptr]) \n\t"
+ :
+ : [src_ptr] "r"(src + i), [dst_ptr] "r"(dst + i)
+ : "f2", "f3", "memory");
+ }
+
+ __asm__ volatile("mova.m.x %[mask_temp] \n\t" : : [mask_temp] "r"(mask_temp));
+}
+
+static inline size_t tensor_bytes(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+int entry_point(struct ggml_et_set_rows_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0; // Source data tensor (F32)
+ struct ggml_tensor * src1 = ¶ms->src1; // Row indices tensor (I64)
+ struct ggml_tensor * dst = ¶ms->dst; // Destination tensor (F32/F16)
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_I64) {
+ return -1; // Invalid source types
+ }
+
+ if (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16) {
+ return -1; // Unsupported destination type
+ }
+
+ float * src0_data = (float *) src0->data;
+ int64_t * src1_data = (int64_t *) src1->data;
+ void * dst_data = dst->data;
+
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ const int64_t ne00 = src0->ne[0]; // Source columns (row width)
+ const int64_t ne01 = src0->ne[1]; // Source rows (number of rows to write)
+ const int64_t ne02 = src0->ne[2]; // Source batch dimension
+ const int64_t ne03 = src0->ne[3]; // Source outer batch dimension
+
+ const int64_t nb01 = src0->nb[1];
+ const int64_t nb02 = src0->nb[2];
+ const int64_t nb03 = src0->nb[3];
+
+ const int64_t ne10 = src1->ne[0]; // Number of indices in dimension 0
+ const int64_t ne11 = src1->ne[1]; // Number of indices in dimension 1
+ const int64_t ne12 = src1->ne[2]; // Batch dimension for indices
+
+ const int64_t nb10 = src1->nb[0];
+ const int64_t nb11 = src1->nb[1];
+ const int64_t nb12 = src1->nb[2];
+
+ const int64_t ne_dst1 = dst->ne[1]; // Number of rows in destination (for bounds checking)
+
+ const int64_t nb1 = dst->nb[1];
+ const int64_t nb2 = dst->nb[2];
+ const int64_t nb3 = dst->nb[3];
+
+ // Validate that number of indices matches number of source rows
+ if (ne10 != ne01) {
+ return -1; // Number of indices must match number of source rows
+ }
+#ifdef ET_UBERKERNEL
+ evict_region_past_l2(params->src0.data, tensor_bytes(¶ms->src0));
+ evict_region_past_l2(params->src1.data, tensor_bytes(¶ms->src1));
+ FENCE;
+ et_barrier(ET_BARRIER_GLOBAL);
+#endif
+ const int64_t total_rows = ne01 * ne02 * ne03;
+
+ // Determine cache-line element count based on destination type
+ const int64_t dst_cl_elems = (dst->type == GGML_TYPE_F16) ? CACHE_LINE_F16_ELEMS : CACHE_LINE_F32_ELEMS;
+
+ // Check if rows are cache-line aligned in the destination
+ const bool row_cache_aligned = (ne00 >= dst_cl_elems) && (ne00 % dst_cl_elems == 0);
+
+ if (row_cache_aligned) {
+ // Cache-aligned path: distribute dst cache lines across threads
+ // Each thread owns complete cache lines -> no coherence conflicts
+ const int64_t cls_per_row = ne00 / dst_cl_elems;
+ const int64_t total_cls = total_rows * cls_per_row;
+ const int64_t cls_per_thread = (total_cls + num_threads - 1) / num_threads;
+ const int64_t my_start = thread_id * cls_per_thread;
+ int64_t my_end = my_start + cls_per_thread;
+ if (my_end > total_cls) {
+ my_end = total_cls;
+ }
+ if (my_start >= total_cls) {
+ return 0;
+ }
+
+ for (int64_t cl = my_start; cl < my_end; cl++) {
+ // Map flat cache-line index -> (row, offset within row)
+ const int64_t row_flat = cl / cls_per_row;
+ const int64_t cl_in_row = cl % cls_per_row;
+
+ // Decompose flat row -> (i03, i02, i01)
+ const int64_t i01 = row_flat % ne01;
+ const int64_t tmp = row_flat / ne01;
+ const int64_t i02 = tmp % ne02;
+ const int64_t i03 = tmp / ne02;
+
+ // Look up destination row index
+ const int64_t i12 = i03 % ne12;
+ const int64_t i11 = i02 % ne11;
+ const int64_t i10 = i01;
+ const int64_t index_byte_offset = i10 * nb10 + i11 * nb11 + i12 * nb12;
+ const int64_t dst_row_index = *(int64_t *) ((char *) src1_data + index_byte_offset);
+
+ if (dst_row_index < 0 || dst_row_index >= ne_dst1) {
+ return -1;
+ }
+
+ // Source pointer: row base + cache-line offset (always F32 source)
+ const int64_t elem_offset = cl_in_row * dst_cl_elems;
+ const float * src_ptr =
+ (const float *) ((char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03) + elem_offset;
+
+ // Destination pointer: scattered row base + cache-line offset
+ char * dst_row_base = (char *) dst_data + dst_row_index * nb1 + i02 * nb2 + i03 * nb3;
+
+ if (dst->type == GGML_TYPE_F32) {
+ float * dst_ptr = (float *) dst_row_base + elem_offset;
+ copy_cache_aligned_f32(dst_ptr, src_ptr);
+ } else {
+ uint16_t * dst_ptr = (uint16_t *) dst_row_base + elem_offset;
+ copy_cache_aligned_f16(dst_ptr, src_ptr);
+ }
+ }
+ } else if (nb1 % CACHE_LINE_SIZE_BYTES == 0) {
+ // LCM-aligned path: destination row stride is cache-line-aligned, so
+ // scattered rows never share a cache line even though ne00 doesn't
+ // fill complete cache lines. Group rows via lcm(ne00, dst_cl_elems)
+ // and distribute cache lines across threads — each thread exclusively
+ // owns its cache lines, so normal stores are safe (no atomics needed).
+ const int64_t g = gcd64(ne00, dst_cl_elems);
+ const int64_t rows_per_group = dst_cl_elems / g; // lcm / ne00
+ const int64_t cls_per_group = ne00 / g; // lcm / dst_cl_elems
+
+ const int64_t total_groups = (total_rows + rows_per_group - 1) / rows_per_group;
+ const int64_t total_cls = total_groups * cls_per_group;
+ const int64_t cls_per_thread = (total_cls + num_threads - 1) / num_threads;
+ const int64_t my_start = thread_id * cls_per_thread;
+ int64_t my_end = my_start + cls_per_thread;
+ if (my_end > total_cls) {
+ my_end = total_cls;
+ }
+ if (my_start >= total_cls) {
+ return 0;
+ }
+
+#ifdef BUILD_FOR_UBERKERNEL
+ et_barrier(ET_BARRIER_GLOBAL);
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+ // evict_region_past_l2(src1_data, tensor_bytes(src1));
+ // // et_barrier(ET_BARRIER_GLOBAL);
+ // FENCE;
+#endif
+
+
+ for (int64_t cl = my_start; cl < my_end; cl++) {
+ const int64_t group_idx = cl / cls_per_group;
+ const int64_t cl_in_group = cl % cls_per_group;
+
+ // Element range [elem_start, elem_end) within the flattened group
+ const int64_t elem_start = cl_in_group * dst_cl_elems;
+ const int64_t elem_end = elem_start + dst_cl_elems;
+
+ // Which row(s) inside this group does the cache line touch?
+ const int64_t r_first = elem_start / ne00;
+ const int64_t r_last = (elem_end - 1) / ne00;
+
+ for (int64_t r = r_first; r <= r_last; r++) {
+ const int64_t row_flat = group_idx * rows_per_group + r;
+ if (row_flat >= total_rows) {
+ break;
+ }
+
+ // Column range within this row
+ int64_t col_begin = (r == r_first) ? (elem_start - r * ne00) : 0;
+ int64_t col_end = (r == r_last) ? (elem_end - r * ne00) : ne00;
+ if (col_end > ne00) {
+ col_end = ne00;
+ }
+
+ // Decompose flat row -> (i03, i02, i01)
+ const int64_t i01 = row_flat % ne01;
+ const int64_t tmp = row_flat / ne01;
+ const int64_t i02 = tmp % ne02;
+ const int64_t i03 = tmp / ne02;
+
+ // Look up destination row index
+ const int64_t i12 = i03 % ne12;
+ const int64_t i11 = i02 % ne11;
+ const int64_t i10 = i01;
+ const int64_t index_byte_offset = i10 * nb10 + i11 * nb11 + i12 * nb12;
+ const int64_t dst_row_index = *(int64_t *) ((char *) src1_data + index_byte_offset);
+
+ if (dst_row_index < 0 || dst_row_index >= ne_dst1) {
+ return -1;
+ }
+
+ const float * src_row = (const float *) ((char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ char * dst_row_base = (char *) dst_data + dst_row_index * nb1 + i02 * nb2 + i03 * nb3;
+
+ // nb1 is cache-line-aligned, so dst_row_base is too.
+ // Use aligned copy when the column range fills a complete
+ // cache line at a cache-line-aligned offset within the row.
+ const bool full_cl = (col_begin % dst_cl_elems == 0) && (col_end - col_begin == dst_cl_elems);
+
+ if (dst->type == GGML_TYPE_F32) {
+ float * dp = (float *) dst_row_base;
+ if (full_cl) {
+ copy_cache_aligned_f32(dp + col_begin, src_row + col_begin);
+ } else {
+ for (int64_t i = col_begin; i < col_end; i++) {
+ dp[i] = src_row[i];
+ }
+ }
+ } else {
+ uint16_t * dp = (uint16_t *) dst_row_base;
+ if (full_cl) {
+ copy_cache_aligned_f16(dp + col_begin, src_row + col_begin);
+ } else {
+ for (int64_t i = col_begin; i < col_end; i++) {
+ dp[i] = fp32_to_fp16(src_row[i]);
+ }
+ }
+ }
+ }
+ }
+
+#ifdef BUILD_FOR_UBERKERNEL
+ et_barrier(ET_BARRIER_GLOBAL);
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+ // evict_region_past_l2(src1_data, tensor_bytes(src1));
+ // // et_barrier(ET_BARRIER_GLOBAL);
+ // FENCE;
+#endif
+
+
+ } else {
+ // Fallback: nb1 not cache-line-aligned, so scattered destination rows
+ // may share a cache line. Use atomic global stores to bypass L1D.
+ for (int64_t row_flat = thread_id; row_flat < total_rows; row_flat += num_threads) {
+ const int64_t i01 = row_flat % ne01;
+ const int64_t tmp = row_flat / ne01;
+ const int64_t i02 = tmp % ne02;
+ const int64_t i03 = tmp / ne02;
+
+ // Look up destination row index
+ const int64_t i12 = i03 % ne12;
+ const int64_t i11 = i02 % ne11;
+ const int64_t i10 = i01;
+ const int64_t index_byte_offset = i10 * nb10 + i11 * nb11 + i12 * nb12;
+ const int64_t dst_row_index = *(int64_t *) ((char *) src1_data + index_byte_offset);
+
+ if (dst_row_index < 0 || dst_row_index >= ne_dst1) {
+ return -1;
+ }
+
+ const float * src_row = (const float *) ((char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ char * dst_row_base = (char *) dst_data + dst_row_index * nb1 + i02 * nb2 + i03 * nb3;
+
+ if (dst->type == GGML_TYPE_F32) {
+ volatile float * dst_row = (volatile float *) dst_row_base;
+ for (int64_t i = 0; i < ne00; i++) {
+ atomic_store_f32(dst_row + i, src_row[i]);
+ }
+ } else {
+ volatile uint16_t * dst_row = (volatile uint16_t *) dst_row_base;
+ for (int64_t i = 0; i < ne00; i++) {
+ atomic_store_f16(dst_row + i, fp32_to_fp16(src_row[i]));
+ }
+ }
+ }
+ }
+
+#ifdef BUILD_FOR_UBERKERNEL
+ et_barrier(ET_BARRIER_GLOBAL);
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+ // evict_region_past_l2(src1_data, tensor_bytes(src1));
+ // // et_barrier(ET_BARRIER_GLOBAL);
+ // FENCE;
+#endif
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Bare Metal Softmax F32 Kernel
+// Softmax function: y[i] = exp(x[i] - max) / sum(exp(x[j] - max))
+//
+// Algorithm:
+// 1. Apply scaling: x' = x * scale
+// 2. Add mask/bias if present: x' = x' + mask * slope (ALiBi support)
+// 3. Find max value for numerical stability: max = max(x')
+// 4. Compute exponentials: exp_vals[i] = exp(x'[i] - max)
+// 5. Compute sum: sum = sum(exp_vals)
+// 6. Normalize: y[i] = exp_vals[i] / sum
+//
+// Features supported:
+// - Temperature scaling via scale parameter
+// - Attention masking (transformer masks)
+// - ALiBi (Attention with Linear Biases) positional encoding
+// - Numerical stability (subtract max before exp)
+// - ggml broadcasting rules for mask tensors
+//
+// Mask Broadcasting Rules (ggml-specific, not standard numpy):
+// - Dimension 0: mask.ne[0] == input.ne[0] (exact match required)
+// - Dimension 1: mask.ne[1] >= input.ne[1] (allows larger pre-allocated masks)
+// - Dimension 2: input.ne[2] % mask.ne[2] == 0 (modulo broadcasting)
+// - Dimension 3: input.ne[3] % mask.ne[3] == 0 (modulo broadcasting)
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <assert.h>
+#include <math.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+// Softmax kernel parameters structure (from ggml-et-ops.h)
+struct ggml_et_softmax_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor src1; // F32 mask tensor (optional, may be zeroed if not used)
+ struct ggml_tensor src2; // F32 sinks tensor (optional, may be zeroed if not used)
+ struct ggml_tensor dst; // F32 output tensor
+ float scale; // Scale factor (temperature scaling)
+ float max_bias; // Max bias for ALiBi (0.0f if not used)
+};
+
+#define LOG2E_F 1.4426950408889634f
+
+typedef struct {
+ float max_val;
+ float sum_val;
+ uint32_t valid_mask;
+} softmax_params_t;
+
+static inline bool softmax_lane_is_valid(float x) {
+ return (x == x) && (x != -INFINITY) && (x != INFINITY);
+}
+
+static inline softmax_params_t softmax_params_empty(void) {
+ softmax_params_t p;
+ p.max_val = -INFINITY;
+ p.sum_val = 0.0f;
+ p.valid_mask = 0;
+ return p;
+}
+
+// chunk_transform_ps_8_branchless_mask
+//
+// Vector transform for 8 logits:
+//
+// x = src * scale + (mask ? mask * slope : 0)
+//
+// Implemented branchlessly so masked and unmasked paths share the same
+// instruction stream. Used by pass1 and pass2 vector loops.
+static inline void chunk_transform_ps_8_branchless_mask(float * tmp8,
+ const float * src,
+ const float * mask,
+ float scale,
+ float slope) {
+ unsigned long ms;
+ const float zero = 0.0f;
+ const unsigned long mask_load_m0 = (mask != NULL) ? 0xFFul : 0x00ul;
+ const float * mp = (mask != NULL) ? mask : &zero;
+
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f10, 0(%[p_scale]) \n\t"
+ "fbc.ps f11, 0(%[p_slope]) \n\t"
+ "fbc.ps f1, 0(%[p_zero]) \n\t"
+
+ "mov.m.x m0, %[maskm0], 0 \n\t" // load mask if needed
+ "flw.ps f1, 0(%[mp]) \n\t"
+
+ "mov.m.x m0, x0, 0xFF \n\t"
+
+ "flw.ps f0, 0(%[sp]) \n\t"
+ "fmul.ps f0, f0, f10 \n\t"
+ "fmul.ps f1, f1, f11 \n\t"
+ "fadd.ps f0, f0, f1, rne \n\t"
+ "fsw.ps f0, 0(%[tp]) \n\t"
+
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(ms)
+ : [tp] "r"(tmp8), [sp] "r"(src), [mp] "r"(mp), [p_zero] "r"(&zero), [p_scale] "r"(&scale),
+ [p_slope] "r"(&slope), [maskm0] "r"(mask_load_m0)
+ : "f0", "f1", "f10", "f11", "memory");
+}
+
+// chunk_transform_ps_8_tail
+//
+// Same as chunk_transform_ps_8_branchless_mask but gates loads, compute,
+// and stores with a caller-supplied m0 mask so that only `count` elements
+// (1-7) are touched. Used for the last sub-8 chunk of a non-aligned row.
+static inline void chunk_transform_ps_8_tail(float * tmp8,
+ const float * src,
+ const float * mask,
+ float scale,
+ float slope,
+ unsigned long tail_m0) {
+ unsigned long ms;
+ const float zero = 0.0f;
+ const unsigned long mask_load_m0 = (mask != NULL) ? tail_m0 : 0x00ul;
+ const float * mp = (mask != NULL) ? mask : &zero;
+
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+
+ // Broadcast constants with all lanes enabled
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f10, 0(%[p_scale]) \n\t"
+ "fbc.ps f11, 0(%[p_slope]) \n\t"
+ "fbc.ps f1, 0(%[p_zero]) \n\t"
+
+ // Load mask data gated by tail mask
+ "mov.m.x m0, %[maskm0], 0 \n\t"
+ "flw.ps f1, 0(%[mp]) \n\t"
+
+ // Load source, compute, and store gated by tail mask
+ "mov.m.x m0, %[tailm0], 0 \n\t"
+
+ "flw.ps f0, 0(%[sp]) \n\t"
+ "fmul.ps f0, f0, f10 \n\t"
+ "fmul.ps f1, f1, f11 \n\t"
+ "fadd.ps f0, f0, f1, rne \n\t"
+ "fsw.ps f0, 0(%[tp]) \n\t"
+
+ "mova.m.x %[ms] \n\t"
+ : [ms] "=&r"(ms)
+ : [tp] "r"(tmp8), [sp] "r"(src), [mp] "r"(mp), [p_zero] "r"(&zero), [p_scale] "r"(&scale),
+ [p_slope] "r"(&slope), [maskm0] "r"(mask_load_m0), [tailm0] "r"(tail_m0)
+ : "f0", "f1", "f10", "f11", "memory");
+}
+
+// softmax_pass1_range
+//
+// Computes the numerically-stable softmax scan over a sub-range of a row.
+//
+// This implements the 1st pass of online softmax
+//
+// max' = max(max, x)
+// sum' = sum * exp(old_max - max') + exp(x - max')
+//
+// and returns a partial result containing:
+//
+// - max_val : maximum logit observed in this range
+// - sum_val : exp-normalized sum relative to max_val
+//
+// These partial results can be merged with softmax_params_merge() to obtain
+// the result for the full row.
+static inline softmax_params_t softmax_pass1_range(const float * src,
+ const float * mask,
+ int begin,
+ int end,
+ float scale,
+ float slope) {
+ __attribute__((aligned(32))) float lane_max[8];
+ __attribute__((aligned(32))) float lane_sum[8];
+ __attribute__((aligned(32))) float tmp[8];
+
+ uint8_t valid_mask = 0;
+
+ const float one_f = 1.0f;
+ const float zero_f = 0.0f;
+ const float neg_inf = -INFINITY;
+ const float log2e = LOG2E_F;
+
+ unsigned long ms;
+
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f20, 0(%[p_ninf]) \n\t"
+ "fbc.ps f21, 0(%[p_zero]) \n\t"
+ "fbc.ps f22, 0(%[p_one]) \n\t"
+ "fbc.ps f23, 0(%[p_log2e]) \n\t"
+ : [ms] "=&r"(ms)
+ : [p_ninf] "r"(&neg_inf), [p_zero] "r"(&zero_f), [p_one] "r"(&one_f), [p_log2e] "r"(&log2e)
+ : "f20", "f21", "f22", "f23");
+
+ const int aligned_end = begin + ((end - begin) & ~7);
+
+ // Process full 8-element chunks
+ int i = begin;
+ for (; i < aligned_end; i += 8) {
+ chunk_transform_ps_8_branchless_mask(tmp, src + i, mask ? (mask + i) : NULL, scale, slope);
+
+ uint8_t cur_mask = 0;
+ for (int j = 0; j < 8; ++j) {
+ if (softmax_lane_is_valid(tmp[j])) {
+ cur_mask |= (uint8_t) (1u << j);
+ }
+ }
+
+ const uint8_t init_mask = (uint8_t) (cur_mask & ~valid_mask);
+ const uint8_t upd_mask = (uint8_t) (cur_mask & valid_mask);
+
+ if (init_mask || upd_mask) {
+ __asm__ volatile(
+ "flw.ps f0, 0(%[p_tmp]) \n\t"
+
+ "mov.m.x m0, %[initm], 0 \n\t"
+ "fcmovm.ps f20, f0, f20 \n\t"
+ "fcmovm.ps f21, f22, f21 \n\t"
+
+ "mov.m.x m0, %[updm], 0 \n\t"
+ "fmax.ps f1, f20, f0 \n\t"
+
+ "fsub.ps f2, f20, f1, rne \n\t"
+ "fmul.ps f2, f2, f23 \n\t"
+ "fexp.ps f2, f2 \n\t"
+
+ "fsub.ps f3, f0, f1, rne \n\t"
+ "fmul.ps f3, f3, f23 \n\t"
+ "fexp.ps f3, f3 \n\t"
+
+ "fmul.ps f21, f21, f2 \n\t"
+ "fadd.ps f21, f21, f3, rne \n\t"
+ "fcmovm.ps f20, f1, f20 \n\t"
+
+ "mov.m.x m0, x0, 0xFF \n\t"
+ :
+ : [p_tmp] "r"(tmp), [initm] "r"((unsigned long) init_mask), [updm] "r"((unsigned long) upd_mask)
+ : "f0", "f1", "f2", "f3", "memory");
+
+ valid_mask |= cur_mask;
+ }
+ }
+
+ // Tail chunk: m0-gated load/compute/store for remaining 1-7 elements
+ if (i < end) {
+ const unsigned long tail_m0 = (1ul << (end - i)) - 1;
+
+ // Fill tmp with NaN so invalid lanes fail softmax_lane_is_valid
+ for (int j = 0; j < 8; j++) {
+ tmp[j] = __builtin_nanf("");
+ }
+
+ chunk_transform_ps_8_tail(tmp, src + i, mask ? (mask + i) : NULL, scale, slope, tail_m0);
+
+ uint8_t cur_mask = 0;
+ for (int j = 0; j < 8; ++j) {
+ if (softmax_lane_is_valid(tmp[j])) {
+ cur_mask |= (uint8_t) (1u << j);
+ }
+ }
+
+ const uint8_t init_mask = (uint8_t) (cur_mask & ~valid_mask);
+ const uint8_t upd_mask = (uint8_t) (cur_mask & valid_mask);
+
+ if (init_mask || upd_mask) {
+ __asm__ volatile(
+ "flw.ps f0, 0(%[p_tmp]) \n\t"
+
+ "mov.m.x m0, %[initm], 0 \n\t"
+ "fcmovm.ps f20, f0, f20 \n\t"
+ "fcmovm.ps f21, f22, f21 \n\t"
+
+ "mov.m.x m0, %[updm], 0 \n\t"
+ "fmax.ps f1, f20, f0 \n\t"
+
+ "fsub.ps f2, f20, f1, rne \n\t"
+ "fmul.ps f2, f2, f23 \n\t"
+ "fexp.ps f2, f2 \n\t"
+
+ "fsub.ps f3, f0, f1, rne \n\t"
+ "fmul.ps f3, f3, f23 \n\t"
+ "fexp.ps f3, f3 \n\t"
+
+ "fmul.ps f21, f21, f2 \n\t"
+ "fadd.ps f21, f21, f3, rne \n\t"
+ "fcmovm.ps f20, f1, f20 \n\t"
+
+ "mov.m.x m0, x0, 0xFF \n\t"
+ :
+ : [p_tmp] "r"(tmp), [initm] "r"((unsigned long) init_mask), [updm] "r"((unsigned long) upd_mask)
+ : "f0", "f1", "f2", "f3", "memory");
+
+ valid_mask |= cur_mask;
+ }
+ }
+
+ __asm__ volatile(
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fsw.ps f20, 0(%[p_lmax]) \n\t"
+ "fsw.ps f21, 0(%[p_lsum]) \n\t"
+ "mova.m.x %[ms] \n\t"
+ :
+ : [p_lmax] "r"(lane_max), [p_lsum] "r"(lane_sum), [ms] "r"(ms)
+ : "memory");
+
+ softmax_params_t out = softmax_params_empty();
+ out.valid_mask = valid_mask;
+
+ for (int k = 0; k < 8; ++k) {
+ if (valid_mask & (1u << k)) {
+ if (out.valid_mask == (1u << k) || out.max_val == -INFINITY || lane_max[k] > out.max_val) {
+ out.max_val = lane_max[k];
+ }
+ }
+ }
+
+ if (out.max_val != -INFINITY) {
+ // Compute lane correction factors via fexp.ps to stay consistent
+ // with the fexp.ps used inside the online softmax loop above.
+ // corr[k] = exp2((lane_max[k] - out.max_val) * LOG2E) = exp(lane_max[k] - out.max_val)
+ const float neg_max_l2 = -out.max_val * LOG2E_F;
+ __attribute__((aligned(32))) float corr[8];
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f0, 0(%[p_nml2]) \n\t"
+ "fbc.ps f2, 0(%[p_l2e]) \n\t"
+ "flw.ps f1, 0(%[p_lmax]) \n\t"
+ "fmadd.ps f0, f1, f2, f0 \n\t"
+ "fexp.ps f0, f0 \n\t"
+ "fsw.ps f0, 0(%[p_corr]) \n\t"
+ "mova.m.x %[ms] \n\t"
+ :
+ : [p_nml2] "r"(&neg_max_l2), [p_l2e] "r"(&log2e), [p_lmax] "r"(lane_max), [p_corr] "r"(corr), [ms] "r"(ms)
+ : "f0", "f1", "f2", "memory");
+ for (int k = 0; k < 8; ++k) {
+ if (valid_mask & (1u << k)) {
+ out.sum_val += lane_sum[k] * corr[k];
+ }
+ }
+ }
+
+ return out;
+}
+
+// Pass 2 (normalize) over [begin, end).
+//
+// Computes: dst[i] = exp(x[i]*scale + mask[i]*slope - max) / sum
+//
+// Uses fexp.ps for the numerator; the denominator (params.sum_val) must
+// already be fully computed by the caller (pass1 + any sink merge).
+static inline void softmax_pass2_range(float * dst,
+ const float * src,
+ const float * mask,
+ int begin,
+ int end,
+ float scale,
+ float slope,
+ softmax_params_t params) {
+ const float s2 = scale * LOG2E_F;
+ const float sl2 = slope * LOG2E_F;
+ const float neg_ml2 = -params.max_val * LOG2E_F;
+ const float inv_sum = et_fdiv(1.0f, params.sum_val);
+
+ unsigned long ms;
+
+ __asm__ volatile(
+ "mova.x.m %[ms] \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ "fbc.ps f10, 0(%[p_s2]) \n\t"
+ "fbc.ps f12, 0(%[p_nml2]) \n\t"
+ "fbc.ps f13, 0(%[p_inv]) \n\t"
+ : [ms] "=&r"(ms)
+ : [p_s2] "r"(&s2), [p_nml2] "r"(&neg_ml2), [p_inv] "r"(&inv_sum)
+ : "f10", "f12", "f13");
+
+ const int aligned_end = begin + ((end - begin) & ~7);
+
+ if (mask != NULL) {
+ __asm__ volatile("fbc.ps f11, 0(%[p_sl2]) \n\t" : : [p_sl2] "r"(&sl2) : "f11");
+
+ for (int c = begin; c < aligned_end; c += 8) {
+ __asm__ volatile(
+ "flw.ps f0, 0(%[sp]) \n\t"
+ "flw.ps f1, 0(%[mp]) \n\t"
+ "fmadd.ps f0, f0, f10, f12 \n\t"
+ "fmadd.ps f0, f1, f11, f0 \n\t"
+ "fexp.ps f0, f0 \n\t"
+ "fmul.ps f0, f0, f13 \n\t"
+ "fsw.ps f0, 0(%[dp]) \n\t"
+ :
+ : [sp] "r"(src + c), [mp] "r"(mask + c), [dp] "r"(dst + c)
+ : "f0", "f1", "memory");
+ }
+
+ // Tail chunk with m0 gating
+ if (aligned_end < end) {
+ const unsigned long tail_m0 = (1ul << (end - aligned_end)) - 1;
+ __asm__ volatile(
+ "mov.m.x m0, %[tm], 0 \n\t"
+ "flw.ps f0, 0(%[sp]) \n\t"
+ "flw.ps f1, 0(%[mp]) \n\t"
+ "fmadd.ps f0, f0, f10, f12 \n\t"
+ "fmadd.ps f0, f1, f11, f0 \n\t"
+ "fexp.ps f0, f0 \n\t"
+ "fmul.ps f0, f0, f13 \n\t"
+ "fsw.ps f0, 0(%[dp]) \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ :
+ : [sp] "r"(src + aligned_end), [mp] "r"(mask + aligned_end), [dp] "r"(dst + aligned_end),
+ [tm] "r"(tail_m0)
+ : "f0", "f1", "memory");
+ }
+ } else {
+ for (int c = begin; c < aligned_end; c += 8) {
+ __asm__ volatile(
+ "flw.ps f0, 0(%[sp]) \n\t"
+ "fmadd.ps f0, f0, f10, f12 \n\t"
+ "fexp.ps f0, f0 \n\t"
+ "fmul.ps f0, f0, f13 \n\t"
+ "fsw.ps f0, 0(%[dp]) \n\t"
+ :
+ : [sp] "r"(src + c), [dp] "r"(dst + c)
+ : "f0", "memory");
+ }
+
+ // Tail chunk with m0 gating
+ if (aligned_end < end) {
+ const unsigned long tail_m0 = (1ul << (end - aligned_end)) - 1;
+ __asm__ volatile(
+ "mov.m.x m0, %[tm], 0 \n\t"
+ "flw.ps f0, 0(%[sp]) \n\t"
+ "fmadd.ps f0, f0, f10, f12 \n\t"
+ "fexp.ps f0, f0 \n\t"
+ "fmul.ps f0, f0, f13 \n\t"
+ "fsw.ps f0, 0(%[dp]) \n\t"
+ "mov.m.x m0, x0, 0xFF \n\t"
+ :
+ : [sp] "r"(src + aligned_end), [dp] "r"(dst + aligned_end), [tm] "r"(tail_m0)
+ : "f0", "memory");
+ }
+ }
+
+ __asm__ volatile("mova.m.x %[ms] \n\t" ::[ms] "r"(ms));
+}
+
+// Single-core row path.
+// pass1_range and pass2_range handle non-8-aligned cols internally via
+// m0-gated tail chunks, so this function just passes cols directly.
+static inline void compute_softmax_row(float * dst,
+ const float * src,
+ const float * mask,
+ int cols,
+ float scale,
+ float slope,
+ float sink_value,
+ bool use_sinks) {
+ softmax_params_t params = softmax_pass1_range(src, mask, 0, cols, scale, slope);
+
+ if (use_sinks) {
+ // For sinks, use fully scalar et_expf to match the reference CPU
+ // backend's expf precision. Sink tests use small arrays (ne<=32)
+ // so the scalar path has negligible performance impact.
+ float max_val = params.max_val;
+ if (sink_value > max_val) {
+ max_val = sink_value;
+ }
+
+ // Compute sum = Σ exp(x'[i] - max) + exp(sink - max) (scalar)
+ float sum = 0.0f;
+ for (int i = 0; i < cols; ++i) {
+ float x = src[i] * scale;
+ if (mask != NULL) {
+ x += mask[i] * slope;
+ }
+ sum += et_expf(x - max_val);
+ }
+ sum += et_expf(sink_value - max_val);
+
+ // Normalize: dst[i] = exp(x'[i] - max) / sum (scalar)
+ float inv_sum = et_fdiv(1.0f, sum);
+ for (int i = 0; i < cols; ++i) {
+ float x = src[i] * scale;
+ if (mask != NULL) {
+ x += mask[i] * slope;
+ }
+ dst[i] = et_expf(x - max_val) * inv_sum;
+ }
+ } else {
+ if (!params.valid_mask) {
+ return;
+ }
+ softmax_pass2_range(dst, src, mask, 0, cols, scale, slope, params);
+ }
+}
+
+// Main entry point for Softmax kernel
+int entry_point(struct ggml_et_softmax_params * params, void * env) {
+ // Cast env to proper type
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ // Validate environment pointer
+ if (!kernel_env) {
+ return -1;
+ }
+
+ // Get thread info using shire mask from environment
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ // Return early if this hart is not active
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ // Basic safety check on params
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ // Extract tensor references
+ struct ggml_tensor * src0 = ¶ms->src0; // Input tensor
+ struct ggml_tensor * src1 = ¶ms->src1; // Mask tensor (optional)
+ struct ggml_tensor * src2 = ¶ms->src2; // Sinks tensor (optional)
+ struct ggml_tensor * dst = ¶ms->dst; // Output tensor
+ float scale = params->scale; // Scale factor
+ float max_bias = params->max_bias; // ALiBi max bias
+
+ // Validate tensor types (F32 only)
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ // Check if mask is used and validate type
+ bool use_mask = (src1->data != NULL && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16));
+
+ bool use_sinks = (src2->data != NULL && src2->type == GGML_TYPE_F32);
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+ float * mask_data = use_mask ? (float *) src1->data : NULL;
+ float * sinks_data = use_sinks ? (float *) src2->data : NULL;
+
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+
+ const int64_t ne00 = src0->ne[0]; // Sequence length (columns)
+ const int64_t ne01 = src0->ne[1]; // Number of rows
+ const int64_t ne02 = src0->ne[2]; // Batch/head dimension
+ const int64_t ne03 = src0->ne[3]; // Outer batch dimension
+
+ // Fast path: softmax of a single element is always 1.0
+ // (exp(x) / exp(x) == 1 for any x, regardless of scale/mask/bias)
+ // Skip all ALiBi, mask, and sink setup.
+ //
+ // Each output element is 4 bytes. A cache line is 64 bytes = 16 floats.
+ // L1 is not coherent across harts, so each thread must own whole cache
+ // lines to avoid cross-hart conflicts.
+ if (ne00 == 1) {
+ const int64_t total_elems = ne01 * ne02 * ne03;
+ const int64_t elems_per_cl = ET_CACHE_LINE_SIZE_BYTES / (int64_t) sizeof(float); // 16
+ const int64_t total_cls = (total_elems + elems_per_cl - 1) / elems_per_cl;
+
+ for (int64_t cl = thread_id; cl < total_cls; cl += num_threads) {
+ const int64_t start = cl * elems_per_cl;
+ int64_t end = start + elems_per_cl;
+ if (end > total_elems) {
+ end = total_elems;
+ }
+ for (int64_t idx = start; idx < end; idx++) {
+ dst_data[idx] = 1.0f;
+ }
+ }
+ return 0;
+ }
+
+ const int64_t ne10 = use_mask ? src1->ne[0] : 0; // Mask sequence length
+ const int64_t ne11 = use_mask ? src1->ne[1] : 0; // Mask rows
+ const int64_t ne12 = use_mask ? src1->ne[2] : 0; // Mask batch/head dimension
+ const int64_t ne13 = use_mask ? src1->ne[3] : 0; // Mask outer batch dimension
+
+ if (use_mask) {
+ // - Dimension 0: mask must equal input exactly
+ // - Dimension 1: mask must be >= input (allows larger pre-allocated masks)
+ // - Dimension 2: input must be divisible by mask (modulo broadcasting)
+ // - Dimension 3: input must be divisible by mask (modulo broadcasting)
+ if (ne10 != ne00 || // Dimension 0: exact match required
+ ne11 < ne01 || // Dimension 1: mask >= input
+ (ne12 > 0 && ne02 % ne12 != 0) || // Dimension 2: input % mask == 0
+ (ne13 > 0 && ne03 % ne13 != 0)) { // Dimension 3: input % mask == 0
+ return -1; // Incompatible dimensions for ggml softmax broadcasting
+ }
+ }
+
+ // ALiBi slope calculation - compute per attention head
+ const uint32_t n_head = (uint32_t) ne02;
+ uint32_t n_head_log2 = 0;
+ float m0 = 1.0f;
+ float m1 = 1.0f;
+
+ if (max_bias > 0.0f) {
+ // This is equivalent to: 1 << floor(log2(n_head))
+ n_head_log2 = 1;
+ while (n_head_log2 < n_head) {
+ n_head_log2 <<= 1;
+ }
+ if (n_head_log2 > n_head) {
+ n_head_log2 >>= 1;
+ }
+
+ // Compute base slopes for ALiBi
+ // m0 = 2^(-max_bias / n_head_log2)
+ // m1 = 2^(-max_bias / (2 * n_head_log2))
+ float inv_n_head_log2 = et_fdiv(1.0f, (float) n_head_log2);
+ m0 = et_expf(-max_bias * 0.69314718f * inv_n_head_log2); // 0.69314718 = ln(2)
+ m1 = et_expf(-max_bias * 0.69314718f * inv_n_head_log2 * 0.5f);
+ }
+
+ // Process tensor row by row in parallel across flattened rows.
+ // Flattened row index spans [i03, i02, i01] with row length ne00.
+ //
+ // When ne00 * sizeof(float) is not a multiple of the cache line size,
+ // adjacent rows share cache lines. Assign contiguous write groups to
+ // each thread so every thread's write footprint covers whole cache
+ // lines, preventing cross-hart L1 coherency issues. When rows ARE
+ // cache-line aligned, rows_per_wg == 1 and this degenerates to the
+ // original stride-by-num_threads distribution.
+ const int64_t rows_per_i03 = ne02 * ne01;
+ const int64_t total_rows = ne03 * rows_per_i03;
+ const int64_t rows_per_wg = et_rows_per_cacheline_group(ne00, sizeof(float));
+ const int64_t total_wgs = (total_rows + rows_per_wg - 1) / rows_per_wg;
+
+ for (int64_t wg = thread_id; wg < total_wgs; wg += num_threads) {
+ const int64_t row_start = wg * rows_per_wg;
+ int64_t row_end = row_start + rows_per_wg;
+ if (row_end > total_rows) {
+ row_end = total_rows;
+ }
+
+ for (int64_t row = row_start; row < row_end; row++) {
+ const int64_t i03 = row / rows_per_i03;
+ const int64_t rem = row % rows_per_i03;
+ const int64_t i02 = rem / ne01;
+ const int64_t i01 = rem % ne01;
+
+ // Calculate ALiBi slope for this attention head
+ float slope = 1.0f;
+ if (max_bias > 0.0f) {
+ const uint32_t h = (uint32_t) i02; // head index
+ if (h < n_head_log2) {
+ slope = m0;
+ for (uint32_t i = 0; i < h; i++) {
+ slope *= m0;
+ }
+ } else {
+ const uint32_t exp = 2 * (h - n_head_log2) + 1;
+ slope = m1;
+ for (uint32_t i = 1; i < exp; i++) {
+ slope *= m1;
+ }
+ }
+ }
+
+ float sink_value = 0.0f;
+ if (use_sinks && sinks_data) {
+ sink_value = sinks_data[i02];
+ }
+
+ const int64_t src_offset = i03 * ne02 * ne01 * ne00 + i02 * ne01 * ne00 + i01 * ne00;
+
+ const float * src_row = src0_data + src_offset;
+ float * dst_row = dst_data + src_offset;
+ const float * mask_row = NULL;
+
+ if (use_mask && mask_data) {
+ const int64_t mask_i03 = (ne13 > 0) ? i03 % ne13 : 0;
+ const int64_t mask_i02 = (ne12 > 0) ? i02 % ne12 : 0;
+ const int64_t mask_i01 = i01;
+
+ const int64_t mask_offset = mask_i03 * ne12 * ne11 * ne10 + mask_i02 * ne11 * ne10 + mask_i01 * ne10;
+
+ mask_row = mask_data + mask_offset;
+ }
+
+ compute_softmax_row(dst_row, src_row, mask_row, (int) ne00, scale, slope, sink_value, use_sinks);
+ }
+ }
+
+ return 0; // Success
+}
--- /dev/null
+//******************************************************************************
+// Solve Triangular F32 Kernel
+// Forward substitution: solve AX = B where A is lower-triangular.
+//
+// src0 (A): [n, n, B1, B2] lower-triangular matrix
+// src1 (B): [k, n, B1, B2] right-hand side
+// dst (X): [k, n, B1, B2] solution
+//
+// For each column j (parallelized across threads):
+// For i = 0..n-1:
+// X[i,j] = (B[i,j] - dot(A[i,0..i-1], X[0..i-1,j])) / A[i,i]
+//
+// Lower-triangular, left-side, non-unit variant implemented.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_solve_tri_params {
+ struct ggml_tensor src0; // A: lower-triangular [n, n, B1, B2]
+ struct ggml_tensor src1; // B: RHS [k, n, B1, B2]
+ struct ggml_tensor dst; // X: solution [k, n, B1, B2]
+};
+
+int entry_point(struct ggml_et_solve_tri_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0; // A
+ struct ggml_tensor * src1 = ¶ms->src1; // B
+ struct ggml_tensor * dst = ¶ms->dst; // X
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * A_data = (const float *) src0->data;
+ const float * B_data = (const float *) src1->data;
+ float * X_data = (float *) dst->data;
+
+ if (!A_data || !B_data || !X_data) {
+ return -1;
+ }
+
+ const int64_t n = src0->ne[1]; // A is n×n
+ const int64_t k = src1->ne[0]; // number of RHS columns
+ const int64_t ne2 = src0->ne[2];
+ const int64_t ne3 = src0->ne[3];
+
+ // Strides in bytes
+ const size_t nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3];
+ const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+
+ // k % 16 == 0 guaranteed by supports_op. Rows are cache-line aligned,
+ // so column groups of 16 map to exclusive cache lines.
+ // TODO: Vectorize the thing
+ const int64_t cols_per_cl = 16;
+ const int64_t num_col_groups = k / cols_per_cl;
+ const int64_t total_work = num_col_groups * ne2 * ne3;
+
+ for (int64_t work = thread_id; work < total_work; work += num_threads) {
+ const int64_t cg = work % num_col_groups;
+ const int64_t i2 = (work / num_col_groups) % ne2;
+ const int64_t i3 = work / (num_col_groups * ne2);
+
+ const int64_t j_start = cg * cols_per_cl;
+ const int64_t j_end = j_start + cols_per_cl;
+
+ const float * A_batch = (const float *) ((const char *) A_data + i2 * nb02 + i3 * nb03);
+ const float * B_batch = (const float *) ((const char *) B_data + i2 * nb12 + i3 * nb13);
+ float * X_batch = (float *) ((char *) X_data + i2 * nb2 + i3 * nb3);
+
+ for (int64_t j = j_start; j < j_end; j++) {
+ for (int64_t i = 0; i < n; i++) {
+ const float * A_row = (const float *) ((const char *) A_batch + i * nb01);
+ float * X_row = (float *) ((char *) X_batch + i * nb1);
+ const float * B_row = (const float *) ((const char *) B_batch + i * nb11);
+
+ float sum = 0.0f;
+ for (int64_t t = 0; t < i; t++) {
+ const float * X_t = (const float *) ((const char *) X_batch + t * nb1);
+ sum += A_row[t] * X_t[j];
+ }
+
+ X_row[j] = et_fdiv(B_row[j] - sum, A_row[i]);
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// SQR F32 Kernel
+// Element-wise square: y[i] = x[i] * x[i]
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+// SQR kernel parameters structure (unary op: src0 -> dst)
+struct ggml_et_sqr_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor dst; // F32 output tensor
+};
+
+int entry_point(struct ggml_et_sqr_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1; // Invalid pointer
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1; // Unsupported type combination
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1; // Null data pointer
+ }
+
+ // Both src and dst are contiguous F32: flatten and distribute by cache lines
+ const int64_t total_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+ const int64_t elements_per_cacheline = 16; // 64 bytes / 4 bytes per float
+ const int64_t total_cachelines = (total_elements + elements_per_cacheline - 1) / elements_per_cacheline;
+
+ const int64_t cl_per_thread = (total_cachelines + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cachelines) {
+ cl_end = total_cachelines;
+ }
+
+ if (cl_start >= total_cachelines) {
+ return 0;
+ }
+
+ const int64_t elem_start = cl_start * elements_per_cacheline;
+ int64_t elem_end = cl_end * elements_per_cacheline;
+ if (elem_end > total_elements) {
+ elem_end = total_elements;
+ }
+
+ const float * src_ptr = src0_data + elem_start;
+ float * dst_ptr = dst_data + elem_start;
+ const int32_t count = (int32_t) (elem_end - elem_start);
+
+ // Process 8 elements at a time: dst[i] = src[i] * src[i]
+ for (int32_t i0 = 0; i0 < count; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x_vec]\n" // Load 8 input values
+ "fmul.ps f11, f10, f10\n" // x * x (8-wide)
+ "fsw.ps f11, %[result]\n" // Store 8 results
+
+ : [result] "=m"(*(float (*)[8]) & dst_ptr[i0])
+ : [x_vec] "m"(*(const float (*)[8]) & src_ptr[i0])
+ : "f10", "f11");
+ }
+
+ return 0;
+}
--- /dev/null
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_ssm_conv_params {
+ struct ggml_tensor src0; // conv_x: [d_conv - 1 + n_t, d_inner, n_seqs]
+ struct ggml_tensor src1; // conv1d.weight: [d_conv, d_inner]
+ struct ggml_tensor dst; // output: [d_inner, n_t, n_seqs]
+};
+
+int entry_point(struct ggml_et_ssm_conv_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * src0_data = (const float *) src0->data;
+ const float * src1_data = (const float *) src1->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !src1_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t nc = src1->ne[0];
+ const int64_t ncs = src0->ne[0];
+ const int64_t nr = src0->ne[1];
+ const int64_t n_t = dst->ne[1];
+ const int64_t n_s = dst->ne[2];
+
+ if (dst->ne[0] != nr || src1->ne[1] != nr || ncs != nc - 1 + n_t || src0->nb[0] != sizeof(float) ||
+ src1->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float) || src0->nb[1] != (size_t) ncs * sizeof(float) ||
+ src1->nb[1] != (size_t) nc * sizeof(float)) {
+ return -1;
+ }
+
+ // Parallelize over d_inner in cache-line-aligned chunks (16 floats = 64B)
+ const int64_t chunk = 16;
+ const int64_t n_chunks = (nr + chunk - 1) / chunk;
+
+ // Save and set vector mask to all 8 lanes
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ for (int64_t i3 = 0; i3 < n_s; ++i3) {
+ for (int64_t i2 = 0; i2 < n_t; ++i2) {
+ const float * s = (const float *) ((const char *) src0_data + i2 * src0->nb[0] + i3 * src0->nb[2]);
+ float * x = (float *) ((char *) dst_data + i2 * dst->nb[1] + i3 * dst->nb[2]);
+
+ for (int64_t ci = thread_id; ci < n_chunks; ci += num_threads) {
+ const int64_t i1_start = ci * chunk;
+ const int64_t i1_end = i1_start + chunk < nr ? i1_start + chunk : nr;
+
+ // Process 8 channels at a time with SIMD
+ int64_t i1 = i1_start;
+ for (; i1 + 8 <= i1_end; i1 += 8) {
+ // Gather 8 channels' data into contiguous buffers for each tap
+ float tmp_s[8], tmp_c[8];
+ float acc[8] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
+
+ for (int64_t i0 = 0; i0 < nc; ++i0) {
+ // TODO: Some way to get rid of this gather
+ for (int j = 0; j < 8; ++j) {
+ tmp_s[j] = s[(i1 + j) * ncs + i0];
+ tmp_c[j] = src1_data[(i1 + j) * nc + i0];
+ }
+
+ __asm__ volatile(
+ "flw.ps f10, %[acc]\n"
+ "flw.ps f11, %[sv]\n"
+ "flw.ps f12, %[cv]\n"
+ "fmadd.ps f10, f11, f12, f10\n"
+ "fsw.ps f10, %[out]\n"
+ : [out] "=m"(*(float (*)[8]) acc)
+ : [acc] "m"(*(const float (*)[8]) acc), [sv] "m"(*(const float (*)[8]) tmp_s),
+ [cv] "m"(*(const float (*)[8]) tmp_c)
+ : "f10", "f11", "f12");
+ }
+
+ // Store 8 results — dst is contiguous along d_inner
+ __asm__ volatile(
+ "flw.ps f10, %[acc]\n"
+ "fsw.ps f10, %[dst]\n"
+ : [dst] "=m"(*(float (*)[8])(x + i1))
+ : [acc] "m"(*(const float (*)[8]) acc)
+ : "f10");
+ }
+
+ // Scalar tail for remaining channels
+ for (; i1 < i1_end; ++i1) {
+ const float * c = src1_data + i1 * nc;
+ const float * s_row = s + i1 * ncs;
+ float sumf = 0.0f;
+ for (int64_t i0 = 0; i0 < nc; ++i0) {
+ sumf += s_row[i0] * c[i0];
+ }
+ x[i1] = sumf;
+ }
+ }
+ }
+ }
+
+ // Restore mask
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+
+ return 0;
+}
--- /dev/null
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_ssm_scan_params {
+ struct ggml_tensor src0; // s: [d_state, head_dim, n_head, n_seqs]
+ struct ggml_tensor src1; // x: [head_dim, n_head, n_seq_tokens, n_seqs]
+ struct ggml_tensor src2; // dt: [n_head, n_seq_tokens, n_seqs]
+ struct ggml_tensor src3; // A: [d_state, n_head] or [1, n_head]
+ struct ggml_tensor src4; // B: [d_state, n_group, n_seq_tokens, n_seqs]
+ struct ggml_tensor src5; // C: [d_state, n_group, n_seq_tokens, n_seqs]
+ struct ggml_tensor src6; // ids: [n_seqs] i32
+ struct ggml_tensor dst; // packed [y, final_state]
+};
+
+static inline float softplus_f32(float x) {
+ return x <= 20.0f ? et_logf(1.0f + et_expf(x)) : x;
+}
+
+int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ const int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ const int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * src1 = ¶ms->src1;
+ struct ggml_tensor * src2 = ¶ms->src2;
+ struct ggml_tensor * src3 = ¶ms->src3;
+ struct ggml_tensor * src4 = ¶ms->src4;
+ struct ggml_tensor * src5 = ¶ms->src5;
+ struct ggml_tensor * src6 = ¶ms->src6;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || src2->type != GGML_TYPE_F32 ||
+ src3->type != GGML_TYPE_F32 || src4->type != GGML_TYPE_F32 || src5->type != GGML_TYPE_F32 ||
+ src6->type != GGML_TYPE_I32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ const float * s_data = (const float *) src0->data;
+ const float * x_data = (const float *) src1->data;
+ const float * dt_data = (const float *) src2->data;
+ const float * A_data = (const float *) src3->data;
+ const float * B_data = (const float *) src4->data;
+ const float * C_data = (const float *) src5->data;
+ const int32_t * ids = (const int32_t *) src6->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!s_data || !x_data || !dt_data || !A_data || !B_data || !C_data || !ids || !dst_data) {
+ return -1;
+ }
+
+ const int64_t d_state = src0->ne[0];
+ const int64_t head_dim = src0->ne[1];
+ const int64_t n_head = src1->ne[1];
+ const int64_t n_group = src4->ne[1];
+ const int64_t n_seq_tokens = src1->ne[2];
+ const int64_t n_seqs = src1->ne[3];
+ const int64_t y_elems = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+
+ if (src0->nb[0] != sizeof(float) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(float) ||
+ src3->nb[0] != sizeof(float) || src4->nb[0] != sizeof(float) || src5->nb[0] != sizeof(float) ||
+ src6->nb[0] != sizeof(int32_t) || dst->nb[0] != sizeof(float)) {
+ return -1;
+ }
+
+ if (n_group <= 0 || n_head % n_group != 0) {
+ return -1;
+ }
+
+ // Cache-line bundling on the dst output (1 dst float per (head, dim, token)).
+ // - When head_dim < 16: bundle 16/head_dim heads per work-unit (1 line of dst).
+ // - When head_dim >= 16: each head's dim slice spans head_dim/16 lines, so we
+ // can split dims into chunks of 16 across threads without false sharing.
+ const int64_t dst_lanes_per_cl = 16;
+ const int64_t heads_per_cacheline = head_dim >= dst_lanes_per_cl ? 1 : (dst_lanes_per_cl / head_dim);
+ const int64_t heads_per_block = heads_per_cacheline > 0 ? heads_per_cacheline : 1;
+ const int64_t blocks_per_seq = (n_head + heads_per_block - 1) / heads_per_block;
+ const int64_t dim_chunk_lanes = head_dim >= dst_lanes_per_cl ? dst_lanes_per_cl : head_dim;
+ const int64_t dim_chunks_per_head = (head_dim + dim_chunk_lanes - 1) / dim_chunk_lanes;
+
+ // A "unit" = (seq, head_block, dim_chunk). This expands the parallelism by a
+ // factor of dim_chunks_per_head over the prior block-only scheme; for Mamba-2
+ // shapes (head_dim=64) that's a 4x bump in active threads.
+ const int64_t units_per_seq = blocks_per_seq * dim_chunks_per_head;
+ const int64_t total_units = n_seqs * units_per_seq;
+ const int64_t units_per_thread = (total_units + num_threads - 1) / num_threads;
+ const int64_t unit_begin = (int64_t) thread_id * units_per_thread;
+ int64_t unit_end = unit_begin + units_per_thread;
+
+ if (unit_begin >= total_units) {
+ return 0;
+ }
+
+ if (unit_end > total_units) {
+ unit_end = total_units;
+ }
+
+ const int A_broadcast = (src3->ne[0] == 1);
+ const int64_t d_state_vec = (d_state / 8) * 8; // largest multiple of 8 <= d_state
+ const float log2e_const = 1.4426950408889634f;
+
+ for (int64_t unit = unit_begin; unit < unit_end; ++unit) {
+ const int64_t seq_idx = unit / units_per_seq;
+ const int64_t unit_in_seq = unit % units_per_seq;
+ const int64_t block_in_seq = unit_in_seq / dim_chunks_per_head;
+ const int64_t dim_chunk_idx = unit_in_seq % dim_chunks_per_head;
+ const int64_t head_begin = block_in_seq * heads_per_block;
+ int64_t head_end = head_begin + heads_per_block;
+
+ if (head_end > n_head) {
+ head_end = n_head;
+ }
+
+ const int64_t dim_begin = dim_chunk_idx * dim_chunk_lanes;
+ int64_t dim_end = dim_begin + dim_chunk_lanes;
+ if (dim_end > head_dim) {
+ dim_end = head_dim;
+ }
+
+ const int32_t state_seq = ids[seq_idx];
+
+ for (int64_t head_idx = head_begin; head_idx < head_end; ++head_idx) {
+ const int64_t group_idx = head_idx / (n_head / n_group);
+
+ // A pointer for this head: contiguous over state_idx when not broadcast
+ const float * A_row = (const float *) ((const char *) A_data + (size_t) head_idx * src3->nb[1]);
+
+ for (int64_t dim_idx = dim_begin; dim_idx < dim_end; ++dim_idx) {
+ const float * state_src =
+ (const float *) ((const char *) s_data + (size_t) dim_idx * src0->nb[1] +
+ (size_t) head_idx * src0->nb[2] + (size_t) state_seq * src0->nb[3]);
+
+ float * state_dst =
+ (float *) ((char *) dst_data + (size_t) y_elems * sizeof(float) + (size_t) dim_idx * src0->nb[1] +
+ (size_t) head_idx * src0->nb[2] + (size_t) seq_idx * src0->nb[3]);
+
+ for (int64_t token_idx = 0; token_idx < n_seq_tokens; ++token_idx) {
+ const float * x_ptr =
+ (const float *) ((const char *) x_data + (size_t) dim_idx * src1->nb[0] +
+ (size_t) head_idx * src1->nb[1] + (size_t) token_idx * src1->nb[2] +
+ (size_t) seq_idx * src1->nb[3]);
+
+ const float * dt_ptr =
+ (const float *) ((const char *) dt_data + (size_t) head_idx * src2->nb[0] +
+ (size_t) token_idx * src2->nb[1] + (size_t) seq_idx * src2->nb[2]);
+
+ const float * B_row =
+ (const float *) ((const char *) B_data + (size_t) group_idx * src4->nb[1] +
+ (size_t) token_idx * src4->nb[2] + (size_t) seq_idx * src4->nb[3]);
+
+ const float * C_row =
+ (const float *) ((const char *) C_data + (size_t) group_idx * src5->nb[1] +
+ (size_t) token_idx * src5->nb[2] + (size_t) seq_idx * src5->nb[3]);
+
+ const float dt_softplus = softplus_f32(*dt_ptr);
+ const float x_dt = (*x_ptr) * dt_softplus;
+ const float dt_log2e = dt_softplus * log2e_const;
+
+ // Source of "previous state" for this token: input state on token 0,
+ // last token's state thereafter (we wrote it into state_dst).
+ const float * prev_row = (token_idx == 0) ? state_src : state_dst;
+
+ float sumf = 0.0f;
+ int64_t state_idx = 0;
+
+ if (d_state_vec > 0) {
+ // Save mask, enable all 8 vector lanes for the state loop.
+ unsigned long saved_mask;
+ __asm__ volatile("mova.x.m %0" : "=r"(saved_mask));
+ __asm__ volatile("mov.m.x m0, x0, 0xFF");
+
+ // Per-token broadcasts:
+ // f20 = x_dt (B*x_dt)
+ // f21 = dt_log2e (for fexp.ps when A is per-state)
+ // f22 = dA (only when A is broadcast scalar)
+ // f23 = sum-of-products accumulator (zeroed)
+ __asm__ volatile(
+ "fbc.ps f20, %[xdt]\n\t"
+ "fbc.ps f21, %[dtl]\n\t"
+ "fbci.pi f23, 0\n\t"
+ :
+ : [xdt] "m"(x_dt), [dtl] "m"(dt_log2e)
+ : "f20", "f21", "f23");
+
+ if (A_broadcast) {
+ // dA is a per-head scalar — compute once and splat.
+ const float dA_scalar = et_expf(dt_softplus * (*A_row));
+ __asm__ volatile("fbc.ps f22, %[da]\n\t" : : [da] "m"(dA_scalar) : "f22");
+ }
+
+ for (; state_idx < d_state_vec; state_idx += 8) {
+ if (!A_broadcast) {
+ // f22 = exp(dt_softplus * A[state..state+7])
+ // = 2^((dt_softplus * A) * log2e) via fexp.ps
+ __asm__ volatile(
+ "flw.ps f24, %[av]\n\t"
+ "fmul.ps f24, f24, f21\n\t" // A * dt_log2e
+ "fexp.ps f22, f24\n\t" // dA = 2^(...)
+ :
+ : [av] "m"(*(const float (*)[8]) & A_row[state_idx])
+ : "f22", "f24");
+ }
+
+ // state = prev * dA + B * x_dt
+ // sumf += state * C
+ // Reads prev before writing state_dst — safe even when
+ // prev_row == state_dst (write-after-read, same index).
+ __asm__ volatile(
+ "flw.ps f25, %[prev]\n\t"
+ "flw.ps f26, %[bv]\n\t"
+ "flw.ps f27, %[cv]\n\t"
+ "fmul.ps f26, f26, f20\n\t" // B * x_dt
+ "fmadd.ps f25, f25, f22, f26\n\t" // state = prev*dA + B*x_dt
+ "fsw.ps f25, %[sd]\n\t"
+ "fmadd.ps f23, f25, f27, f23\n\t" // sum += state*C
+ : [sd] "=m"(*(float (*)[8]) & state_dst[state_idx])
+ : [prev] "m"(*(const float (*)[8]) & prev_row[state_idx]),
+ [bv] "m"(*(const float (*)[8]) & B_row[state_idx]),
+ [cv] "m"(*(const float (*)[8]) & C_row[state_idx])
+ : "f25", "f26", "f27");
+ }
+
+ // Horizontal reduce f23 (8 lanes) -> scalar sumf.
+ __asm__ volatile(
+ "fswizz.ps f1, f23, 0xB1\n\t"
+ "fadd.ps f2, f23, f1, rne\n\t"
+ "fswizz.ps f3, f2, 0x4E\n\t"
+ "fadd.ps f4, f2, f3, rne\n\t"
+ "fmvz.x.ps t0, f4, 4\n\t"
+ "fbcx.ps f5, t0\n\t"
+ "fadd.ps %[vout], f4, f5, rne\n\t"
+ : [vout] "=f"(sumf)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ __asm__ volatile("mova.m.x %0" ::"r"(saved_mask));
+ }
+
+ // Scalar tail (d_state not a multiple of 8).
+ for (; state_idx < d_state; ++state_idx) {
+ const float prev_state = prev_row[state_idx];
+ const float A_val = A_broadcast ? *A_row : A_row[state_idx];
+ const float dA = et_expf(dt_softplus * A_val);
+ const float st = prev_state * dA + B_row[state_idx] * x_dt;
+ state_dst[state_idx] = st;
+ sumf += st * C_row[state_idx];
+ }
+
+ dst_data[seq_idx * (n_seq_tokens * n_head * head_dim) + token_idx * (n_head * head_dim) +
+ head_idx * head_dim + dim_idx] = sumf;
+ }
+ }
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// SUM_ROWS F32 Kernel
+// Row-wise sum reduction: dst[0, i1, i2, i3] = sum(src0[0..ne00-1, i1, i2, i3])
+// Vectorized 8-wide accumulation with horizontal reduction.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_sum_rows_params {
+ struct ggml_tensor src0; // F32 input tensor [ne00, ne01, ne02, ne03]
+ struct ggml_tensor dst; // F32 output tensor [1, ne01, ne02, ne03]
+};
+
+int entry_point(struct ggml_et_sum_rows_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+
+ const int64_t ne00 = src0->ne[0]; // Row length (to be summed)
+ const int64_t ne01 = src0->ne[1];
+ const int64_t ne02 = src0->ne[2];
+ const int64_t ne03 = src0->ne[3];
+
+ const size_t nb01 = src0->nb[1];
+ const size_t nb02 = src0->nb[2];
+ const size_t nb03 = src0->nb[3];
+
+ const size_t nb1 = dst->nb[1];
+ const size_t nb2 = dst->nb[2];
+ const size_t nb3 = dst->nb[3];
+
+ // Flatten rows across dimensions 1,2,3 and distribute across threads
+ const int64_t total_rows = ne01 * ne02 * ne03;
+
+ for (int64_t ir = thread_id; ir < total_rows; ir += num_threads) {
+ const int64_t i03 = ir / (ne02 * ne01);
+ const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01;
+ const int64_t i01 = ir - i03 * ne02 * ne01 - i02 * ne01;
+
+ const float * src_row = (const float *) ((const char *) src0_data + i01 * nb01 + i02 * nb02 + i03 * nb03);
+ float * dst_ptr = (float *) ((char *) dst_data + i01 * nb1 + i02 * nb2 + i03 * nb3);
+
+ // Vectorized 8-wide sum accumulation
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int32_t i0 = 0; i0 < (int32_t) ne00; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[x_vec]\n"
+ "fadd.ps f10, f10, f11\n"
+ :
+ : [x_vec] "m"(*(const float (*)[8]) & src_row[i0])
+ : "f10", "f11");
+ }
+
+ // Horizontal sum of 8 accumulated values in f10
+ float row_sum;
+ __asm__ __volatile__(
+ "fswizz.ps f1, f10, 0xB1 \n\t"
+ "fadd.ps f2, f10, f1, rne \n\t"
+ "fswizz.ps f3, f2, 0x4E \n\t"
+ "fadd.ps f4, f2, f3, rne \n\t"
+ "fmvz.x.ps t0, f4, 4 \n\t"
+ "fbcx.ps f5, t0 \n\t"
+ "fadd.ps %[vout], f4, f5, rne \n\t"
+ : [vout] "=f"(row_sum)::"t0", "f1", "f2", "f3", "f4", "f5");
+
+ atomic_store_f32(dst_ptr, row_sum);
+ }
+
+ return 0;
+}
--- /dev/null
+#ifndef __TENSORS_H
+#define __TENSORS_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#if defined(__cplusplus) && (__cplusplus >= 201103L)
+# include <cinttypes>
+# if (__cplusplus < 202002L)
+# include <cstdbool>
+# endif
+#else
+# include <inttypes.h>
+# include <stdbool.h>
+#endif
+
+/*! \def QUANT_LAST_TRANS
+ \brief Tensor Quant instruction: Do not perform any more transformations.
+*/
+#define QUANT_LAST_TRANS 0
+
+/*! \def QUANT_INT32_TO_FP32
+ \brief Tensor Quant instruction: Convert all elements of A from 32-bit signed integer values to single-precision
+ floating-point values.
+*/
+#define QUANT_INT32_TO_FP32 1
+
+/*! \def QUANT_FP32_TO_INT32
+ \brief Tensor Quant instruction: Convert all elements of A from single-precision floating-point values to 32-
+ bit signed integer values.
+*/
+#define QUANT_FP32_TO_INT32 2
+
+/*! \def QUANT_RELU
+ \brief Tensor Quant instruction: Convert all negative INT32 values in A to 0
+*/
+#define QUANT_RELU 3
+
+/*! \def QUANT_INT32_ADD_ROW
+ \brief Tensor Quant instruction: Read the low-order COLS+1 32-bit signed integer values from an L1
+ scratchpad line, and add this vector to every row of the 32-bit signed integer
+ matrix A.
+*/
+#define QUANT_INT32_ADD_ROW 4
+
+/*! \def QUANT_INT32_ADD_COL
+ \brief Tensor Quant instruction: Read the low-order ROWS+1 32-bit signed integer values from an L1
+ scratchpad line, and add this vector to every column of the 32-bit signed
+ integer matrix A.
+*/
+#define QUANT_INT32_ADD_COL 5
+
+/*! \def QUANT_FP32_MUL_ROW
+ \brief Tensor Quant instruction: Read the low-order COLS+1 single-precision floating-point values from an
+ L1 scratchpad line, and multiply the single-precision elements of each row
+ of matrix A element-wise by this vector.
+*/
+#define QUANT_FP32_MUL_ROW 6
+
+/*! \def QUANT_FP32_MUL_COL
+ \brief Tensor Quant instruction: Read the low-order ROWS+1 single-precision floating-point values from an
+ L1 scratchpad line, and multiply the single-precision elements of each col-
+ umn of matrix A element-wise by this vector.
+*/
+#define QUANT_FP32_MUL_COL 7
+
+/*! \def QUANT_SATINT8
+ \brief Tensor Quant instruction: Clamp all 32-bit signed integer values in A to the range [-128, 127].
+ The values are written in bits 7:0 of each element, with bits 31:8 set to zero.
+*/
+#define QUANT_SATINT8 8
+
+/*! \def QUANT_SATUINT8
+ \brief Tensor Quant instruction: Clamp all 32-bit signed integer values in A to the range [0, 255]. The values
+ are written in bits 7:0 of each element, with bits 31:8 set to zero.
+*/
+#define QUANT_SATUINT8 9
+
+/*! \def QUANT_PACK_128B
+ \brief Tensor Quant instruction: Copy the low-order byte of the n-th 32-bit value in each row of A to the n-th
+ byte of the row.
+*/
+#define QUANT_PACK_128B 10
+
+/*! \def TENSOR_REDUCE_OP_FADD
+ \brief Tensor Reduce instruction: The result is the addition of the incoming single-precision floating-point data
+ and the single-precision floating-point values in the vector register file.
+*/
+#define TENSOR_REDUCE_OP_FADD 0
+
+// #define TENSOR_REDUCE_OP_FSUB 1 -- Not supported
+
+/*! \def TENSOR_REDUCE_OP_FMAX
+ \brief Tensor Reduce instruction: The result is the maximum of the incoming single-precision floating-point data
+and the single-precision floating-point values in the vector register file.
+*/
+#define TENSOR_REDUCE_OP_FMAX 2
+
+/*! \def TENSOR_REDUCE_OP_FMIN
+ \brief Tensor Reduce instruction: The result is the minimum of the incoming single-precision floating-point data
+and the single-precision floating-point values in the vector register file..
+*/
+#define TENSOR_REDUCE_OP_FMIN 3
+
+/*! \def TENSOR_REDUCE_OP_IADD
+ \brief Tensor Reduce instruction: The result is the addition of the incoming 32-bit integer data and the 32-bit inte-
+ger values in the vector register file.
+*/
+#define TENSOR_REDUCE_OP_IADD 4
+
+// #define TENSOR_REDUCE_OP_ISUB 5 -- Not supported
+
+/*! \def TENSOR_REDUCE_OP_IMAX
+ \brief Tensor Reduce instruction: The result is the maximum of the incoming 32-bit signed integer data and the
+32-bit signed integer values in the vector register file.
+*/
+#define TENSOR_REDUCE_OP_IMAX 6
+
+/*! \def TENSOR_REDUCE_OP_IMIN
+ \brief Tensor Reduce instruction: The result is the minimum of the incoming 32-bit signed integer data and the
+32-bit signed integer values in the vector register file.
+*/
+#define TENSOR_REDUCE_OP_IMIN 7
+
+/*! \def TENSOR_REDUCE_OP_FGET
+ \brief Tensor Reduce instruction get function to be performed
+*/
+#define TENSOR_REDUCE_OP_FGET 8
+
+/*! \def TENSOR_LOAD_WAIT_0
+ \brief Tensor load to L1 Scratchpad with ID = 0 is complete.
+*/
+#define TENSOR_LOAD_WAIT_0 0
+
+/*! \def TENSOR_LOAD_WAIT_1
+ \brief Tensor load to L1 Scratchpad with ID = 1 is complete.
+*/
+#define TENSOR_LOAD_WAIT_1 1
+
+/*! \def TENSOR_FMA_WAIT
+ \brief All previous tensor matrix multiplication instructions are complete.
+*/
+#define TENSOR_FMA_WAIT 7
+
+/*! \def TENSOR_STORE_WAIT
+ \brief All previous tensor store instructions are complete.
+*/
+#define TENSOR_STORE_WAIT 8
+
+/*! \def TENSOR_REDUCE_WAIT
+ \brief All previous tensor reduction instructions are complete
+*/
+#define TENSOR_REDUCE_WAIT 9
+
+/*! \def TENSOR_QUANT_WAIT
+ \brief TensorQuant is complete
+*/
+#define TENSOR_QUANT_WAIT 10
+
+// TensorFMA opcode values (tensor_fma CSR 0x801, bits 3:1)
+#define TENSOR_FMA_OP_FP32 0 // TensorFMA32: FP32 x FP32 -> FP32
+#define TENSOR_FMA_OP_FP16 1 // TensorFMA16A32: FP16 x FP16 -> FP32
+// opcode 2 is reserved
+#define TENSOR_FMA_OP_INT8 3 // TensorIMA8A32: INT8 x INT8 -> INT32
+
+// TensorLoad transformation values (tensor_load CSR 0x83F, bits 61:59)
+#define TENSOR_LOAD_PLAIN 0 // TensorLoad: 64B rows
+#define TENSOR_LOAD_INTERLEAVE8 1 // TensorLoadInterleave8: for TensorIMA8A32 B
+#define TENSOR_LOAD_INTERLEAVE16 2 // TensorLoadInterleave16: for TensorFMA16A32 B
+// transformations 3-4 are reserved
+#define TENSOR_LOAD_TRANSPOSE8 5 // TensorLoadTranspose8: 8-bit transpose
+#define TENSOR_LOAD_TRANSPOSE16 6 // TensorLoadTranspose16: 16-bit transpose
+#define TENSOR_LOAD_TRANSPOSE32 7 // TensorLoadTranspose32: 32-bit transpose
+
+/*! \def TENSOR_ERROR_LOAD_TRANSFORM
+ \brief Define for tensor load transform error.
+*/
+#define TENSOR_ERROR_LOAD_TRANSFORM 1
+
+/*! \def TENSOR_ERROR_FCC_OVERFLOW
+ \brief Define for tensor fcc overflow error.
+*/
+#define TENSOR_ERROR_FCC_OVERFLOW 3
+
+/*! \def TENSOR_ERROR_SCP_DISABLED
+ \brief Define for tensor scp disabled error.
+*/
+#define TENSOR_ERROR_SCP_DISABLED 4
+
+/*! \def TENSOR_ERROR_LOCKSW
+ \brief Define for tensor locksw error.
+*/
+#define TENSOR_ERROR_LOCKSW 5
+
+/*! \def TENSOR_ERROR_TL1_FMA
+ \brief Define for L1 FMA error.
+*/
+#define TENSOR_ERROR_TL1_FMA 6
+
+/*! \def TENSOR_ERROR_MEM_FAULT
+ \brief Define for Memory fault error.
+*/
+#define TENSOR_ERROR_MEM_FAULT 7
+
+/*! \def TENSOR_ERROR_STORE_COOP
+ \brief Define for store coop error.
+*/
+#define TENSOR_ERROR_STORE_COOP 8
+
+/*! \def TENSOR_ERROR_REDUCE
+ \brief Define for tensor reduce error.
+*/
+#define TENSOR_ERROR_REDUCE 9
+
+/*! \struct et_tensor_load_l2scp_conf
+ \brief Tensor load from scp instruction configuration structure.
+*/
+typedef struct et_tensor_load_l2scp_conf {
+ bool use_tmask;
+ uint64_t dst_start;
+ uint64_t addr;
+ uint64_t num_lines;
+ uint64_t stride;
+ uint64_t id;
+} et_tensor_load_l2scp_conf_t;
+
+/*! \enum reduce_transform_t
+ \brief enum transform mode for tensor reduce.
+*/
+typedef enum {
+ FADD = 0x0ULL,
+ FSUB = 0x1ULL,
+ FMAX = 0x2ULL,
+ FMIN = 0x3ULL,
+ IADD = 0x4ULL,
+ ISUB = 0x5ULL,
+ IMAX = 0x6ULL,
+ IMIN = 0x7ULL,
+ FGET = 0x8ULL
+} reduce_transform_t;
+
+/*! \struct et_tensor_load_conf
+ \brief Tensor load instruction configuration structure.
+*/
+typedef struct et_tensor_load_conf {
+ bool use_tmask;
+ bool use_coop;
+ bool use_tenb;
+ uint64_t dst_start;
+ uint64_t transformation;
+ uint64_t rd_l2scp;
+ uint64_t addr;
+ uint64_t offset;
+ uint64_t num_lines;
+ uint64_t stride;
+ uint64_t id;
+} et_tensor_load_conf_t;
+
+/*! \fn inline void tensor_wait(long id)
+ \brief Tensor wait instruction, Tensor Wait can be used to stall execution until
+ a previously issued tensor instruction completes.
+ \param id tensor ID
+ \return none
+ \tensorops Implementation of tensor_wait api
+*/
+inline __attribute__((always_inline)) void tensor_wait(long id) {
+ __asm__ __volatile__(" csrw 0x830, %[id]\n" : : [id] "r"(id) : "memory");
+}
+
+/*! \fn inline void tensor_load (tensor_load *conf)
+ \brief Tensor load instruction, it loads data from memory (bypass-ing the L1 cache)
+ into the L1 scratchpad. Input parameter defines the configuration to tensor load.
+ \param use_tmask the tensor_mask register is used for this operation
+ \param use_coop the operation is a cooperative tensor load.
+ \param dst_start L1 Scratchpad starting cache line
+ \param transformation These bits, along with bit 52, decodes the type of tensor operation.
+ \param use_tenb This bit, along with transformation, decodes the type of tensor operation.
+ \param addr tensor load address
+ \param offset tensor load address offset
+ \param num_lines tensor load number of cache lines
+ \param stride tensor load stride value
+ \param id tensor load id
+ \return none
+ \tensorops Implementation of tensor_load api
+
+*/
+// 1. Load Matrix A segment (1 row x 16 cols) into SCP ID 0
+// dst_start 0 refers to the first line of L1 Scratchpad
+// tensor_load(false, false, 0, 0, 0,
+// (uint64_t)(src0_data + m * K + kb), 0, 1, 0, 0);
+
+inline void __attribute__((always_inline)) tensor_load(bool use_tmask,
+ bool use_coop,
+ uint64_t dst_start,
+ uint64_t transformation,
+ uint64_t use_tenb,
+ uint64_t addr,
+ uint64_t offset,
+ uint64_t num_lines,
+ uint64_t stride,
+ uint64_t id) {
+ // Address alignment depends on transformation type:
+ // Interleave8, Transpose8 (1,5): 16B aligned, addr bits 47:4
+ // Interleave16, Transpose16 (2,6): 32B aligned, addr bits 47:5
+ // Load, Transpose32, LoadB (0,7): 64B aligned, addr bits 47:6
+ uint64_t addr_mask = (transformation == 1 || transformation == 5) ? 0xFFFFFFFFFFF0ULL :
+ (transformation == 2 || transformation == 6) ? 0xFFFFFFFFFFE0ULL :
+ 0xFFFFFFFFFFC0ULL;
+ uint64_t csr_enc = (((uint64_t) use_tmask & 1) << 63) | (((uint64_t) use_coop & 1) << 62) |
+ ((transformation & 0x7) << 59) | ((dst_start & 0x3F) << 53) | ((use_tenb & 0x1) << 52) |
+ ((addr & addr_mask)) | ((offset & 0x3) << 4) | ((num_lines & 0xF));
+
+ uint64_t x31_enc = (stride & 0xFFFFFFFFFFC0ULL) | (id & 0x1);
+
+ __asm__ __volatile__(
+ "mv x31, %[x31v]\n"
+ "csrw 0x83f, %[csrv]\n"
+ :
+ : [x31v] "r"(x31_enc), [csrv] "r"(csr_enc)
+ : "x31", "memory");
+}
+
+/*! \fn inline void et_tensor_load (et_tensor_load_conf_t *conf)
+ \brief Tensor load instruction, it loads data from memory (bypass-ing the L1 cache)
+ into the L1 scratchpad. Input parameter defines the configuration to tensor load.
+ \param conf tensor load configuration
+ \return none
+ \tensorops Implementation of et_tensor_load api
+*/
+inline void __attribute__((always_inline)) et_tensor_load(et_tensor_load_conf_t * conf) {
+ tensor_load(conf->use_tmask, conf->use_coop, conf->dst_start, conf->transformation, (uint64_t) conf->use_tenb,
+ conf->addr, conf->offset, conf->num_lines, conf->stride, conf->id);
+}
+
+/*! \fn inline void tensor_load_setup_b(bool use_coop, uint64_t addr, uint64_t num_lines, uint64_t stride, uint64_t id)
+ \brief Tensor load instruction setup
+ \param use_coop the operation is a cooperative tensor load.
+ \param addr tensor load address
+ \param num_lines tensor load number of cache lines
+ \param stride tensor load stride value
+ \param id tensor load id
+ \return none
+ \tensorops Implementation of tensor_load_setup_b api
+*/
+inline void __attribute__((always_inline)) tensor_load_setup_b(bool use_coop,
+ uint64_t addr,
+ uint64_t num_lines,
+ uint64_t stride,
+ uint64_t id) {
+ uint64_t csr_enc =
+ (((uint64_t) use_coop & 1) << 62) | (0x1ULL << 52) | ((addr & 0xFFFFFFFFFFC0ULL)) | ((num_lines & 0xF));
+ uint64_t x31_enc = (stride & 0xFFFFFFFFFFC0ULL) | (id & 0x1);
+
+ __asm__ __volatile__(
+ "mv x31, %[x31v]\n"
+ "csrw 0x83f, %[csrv]\n"
+ :
+ : [x31v] "r"(x31_enc), [csrv] "r"(csr_enc)
+ : "x31", "memory");
+}
+
+/*! \fn inline void et_tensor_load_l2scp (et_tensor_load_l2scp_conf_t *conf)
+ \brief Tensor load l2scp loads data from memory (bypassing the L1 and L2 caches) into the L2 scratchpad.
+ \param conf tensor load configuration
+ \return none
+ \tensorops Implementation of et_tensor_load_l2scp api
+*/
+inline void __attribute__((always_inline)) et_tensor_load_l2scp(et_tensor_load_l2scp_conf_t * conf) {
+ uint64_t csr_enc =
+ (((((uint64_t) conf->use_tmask) & 1) << 63) | ((conf->dst_start & 0x1FFFCUL) << (48 - 2)) |
+ ((conf->dst_start & 0x3UL) << 4) | ((conf->addr & 0xFFFFFFFFFFC0UL)) | ((conf->num_lines & 0x0FUL)));
+ uint64_t x31_enc = (conf->stride & 0xFFFFFFFFFFC0ULL) | (conf->id & 0x1);
+
+ __asm__ __volatile__(
+ "mv x31, %[x31v]\n"
+ "csrw 0x85f, %[csrv]\n"
+ :
+ : [x31v] "r"(x31_enc), [csrv] "r"(csr_enc)
+ : "x31", "memory");
+}
+
+/*! \fn inline void tensor_store_scp(uint64_t entry_stride,
+ uint64_t start_scp_entry,
+ uint64_t Arows,
+ uint64_t addr,
+ uint64_t stride)
+ \brief Tensor Store writes a series of 64-byte blocks of data from the L1 scratchpad into memory.
+ A matrix X can have up to 16 rows, and each row can be up to 64B in size (the number of columns depends on the type of elements of X).
+ \param entry_stride Register stride
+ \param start_scp_entry Start register
+ \param Arows A matrix row size
+ \param addr Virtual Address
+ \param stride This value is the distance in bytes between consecutive tensor rows in memory
+ \return none
+ \tensorops Implementation of tensor_store_scp api
+*/
+inline void __attribute__((always_inline)) tensor_store_scp(uint64_t entry_stride,
+ uint64_t start_scp_entry,
+ uint64_t Arows,
+ uint64_t addr,
+ uint64_t stride) {
+ uint64_t csr_enc = ((entry_stride & 0x3) << 62) | ((start_scp_entry & 0x3F) << 56) | ((addr & 0xFFFFFFFFFFC0ULL)) |
+ ((Arows & 0xF) << 51) | (((uint64_t) 1) << 48);
+ uint64_t x31_enc = (stride & 0xFFFFFFFFFFC0UL);
+
+ __asm__ __volatile__(
+ "mv x31, %[x31v]\n"
+ "csrw 0x87f, %[csrv]\n"
+ :
+ : [x31v] "r"(x31_enc), [csrv] "r"(csr_enc)
+ : "x31", "memory");
+}
+
+/*! \fn inline void tensor_store(uint64_t reg_stride,
+ uint64_t start_reg,
+ uint64_t cols,
+ uint64_t Arows,
+ uint64_t addr,
+ uint64_t coop_store,
+ uint64_t stride)
+ \brief The Tensor store instruction reads a tensor from the vector register files and writes it to memory,
+ bypassing the L1 data cache and the L2 cache. For the purposes of this instruction the tensor has ROWS+1 rows,
+ and each row is 16*SIZE+16 bytes in size.
+ \param reg_stride Register stride
+ \param start_reg start register address
+ \param cols matrix row size.
+ \param Arows matrix row size
+ \param addr Virtual Address
+ \param coop_store Number of minions to cooperate with
+ \param stride This value is the distance in bytes between consecutive tensor rows in memory
+ \return none
+ \tensorops Implementation of tensor_store api
+*/
+inline void __attribute__((always_inline)) tensor_store(uint64_t reg_stride,
+ uint64_t start_reg,
+ uint64_t cols,
+ uint64_t Arows,
+ uint64_t addr,
+ uint64_t coop_store,
+ uint64_t stride) {
+ uint64_t warl = 0;
+ uint64_t csr_enc = ((reg_stride & 0x3) << 62) | ((start_reg & 0x1F) << 57) | ((cols & 0x3) << 55) |
+ ((addr & 0xFFFFFFFFFFF0)) | ((Arows & 0xF) << 51) | ((coop_store & 0x3) << 49) | ((warl & 0xF));
+
+ uint64_t x31_enc = (stride & 0xFFFFFFFFFF0UL);
+
+ __asm__ __volatile__(
+ "mv x31, %[x31v]\n"
+ "csrw 0x87f, %[csrv]\n"
+ :
+ : [x31v] "r"(x31_enc), [csrv] "r"(csr_enc)
+ : "x31", "memory");
+}
+
+/*! \fn inline void tensor_fma(bool use_tmask,
+ uint64_t b_num_col,
+ uint64_t a_num_rows,
+ uint64_t a_num_cols,
+ uint64_t offset,
+ bool tenc_loc,
+ bool tenb_unsigned,
+ bool tena_unsigned,
+ bool tenb_loc,
+ uint64_t scp_loc_b,
+ uint64_t scp_loc_a,
+ uint64_t opcode,
+ bool first_pass)
+ \brief The Tensor FMA instruction multiplies two matrices A and B, optionally adds the resulting matrix
+ to a third matrix C, and writes the result back onto matrix C
+ \param use_tmask Use tensor_mask CSR to skip operations in an A row granularity.
+ \param b_num_col B matrix number of columns
+ \param a_num_rows A matrix number of rows
+ \param a_num_cols A matrix number of columns
+ \param offset A matrix starting column for the operation.
+ \param tenc_loc Location of matrix C (0 = L1 scratchpad, 1 = memory).
+ \param tenb_unsigned TenB is signed (0) or unsigned (1).
+ \param tena_unsigned TenA is signed (0) or unsigned (1).
+ \param tenb_loc Location of matrix B (0 = L1 scratchpad, 1 = memory).
+ \param scp_loc_b Starting L1 scratchpad cache line where matrix B is stored, ignored when xs[20] = 1.
+ \param scp_loc_a Starting L1 scratchpad cache line where matrix A is stored, ignored when xs[20] = 1.
+ \param opcode 0 = TensorFMA32 (F32xF32->F32), 1 = TensorFMA16A32 (F16xF16->F32), 3 = TensorIMA8A32 (I8xF8->I32).
+ Other opcodes are invalid.
+ \param first_pass if set to 0 then the initial value of TenC is added to the result
+ \return none
+ \tensorops Implementation of tensor_fma api
+*/
+inline void __attribute__((always_inline)) tensor_fma(bool use_tmask,
+ uint64_t b_num_col,
+ uint64_t a_num_rows,
+ uint64_t a_num_cols,
+ uint64_t offset,
+ bool tenc_loc,
+ bool tenb_unsigned,
+ bool tena_unsigned,
+ bool tenb_loc,
+ uint64_t scp_loc_b,
+ uint64_t scp_loc_a,
+ uint64_t opcode,
+ bool first_pass) {
+ uint64_t csr_enc = (((uint64_t) use_tmask & 1) << 63) | ((b_num_col & 0x3) << 55) | ((a_num_rows & 0xF) << 51) |
+ ((a_num_cols & 0xF) << 47) | ((offset & 0xF) << 43) | (((uint64_t) tenc_loc & 1) << 23) |
+ (((uint64_t) tena_unsigned & 1) << 22) | (((uint64_t) tenb_unsigned & 1) << 21) |
+ (((uint64_t) tenb_loc & 1) << 20) | ((scp_loc_b & 0xFF) << 12) | ((scp_loc_a & 0xFF) << 4) |
+ ((opcode & 0x7) << 1) | ((uint64_t) first_pass & 1);
+
+ __asm__ __volatile__("csrw 0x801, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) :);
+}
+
+/*! \fn inline uint32_t tensor_reduce_uint32(uint32_t value, uint64_t operation, uint64_t partnerID, uint64_t action)
+ \brief Tensor reduce allows a group of harts to communicate values held in floating-point registers to collectively calculate a reduction
+ function.
+ \param value Register stride
+ \param operation Function to be performed.
+ \param partnerID Receiver minionID.
+ \param action action value
+ \return uint32_t value after reduction
+ \tensorops Implementation of tensor_reduce_uint32 api
+*/
+inline uint32_t __attribute__((always_inline)) tensor_reduce_uint32(uint32_t value,
+ uint64_t operation,
+ uint64_t partnerID,
+ uint64_t action) {
+ uint64_t warl = 0;
+ uint32_t out;
+ uint64_t csr_enc = ((warl & 0x2) << 62) | ((0ULL & 0x1F) << 57) | ((warl & 0x1FFFFFFF) << 28) |
+ ((operation & 0xF) << 24) | ((1ULL & 0xFF) << 16) | ((partnerID & 0x1FFF) << 3) |
+ ((warl & 0x1) << 2) | ((action & 0x3));
+
+ __asm__ __volatile__(
+ "fmv.s.x f0, %[value]\n"
+ "csrw 0x800, %[csr_enc]\n"
+ "fmv.x.s %[out], f0\n"
+ : [out] "=r"(out)
+ : [csr_enc] "r"(csr_enc), [value] "r"(value)
+ : "f0");
+
+ return out;
+}
+
+/*! \fn inline float tensor_reduce_float(float freg, uint64_t operation, uint64_t num_reg, uint64_t partnerID, uint64_t action) {
+ \brief TensorReduce allows a group of harts to communicate values held in floating-point registers to collectively calculate a reduction
+ function.
+ \param freg Freg register stride
+ \param operation Function to be performed.
+ \param num_reg number of registers to use
+ \param partnerID Receiver minionID.
+ \param action action value
+ \return float value after reduction
+ \tensorops Implementation of tensor_reduce_float api
+*/
+inline float __attribute__((always_inline)) tensor_reduce_float(float freg,
+ uint64_t operation,
+ uint64_t num_reg,
+ uint64_t partnerID,
+ uint64_t action) {
+ uint64_t warl = 0;
+ float out;
+ uint64_t csr_enc = ((warl & 0x2) << 62) | ((0ULL & 0x1F) << 57) | ((warl & 0x1FFFFFFF) << 28) |
+ ((operation & 0xF) << 24) | ((num_reg & 0xFF) << 16) | ((partnerID & 0x1FFF) << 3) |
+ ((warl & 0x1) << 2) | ((action & 0x3));
+
+ __asm__ __volatile__(
+ "fmv.s f0, %[freg]\n"
+ "csrw 0x800, %[csr_enc]\n"
+ "fmv.s %[out], f0\n"
+ : [out] "=f"(out)
+ : [csr_enc] "r"(csr_enc), [freg] "f"(freg)
+ : "f0");
+
+ return out;
+}
+
+//#define tensor_reduce_float1(fval, operation, partnerID, action) do {
+// uint64_t warl = 0;
+// float out;
+// uint64_t csr_enc = ((warl & 0x2 ) << 62) |
+// ((0 & 0x1F ) << 57) |
+// ((warl & 0x1FFFFFFF ) << 28) |
+// ((operation & 0xF ) << 24) |
+// ((1 & 0xFF ) << 16) |
+// ((partnerID & 0x1FFF ) << 3 ) |
+// ((warl & 0x1 ) << 2 ) |
+// ((action & 0x3 ) );
+//
+// register float asm("f0") fval;
+// __asm__ volatile (
+// "csrw 0x800, %[csr_enc]"
+// : "+r" (ftmp)
+// : [csr_enc] "r" (csr_enc)
+// );
+//} while (0)
+//
+//
+//inline float __attribute__((always_inline)) tensor_reduce_float(uint64_t fstart, uint64_t operation, uint64_t num_reg, uint64_t partnerID, uint64_t action) {
+// uint64_t warl = 0;
+// float out;
+// uint64_t csr_enc = ((warl & 0x2 ) << 62) |
+// ((fstart & 0x1F ) << 57) |
+// ((warl & 0x1FFFFFFF ) << 28) |
+// ((operation & 0xF ) << 24) |
+// ((num_reg & 0xFF ) << 16) |
+// ((partnerID & 0x1FFF ) << 3 ) |
+// ((warl & 0x1 ) << 2 ) |
+// ((action & 0x3 ) );
+//
+// __asm__ volatile (
+// "csrw 0x800, %[csr_enc]\n"
+// : /*empty*/
+// : [csr_enc] "r" (csr_enc),
+// : /*"f0", "f1", "f2", "f3", "f4",
+// "f5", "f6", "f7", "f8", "f9",
+// "f10", "f11", "f12", "f13", "f14",
+// "f15", "f16", "f17", "f18", "f19",
+// "f20", "f21", "f22", "f23", "f24",
+// "f25", "f26", "f27", "f28", "f29",
+// "f30", "f31"*/
+// );
+//
+// return out;
+//}
+
+/*! \fn inline void tensor_reduce(uint64_t start_reg, uint64_t operation, uint64_t num_reg, uint64_t partnerID, uint64_t action)
+ \brief The TensorReduce instruction allows up to 216 harts to collectively calculate a reduction function.
+ \param start_reg starting register
+ \param operation Function to be performed.
+ \param num_reg number of registers
+ \param partnerID Receiver minionID.
+ \param action action value
+ \return uint32_t value after reduction
+ \tensorops Implementation of tensor_reduce api
+*/
+
+inline void __attribute__((always_inline)) tensor_reduce(uint64_t start_reg,
+ uint64_t operation,
+ uint64_t num_reg,
+ uint64_t partnerID,
+ uint64_t action) {
+ uint64_t warl = 0;
+
+ uint64_t csr_enc = ((warl & 0x2) << 62) | ((start_reg & 0x1F) << 57) | ((warl & 0x1FFFFFFF) << 28) |
+ ((operation & 0xF) << 24) | ((num_reg & 0xFF) << 16) | ((partnerID & 0x1FFF) << 3) |
+ ((warl & 0x1) << 2) | ((action & 0x3));
+
+ __asm__ __volatile__("csrw 0x800, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) :);
+}
+
+/*! \fn inline void tensor_reduce_send(uint64_t start_reg, uint64_t num_reg, uint64_t partnerID)
+ \brief This function applies reduce instruction to function and then sends to partner minion.
+ \param start_reg starting register
+ \param num_reg number of registers
+ \param partnerID Receiver minionID.
+ \return none
+ \tensorops Implementation of tensor_reduce_send api
+*/
+inline void __attribute__((always_inline)) tensor_reduce_send(uint64_t start_reg,
+ uint64_t num_reg,
+ uint64_t partnerID) {
+ uint64_t warl = 0;
+ tensor_reduce(start_reg, warl, num_reg, partnerID, 0);
+}
+
+/*! \fn inline void tensor_reduce_recv(uint64_t start_reg, uint64_t operation, uint64_t num_reg, uint64_t partnerID)
+ \brief This function recieves reduce function from partner minion.
+ \param start_reg starting register
+ \param operation operation to be performed
+ \param num_reg number of registers
+ \param partnerID Receiver minionID.
+ \return none
+ \tensorops Implementation of tensor_reduce_recv api
+*/
+inline void __attribute__((always_inline)) tensor_reduce_recv(uint64_t start_reg,
+ uint64_t operation,
+ uint64_t num_reg,
+ uint64_t partnerID) {
+ tensor_reduce(start_reg, operation, num_reg, partnerID, 1);
+}
+
+/*! \fn inline void tensor_reduce_auto(uint64_t start_reg, uint64_t operation, uint64_t num_reg, uint64_t tree_depth)
+ \brief The Tensor reduce instruction allows up to 216 harts to collectively calculate a reduction function.
+ \param start_reg starting register
+ \param operation operation to be performed
+ \param num_reg number of registers
+ \param tree_depth tree depth
+ \return none
+ \tensorops Implementation of tensor_reduce_auto api
+*/
+inline void __attribute__((always_inline)) tensor_reduce_auto(uint64_t start_reg,
+ uint64_t operation,
+ uint64_t num_reg,
+ uint64_t tree_depth) {
+ tensor_reduce(start_reg, operation, num_reg, (0ULL << 4) | (tree_depth & 0xF), 3);
+}
+
+/*! \fn inline void tensor_broadcast(uint64_t start_reg, uint64_t operation, uint64_t num_reg, uint64_t tree_depth) {
+ \brief The Tensor broadcast instruction allows up to 216 harts to receive values held in the vector registers
+ of one of the harts in the group. The broadcast operation is performed in a binary-tree fashion, where the source
+ data is originally in the root node and the final result ends up in the leaf nodes.
+ \param start_reg Starting floating-point register
+ \param operation operation to be performed
+ \param num_reg Number of floating-point registers
+ \param tree_depth tree depth
+ \return none
+ \tensorops Implementation of tensor_broadcast api
+*/
+inline void __attribute__((always_inline)) tensor_broadcast(uint64_t start_reg,
+ uint64_t operation,
+ uint64_t num_reg,
+ uint64_t tree_depth) {
+ tensor_reduce(start_reg, operation, num_reg, (0ULL << 4) | (tree_depth & 0xF), 2);
+}
+
+/*! \fn inline void tensor_reduce_autopair(uint64_t start_reg, uint64_t operation, uint64_t num_reg, uint64_t start_lvl, uint64_t end_lvl, uint64_t action) {
+ \brief This function is wrapper of Tensor Reduce (auto-pair variant) instruction.
+ \param start_reg Starting floating-point register
+ \param operation Function to be performed
+ \param num_reg Number of floating-point registers
+ \param start_lvl starting level value
+ \param end_lvl ending level value
+ \param action action value
+ \return none
+ \tensorops Implementation of tensor_reduce_autopair api
+
+*/
+inline void __attribute__((always_inline)) tensor_reduce_autopair(uint64_t start_reg,
+ uint64_t operation,
+ uint64_t num_reg,
+ uint64_t start_lvl,
+ uint64_t end_lvl,
+ uint64_t action) {
+ uint64_t partnerID;
+ // PRM-10 defines the partnerID field for Tensor Reduce (auto-pair variant) as following:
+ // [15:11] WARL(0)
+ // [10: 7] End level for autopair
+ // [ 6: 3] Start level for autopair
+ uint64_t warl = 0;
+ partnerID = ((warl & 0xF) << 11) | ((end_lvl & 0xF) << 7) | ((start_lvl & 0xF) << 3);
+ // Operations encoding:
+ // 0000=fadd, 0001=fsub, 0010=fmax, 0011=fmin, 0100=iadd, 0101=isub, 0110=imax, 0111=imin, 1000=fget
+ //
+ // Action encoding:
+ // 00=send, 01=receive, 10=auto-pair broadcast derive from hartid,11=auto-pair reduce derive from hartid
+ tensor_reduce(start_reg, operation, num_reg, (partnerID >> 3), action);
+}
+
+/*! \fn inline void tensor_quant(uint64_t start_reg, uint64_t col, uint64_t row, uint64_t scp_loc, uint64_t transf9, uint64_t transf8, uint64_t transf7, uint64_t transf6, uint64_t transf5, uint64_t transf4, uint64_t transf3, uint64_t transf2, uint64_t transf1, uint64_t transf0 )
+ \brief Tensor quantization (TensorQuant) instructions are encoded as writes to the tensor_quant CSR. The TensorQuant
+ instruction performs a sequence of up to 10 transformations to a matrix A
+ \param start_reg Starting register
+ \param col A matrix number of columns.
+ \param row A matrix number of rows.
+ \param scp_loc L1 scratchpad cache line where the first vector is stored.
+ \param transf9 Transformation 9.
+ \param transf8 Transformation 8.
+ \param transf7 Transformation 7.
+ \param transf6 Transformation 6.
+ \param transf5 Transformation 5.
+ \param transf4 Transformation 4.
+ \param transf3 Transformation 3.
+ \param transf2 Transformation 2.
+ \param transf1 Transformation 1.
+ \param transf0 Transformation 0.
+ \return none
+ \tensorops Implementation of tensor_quant api
+*/
+inline void __attribute__((always_inline)) tensor_quant(uint64_t start_reg,
+ uint64_t col,
+ uint64_t row,
+ uint64_t scp_loc,
+ uint64_t transf9,
+ uint64_t transf8,
+ uint64_t transf7,
+ uint64_t transf6,
+ uint64_t transf5,
+ uint64_t transf4,
+ uint64_t transf3,
+ uint64_t transf2,
+ uint64_t transf1,
+ uint64_t transf0) {
+ uint64_t csr_enc = ((start_reg & 0x1F) << 57) | ((col & 0x3) << 55) | ((row & 0xF) << 51) |
+ ((scp_loc & 0x3F) << 45) | ((transf9 & 0xF) << 36) | ((transf8 & 0xF) << 32) |
+ ((transf7 & 0xF) << 28) | ((transf6 & 0xF) << 24) | ((transf5 & 0xF) << 20) |
+ ((transf4 & 0xF) << 16) | ((transf3 & 0xF) << 12) | ((transf2 & 0xF) << 8) |
+ ((transf1 & 0xF) << 4) | ((transf0 & 0xF) << 0);
+
+ __asm__ __volatile__("csrw 0x806, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) :);
+}
+
+/*! \fn inline void tensor_mask(uint64_t zeros, uint64_t mask_bits)
+ \brief The TensorLoad, TensorFMA, and CacheOp instructions can operate under the
+ control of the tensor_mask CSR. The tensor_mask CSR contains one bit for each
+ of the destination lines that TensorLoad can potentially write into the scratchpad
+ \param zeros all zeros
+ \param mask_bits tensor bit mask
+ \return none
+ \tensorops Implementation of tensor_mask api
+*/
+inline void __attribute__((always_inline)) tensor_mask(uint64_t zeros, uint64_t mask_bits) {
+ uint64_t csr_enc = ((zeros & 0x000000000000) << 16) | (mask_bits & 0xFFFF);
+
+ __asm__ __volatile__("csrw 0x805, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) :);
+}
+
+/*! \fn inline void tensor_coop(uint64_t val)
+ \brief The tensor_coop instruction specifies which harts participate in cooperative tensor load operations. Only the first hart of each
+ selected Minion core participates in the cooperative operations, since the second hart cannot issue tensor load operations.
+ \param val value contains encoded coop id, minion and neigh mask
+ \return none
+ \tensorops Implementation of tensor_coop api
+*/
+inline void __attribute__((always_inline)) tensor_coop(uint64_t val) {
+ __asm__ __volatile__("csrw 0x804, %[val]\n" : : [val] "r"(val) :);
+}
+
+/*! \fn inline void convolution_ctrl(uint64_t row_start, uint64_t col_start)
+ \brief This function modifies the convolution control register.
+ This register encodes the location of a tensor inside a larger two-dimensional array.
+ \param row_start signed integer value specifying the row inside the array where the first row of the tensor resides
+ \param col_start signed integer value specifying the column inside the array where the first column of the tensor resides
+ \return none
+ \tensorops Implementation of convolution_ctrl api
+*/
+inline void __attribute__((always_inline)) convolution_ctrl(uint64_t row_start, uint64_t col_start) {
+ uint64_t csr_enc = ((row_start & 0xFFFF) << 32) | (col_start & 0xFFFF);
+
+ __asm__ __volatile__("csrw 0x803, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) :);
+}
+
+/*! \fn inline void convolution_size(uint64_t srow, uint64_t nrow, uint64_t scol, uint64_t ncol)
+ \brief This function modifies the convolution size register.
+ This register specifies the layout of a two-dimensional array used for convolutions.
+ \param srow integer value specifying the row inside the array where the first row of the tensor resides
+ \param nrow integer values specifying the number of rows of the array
+ \param scol integer value specifying the distance, in number of columns, between consecutive column accesses to the array during
+ convolution operations
+ \param ncol integer values specifying the number of columns of the array
+ \return none
+ \tensorops Implementation of convolution_size api
+*/
+inline void __attribute__((always_inline)) convolution_size(uint64_t srow,
+ uint64_t nrow,
+ uint64_t scol,
+ uint64_t ncol) {
+ uint64_t csr_enc = ((srow & 0xFF) << 56) | ((nrow & 0xFFFF) << 32) | ((scol & 0xFF) << 24) | ((ncol & 0xFFFF));
+
+ __asm__ __volatile__("csrw 0x802, %[csr_enc]\n" : : [csr_enc] "r"(csr_enc) :);
+}
+
+/*! \fn inline unsigned get_tensor_error()
+ \brief This function returns tensor error register value.
+ The tensor_error register accrues errors that occur during the execution of tensor instructions and cache management operations. When the tensor coprocessor or the cache management coprocessor generates an exception, the exception is recorded in
+ the tensor_error register and execution does not trap. The tensor_error register is never cleared by the implementation. It is the
+ responsibility of the software to clear tensor_error
+ \return Tensor error value
+ \tensorops Implementation of get_tensor_error api
+*/
+inline unsigned long __attribute__((always_inline)) get_tensor_error() {
+ unsigned long error;
+
+ __asm__ __volatile__("csrr %0, 0x808" : "=r"(error));
+
+ return error;
+}
+
+/*! \fn inline uint64_t get_tensor_mask()
+ \brief This function returns tensor mask register value.
+ \return Tensor mask value
+ \tensorops Implementation of get_tensor_mask api
+*/
+inline uint64_t __attribute__((always_inline)) get_tensor_mask() {
+ uint64_t val;
+
+ __asm__ __volatile__("csrr %0, 0x805" : "=r"(val));
+
+ return val;
+}
+
+#define mask_set(msk, val) \
+ do { \
+ __asm__ volatile("mov.m.x m" #msk ", zero, %0" ::"n"(val)); \
+ } while (0)
+
+#define flw_ps(fd, ptr) \
+ do { \
+ __asm__ volatile("flw.ps f" #fd ", (%0)" ::"r"(ptr)); \
+ } while (0)
+
+#define fsw_ps(fd, ptr) \
+ do { \
+ __asm__ volatile("fsw.ps f" #fd ", (%0)" ::"r"(ptr) : "memory"); \
+ } while (0)
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // ! __TENSORS_H
--- /dev/null
+//******************************************************************************
+// Tri F32 Kernel
+// Triangular masking: zero out elements outside the triangular region.
+//
+// tri_type (matches ggml_tri_type enum):
+// 0 = UPPER_DIAG: keep where i0 >= i1
+// 1 = UPPER: keep where i0 > i1
+// 2 = LOWER_DIAG: keep where i0 <= i1
+// 3 = LOWER: keep where i0 < i1
+//
+// Distribution: cache-line aligned chunks of the flat contiguous dst.
+// Each element is individually classified as keep or zero based on its
+// (i0, i1) coordinates. This avoids cache-line sharing between threads
+// when ne0 is not a multiple of 16.
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+#define TRI_TYPE_UPPER_DIAG 0
+#define TRI_TYPE_UPPER 1
+#define TRI_TYPE_LOWER_DIAG 2
+#define TRI_TYPE_LOWER 3
+
+struct ggml_et_tri_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+ int32_t tri_type;
+};
+
+static inline int keep_element(int32_t tri_type, int64_t i0, int64_t i1) {
+ switch (tri_type) {
+ case TRI_TYPE_LOWER:
+ return i0 < i1;
+ case TRI_TYPE_LOWER_DIAG:
+ return i0 <= i1;
+ case TRI_TYPE_UPPER:
+ return i0 > i1;
+ case TRI_TYPE_UPPER_DIAG:
+ return i0 >= i1;
+ default:
+ return 0;
+ }
+}
+
+int entry_point(struct ggml_et_tri_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+ int32_t tri_type = params->tri_type;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ const int64_t ne0 = dst->ne[0];
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t ne3 = dst->ne[3];
+
+ const size_t nb01 = src0->nb[1], nb02 = src0->nb[2], nb03 = src0->nb[3];
+ const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3];
+
+ const int64_t total_rows = ne1 * ne2 * ne3;
+
+ //==========================================================================
+ // Fast path: ne0 % 16 == 0 — rows are cache-line aligned, distribute rows
+ //==========================================================================
+ if (ne0 % 16 == 0) {
+ float zero = 0.0f;
+ __asm__ volatile("fbc.ps f10, %[z]\n" : : [z] "m"(zero) : "f10");
+
+ for (int64_t row = thread_id; row < total_rows; row += num_threads) {
+ const int64_t i1 = row % ne1;
+ const int64_t i2 = (row / ne1) % ne2;
+ const int64_t i3 = row / (ne1 * ne2);
+
+ const float * src_row = (const float *) ((const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+ float * dst_row = (float *) ((char *) dst_data + i1 * nb1 + i2 * nb2 + i3 * nb3);
+
+ int64_t keep_start, keep_end;
+ switch (tri_type) {
+ case TRI_TYPE_LOWER:
+ keep_start = 0;
+ keep_end = i1;
+ break;
+ case TRI_TYPE_LOWER_DIAG:
+ keep_start = 0;
+ keep_end = i1 + 1;
+ break;
+ case TRI_TYPE_UPPER:
+ keep_start = i1 + 1;
+ keep_end = ne0;
+ break;
+ case TRI_TYPE_UPPER_DIAG:
+ keep_start = i1;
+ keep_end = ne0;
+ break;
+ default:
+ return -1;
+ }
+ if (keep_end > ne0) {
+ keep_end = ne0;
+ }
+
+ // Zero prefix [0, keep_start) — SIMD for aligned blocks, scalar tail
+ int64_t i0 = 0;
+ for (; i0 + 8 <= keep_start; i0 += 8) {
+ __asm__ volatile("fsw.ps f10, %[d]\n" : [d] "=m"(*(float (*)[8]) & dst_row[i0])::"f10");
+ }
+ for (; i0 < keep_start; i0++) {
+ dst_row[i0] = 0.0f;
+ }
+
+ // Copy kept region [keep_start, keep_end) — SIMD + scalar tail
+ for (; i0 + 8 <= keep_end; i0 += 8) {
+ __asm__ volatile(
+ "flw.ps f11, %[s]\n"
+ "fsw.ps f11, %[d]\n"
+ : [d] "=m"(*(float (*)[8]) & dst_row[i0])
+ : [s] "m"(*(const float (*)[8]) & src_row[i0])
+ : "f11");
+ }
+ for (; i0 < keep_end; i0++) {
+ dst_row[i0] = src_row[i0];
+ }
+
+ // Zero suffix [keep_end, ne0) — SIMD + scalar tail
+ for (; i0 + 8 <= ne0; i0 += 8) {
+ __asm__ volatile("fsw.ps f10, %[d]\n" : [d] "=m"(*(float (*)[8]) & dst_row[i0])::"f10");
+ }
+ for (; i0 < ne0; i0++) {
+ dst_row[i0] = 0.0f;
+ }
+ }
+ return 0;
+ }
+
+ //==========================================================================
+ // Unaligned fallback: distribute by cache lines, scalar per element
+ //==========================================================================
+ {
+ const int64_t total_elements = ne0 * ne1 * ne2 * ne3;
+ const int64_t elems_per_cl = 16;
+ const int64_t total_cl = (total_elements + elems_per_cl - 1) / elems_per_cl;
+
+ const int64_t cl_per_thread = (total_cl + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cl) {
+ cl_end = total_cl;
+ }
+ if (cl_start >= total_cl) {
+ return 0;
+ }
+
+ const int64_t es = cl_start * elems_per_cl;
+ int64_t ee = cl_end * elems_per_cl;
+ if (ee > total_elements) {
+ ee = total_elements;
+ }
+
+ int64_t row_idx = es / ne0;
+ int64_t col = es % ne0;
+
+ int64_t pos = es;
+ while (pos < ee) {
+ const int64_t i1 = row_idx % ne1;
+ const int64_t i2 = (row_idx / ne1) % ne2;
+ const int64_t i3 = row_idx / (ne1 * ne2);
+
+ const float * src_row = (const float *) ((const char *) src0_data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+
+ int64_t row_remaining = ne0 - col;
+ int64_t chunk_remaining = ee - pos;
+ int64_t n = row_remaining < chunk_remaining ? row_remaining : chunk_remaining;
+
+ int64_t keep_start, keep_end;
+ switch (tri_type) {
+ case TRI_TYPE_LOWER:
+ keep_start = 0;
+ keep_end = i1;
+ break;
+ case TRI_TYPE_LOWER_DIAG:
+ keep_start = 0;
+ keep_end = i1 + 1;
+ break;
+ case TRI_TYPE_UPPER:
+ keep_start = i1 + 1;
+ keep_end = ne0;
+ break;
+ case TRI_TYPE_UPPER_DIAG:
+ keep_start = i1;
+ keep_end = ne0;
+ break;
+ default:
+ return -1;
+ }
+ if (keep_end > ne0) {
+ keep_end = ne0;
+ }
+
+ int64_t end_col = col + n;
+ for (int64_t i0 = col; i0 < end_col; i0++) {
+ if (i0 >= keep_start && i0 < keep_end) {
+ dst_data[pos + (i0 - col)] = src_row[i0];
+ } else {
+ dst_data[pos + (i0 - col)] = 0.0f;
+ }
+ }
+
+ pos += n;
+ col = 0;
+ row_idx++;
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+#include "ggml-et-uberkernel-common.h"
+#include "ggml-et-uberkernel-kernel-map.h"
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+struct ggml_et_glu_params;
+struct ggml_et_unary_params;
+struct ggml_et_rope_params;
+struct ggml_et_rms_norm_params;
+struct ggml_et_rms_norm_mul_params;
+struct ggml_et_softmax_params;
+struct ggml_et_set_rows_params;
+struct ggml_et_get_rows_params;
+struct ggml_et_cont_params;
+struct ggml_et_concat_params;
+struct ggml_et_cumsum_params;
+struct ggml_et_diag_params;
+struct ggml_et_fill_params;
+struct ggml_et_flash_attn_ext_params;
+struct ggml_et_gated_delta_net_params;
+struct ggml_et_group_norm_params;
+struct ggml_et_im2col_params;
+struct ggml_et_l2_norm_params;
+struct ggml_et_mul_mat_id_params;
+struct ggml_et_norm_params;
+struct ggml_et_pad_params;
+struct ggml_et_repeat_params;
+struct ggml_et_rwkv_wkv6_params;
+struct ggml_et_rwkv_wkv7_params;
+struct ggml_et_scale_params;
+struct ggml_et_set_params;
+struct ggml_et_solve_tri_params;
+struct ggml_et_sqr_params;
+struct ggml_et_ssm_conv_params;
+struct ggml_et_ssm_scan_params;
+struct ggml_et_sum_rows_params;
+struct ggml_et_tri_params;
+
+extern int el_map_f32_entry(struct ggml_et_binary_params *, void *);
+extern int glu_f32_entry(struct ggml_et_glu_params *, void *);
+extern int unary_f32_entry(struct ggml_et_unary_params *, void *);
+extern int rope_f32_entry(struct ggml_et_rope_params *, void *);
+extern int rms_norm_f32_entry(struct ggml_et_rms_norm_params *, void *);
+extern int rms_norm_mul_f32_entry(struct ggml_et_rms_norm_mul_params *, void *);
+extern int softmax_f32_entry(struct ggml_et_softmax_params *, void *);
+extern int set_rows_f32_entry(struct ggml_et_set_rows_params *, void *);
+extern int get_rows_f32_entry(struct ggml_et_get_rows_params *, void *);
+extern int cont_f32_entry(struct ggml_et_cont_params *, void *);
+extern int cont_f16_entry(struct ggml_et_cont_params *, void *);
+extern int cpy_f32_f16_entry(struct ggml_et_cont_params *, void *);
+extern int concat_f32_entry(struct ggml_et_concat_params *, void *);
+extern int cumsum_f32_entry(struct ggml_et_cumsum_params *, void *);
+extern int diag_f32_entry(struct ggml_et_diag_params *, void *);
+extern int fill_f32_entry(struct ggml_et_fill_params *, void *);
+extern int flash_attn_ext_f32_entry(struct ggml_et_flash_attn_ext_params *, void *);
+extern int flash_attn_ext_f16_me_entry(struct ggml_et_flash_attn_ext_params *, void *);
+extern int gated_delta_net_f32_entry(struct ggml_et_gated_delta_net_params *, void *);
+extern int group_norm_f32_entry(struct ggml_et_group_norm_params *, void *);
+extern int im2col_entry(struct ggml_et_im2col_params *, void *);
+extern int l2_norm_f32_entry(struct ggml_et_l2_norm_params *, void *);
+extern int mul_mat_id_f32_entry(struct ggml_et_mul_mat_id_params *, void *);
+extern int norm_f32_entry(struct ggml_et_norm_params *, void *);
+extern int pad_f32_entry(struct ggml_et_pad_params *, void *);
+extern int repeat_f32_entry(struct ggml_et_repeat_params *, void *);
+extern int rwkv_wkv6_f32_entry(struct ggml_et_rwkv_wkv6_params *, void *);
+extern int rwkv_wkv7_f32_entry(struct ggml_et_rwkv_wkv7_params *, void *);
+extern int scale_f32_entry(struct ggml_et_scale_params *, void *);
+extern int set_f32_entry(struct ggml_et_set_params *, void *);
+extern int solve_tri_f32_entry(struct ggml_et_solve_tri_params *, void *);
+extern int sqr_f32_entry(struct ggml_et_sqr_params *, void *);
+extern int ssm_conv_f32_entry(struct ggml_et_ssm_conv_params *, void *);
+extern int ssm_scan_f32_entry(struct ggml_et_ssm_scan_params *, void *);
+extern int sum_rows_f32_entry(struct ggml_et_sum_rows_params *, void *);
+extern int tri_f32_entry(struct ggml_et_tri_params *, void *);
+extern int mul_mat_f16_entry(struct ggml_et_binary_params *, void *);
+extern int mul_mat_f16_matrix_engine_entry(struct ggml_et_binary_params *, void *);
+extern int mul_mat_f32_entry(struct ggml_et_binary_params *, void *);
+extern int mul_mat_f32_matrix_engine_entry(struct ggml_et_binary_params *, void *);
+extern int mul_mat_Q8_0_entry(struct ggml_et_mm_q8_params *, void *);
+extern int mul_mat_Q4_0_entry(struct ggml_et_binary_params *, void *);
+
+static inline size_t tensor_bytes(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+struct uber_glu_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+ // trailing scalars omitted — not needed for eviction
+};
+
+struct uber_unary_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+};
+
+struct uber_rope_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor src2;
+ struct ggml_tensor dst;
+};
+
+struct uber_rms_norm_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+};
+
+struct uber_rms_norm_mul_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+struct uber_softmax_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor src2;
+ struct ggml_tensor dst;
+};
+
+struct uber_set_rows_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+struct uber_get_rows_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+struct uber_cont_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor dst;
+};
+
+// src0 + src1 + dst (no trailing scalars needed for eviction)
+struct uber_concat_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+struct uber_ssm_conv_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+struct uber_solve_tri_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor dst;
+};
+
+struct uber_mul_mat_id_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor src2;
+ struct ggml_tensor dst;
+};
+
+// flash_attn_ext: Q=src0, K=src1, V=src2, mask=src3, dst (mask optional)
+struct uber_flash_attn_ext_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor src2;
+ struct ggml_tensor mask;
+ struct ggml_tensor dst;
+};
+
+// ssm_scan: 7 source tensors + dst
+struct uber_ssm_scan_params {
+ struct ggml_tensor src0;
+ struct ggml_tensor src1;
+ struct ggml_tensor src2;
+ struct ggml_tensor src3;
+ struct ggml_tensor src4;
+ struct ggml_tensor src5;
+ struct ggml_tensor src6;
+ struct ggml_tensor dst;
+};
+
+// gated_delta_net: q,k,v,g,beta,state_in,dst
+struct uber_gated_delta_net_params {
+ struct ggml_tensor q;
+ struct ggml_tensor k;
+ struct ggml_tensor v;
+ struct ggml_tensor g;
+ struct ggml_tensor beta;
+ struct ggml_tensor state_in;
+ struct ggml_tensor dst;
+};
+
+static void copy_f32_to_f16_row(uint16_t * dst, const float * src, int64_t num_elements) {
+ for (int64_t i = 0; i < num_elements; i++) {
+ dst[i] = fp32_to_fp16(src[i]);
+ }
+}
+
+static void copy_f32_row(float * dst, const float * src, int64_t num_elements) {
+ for (int64_t i = 0; i < num_elements; i++) {
+ dst[i] = src[i];
+ }
+}
+
+static void evict_region_past_l2_local(const void * addr, size_t bytes) {
+ if (!addr || bytes == 0) {
+ return;
+ }
+
+ const uint64_t CL = 64;
+ uint64_t base = (uint64_t) addr & ~(CL - 1);
+ uint64_t end = ((uint64_t) addr + bytes + CL - 1) & ~(CL - 1);
+ uint64_t nlines = (end - base) / CL;
+ cache_ops_priv_evict_sw(0, /*to_L2*/ 3, 0, 0, CL);
+}
+
+int entry_point(struct ggml_et_uberkernel_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env || !params) {
+ return -1;
+ }
+
+ struct ggml_et_uberkernel_inst * insts = (struct ggml_et_uberkernel_inst *) (uintptr_t) params->insts;
+ uint8_t * params_blob = (uint8_t *) (uintptr_t) params->params_blob;
+
+ if (!insts || !params_blob || params->inst_stride < sizeof(struct ggml_et_uberkernel_inst)) {
+ return -1;
+ }
+
+ for (uint32_t i = 0; i < params->num_insts; ++i) {
+ struct ggml_et_uberkernel_inst * inst =
+ (struct ggml_et_uberkernel_inst *) ((uint8_t *) insts + (i * params->inst_stride));
+ void * inst_params = params_blob + inst->params_offset;
+ int rc = -1;
+
+ et_barrier_global(32ULL);
+
+ switch (inst->kernel_id) {
+ case GGML_ET_UBERKERNEL_KERNEL_EL_MAP_F32:
+ {
+ struct ggml_et_binary_params * p = (struct ggml_et_binary_params *) inst_params;
+ rc = el_map_f32_entry(p, env);
+ break;
+ }
+ // case GGML_ET_UBERKERNEL_KERNEL_UNARY_F32: {
+ // // struct uber_unary_params *p = (struct uber_unary_params *) inst_params;
+ // // et_barrier(ET_BARRIER_GLOBAL);
+ // rc = unary_f32_entry((struct ggml_et_unary_params *) inst_params, env);
+ // break;
+ // }
+ // case GGML_ET_UBERKERNEL_KERNEL_CPY_F32_F16: {
+ // struct uber_unary_params *p = (struct uber_unary_params *) inst_params;
+ // // evict_region_past_l2(p->src0.data, tensor_bytes(&p->src0));
+ // rc = cpy_f32_f16_entry((struct ggml_et_cont_params *) inst_params, env);
+ // break;
+ // }
+ // case GGML_ET_UBERKERNEL_KERNEL_GET_ROWS_F32: {
+ // struct uber_get_rows_params *p = (struct uber_get_rows_params *) inst_params;
+ // rc = get_rows_f32_entry((struct ggml_et_get_rows_params *) inst_params, env);
+ // break;
+ // }
+ // case GGML_ET_UBERKERNEL_KERNEL_CONT_F32: {
+ // struct uber_cont_params *p = (struct uber_cont_params *) inst_params;
+ // // evict_region_past_l2_local(p->src0.data, tensor_bytes(&p->src0));
+ // // evict_region_past_l2(p->dst.data, tensor_bytes(&p->dst));
+ // rc = cont_f32_entry((struct ggml_et_cont_params *) inst_params, env);
+ // break;
+ // }
+ case GGML_ET_UBERKERNEL_KERNEL_GLU_F32:
+ {
+ rc = glu_f32_entry((struct ggml_et_glu_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_ROPE_F32:
+ {
+ rc = rope_f32_entry((struct ggml_et_rope_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_RMS_NORM_F32:
+ {
+ // struct ggml_et_rms_norm_params *p = (struct ggml_et_rms_norm_params *) inst_params;
+ // evict_region_past_l2(p->src0.data, tensor_bytes(&p->src0));
+ rc = rms_norm_f32_entry((struct ggml_et_rms_norm_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_RMS_NORM_MUL_F32:
+ {
+ struct uber_rms_norm_mul_params * p = (struct uber_rms_norm_mul_params *) inst_params;
+ evict_region_past_l2(p->src0.data, tensor_bytes(&p->src0));
+ evict_region_past_l2(p->src1.data, tensor_bytes(&p->src1));
+ rc = rms_norm_mul_f32_entry((struct ggml_et_rms_norm_mul_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_SOFTMAX_F32:
+ {
+ rc = softmax_f32_entry((struct ggml_et_softmax_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_SET_ROWS_F32:
+ {
+ rc = set_rows_f32_entry((struct ggml_et_set_rows_params *) inst_params, env);
+ break;
+ }
+
+ // Single-source ops (src0 → dst)
+ case GGML_ET_UBERKERNEL_KERNEL_SQR_F32:
+ {
+ rc = sqr_f32_entry((struct ggml_et_sqr_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_SCALE_F32:
+ {
+ rc = scale_f32_entry((struct ggml_et_scale_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_SUM_ROWS_F32:
+ {
+ rc = sum_rows_f32_entry((struct ggml_et_sum_rows_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_CUMSUM_F32:
+ {
+ rc = cumsum_f32_entry((struct ggml_et_cumsum_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_NORM_F32:
+ {
+ rc = norm_f32_entry((struct ggml_et_norm_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_L2_NORM_F32:
+ {
+ rc = l2_norm_f32_entry((struct ggml_et_l2_norm_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_GROUP_NORM_F32:
+ {
+ rc = group_norm_f32_entry((struct ggml_et_group_norm_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_REPEAT_F32:
+ {
+ rc = repeat_f32_entry((struct ggml_et_repeat_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_DIAG_F32:
+ {
+ rc = diag_f32_entry((struct ggml_et_diag_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_TRI_F32:
+ {
+ rc = tri_f32_entry((struct ggml_et_tri_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_PAD_F32:
+ {
+ rc = pad_f32_entry((struct ggml_et_pad_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_CONT_F16:
+ {
+ rc = cont_f16_entry((struct ggml_et_cont_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_FILL_F32:
+ {
+ rc = fill_f32_entry((struct ggml_et_fill_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_SET_F32:
+ {
+ rc = set_f32_entry((struct ggml_et_set_params *) inst_params, env);
+ break;
+ }
+
+ // Two-source ops
+ case GGML_ET_UBERKERNEL_KERNEL_CONCAT_F32:
+ {
+ rc = concat_f32_entry((struct ggml_et_concat_params *) inst_params, env);
+ break;
+ }
+ // case GGML_ET_UBERKERNEL_KERNEL_SSM_CONV_F32: {
+ // rc = ssm_conv_f32_entry((struct ggml_et_ssm_conv_params *) inst_params, env);
+ // break;
+ // }
+ case GGML_ET_UBERKERNEL_KERNEL_SOLVE_TRI_F32:
+ {
+ rc = solve_tri_f32_entry((struct ggml_et_solve_tri_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_IM2COL:
+ {
+ rc = im2col_entry((struct ggml_et_im2col_params *) inst_params, env);
+ break;
+ }
+
+ // Three-source ops
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_ID_F32:
+ {
+ rc = mul_mat_id_f32_entry((struct ggml_et_mul_mat_id_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_FLASH_ATTN_EXT_F32:
+ {
+ rc = flash_attn_ext_f32_entry((struct ggml_et_flash_attn_ext_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_FLASH_ATTN_EXT_F16_ME:
+ {
+ rc = flash_attn_ext_f16_me_entry((struct ggml_et_flash_attn_ext_params *) inst_params, env);
+ break;
+ }
+
+ case GGML_ET_UBERKERNEL_KERNEL_GATED_DELTA_NET_F32:
+ {
+ rc = gated_delta_net_f32_entry((struct ggml_et_gated_delta_net_params *) inst_params, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_SSM_SCAN_F32:
+ {
+ rc = ssm_scan_f32_entry((struct ggml_et_ssm_scan_params *) inst_params, env);
+ break;
+ }
+ // rwkv: raw float* params, no ggml_tensor fields to evict via
+ case GGML_ET_UBERKERNEL_KERNEL_RWKV_WKV6_F32:
+ {
+ rc = rwkv_wkv6_f32_entry((struct ggml_et_rwkv_wkv6_params *) inst_params, env);
+ break;
+ }
+
+ case GGML_ET_UBERKERNEL_KERNEL_RWKV_WKV7_F32:
+ {
+ rc = rwkv_wkv7_f32_entry((struct ggml_et_rwkv_wkv7_params *) inst_params, env);
+ break;
+ }
+
+ // MUL_MAT: evict src1 (activations); src0=weights is
+ // read-only so never stale from a prior uberkernel op
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_F16:
+ {
+ struct ggml_et_binary_params * p = (struct ggml_et_binary_params *) inst_params;
+ rc = mul_mat_f16_entry(p, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_F16_MATRIX_ENGINE:
+ {
+ struct ggml_et_binary_params * p = (struct ggml_et_binary_params *) inst_params;
+ rc = mul_mat_f16_matrix_engine_entry(p, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_F32:
+ {
+ struct ggml_et_binary_params * p = (struct ggml_et_binary_params *) inst_params;
+ rc = mul_mat_f32_entry(p, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_F32_MATRIX_ENGINE:
+ {
+ struct ggml_et_binary_params * p = (struct ggml_et_binary_params *) inst_params;
+ rc = mul_mat_f32_matrix_engine_entry(p, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_Q8_0:
+ {
+ struct ggml_et_mm_q8_params * p = (struct ggml_et_mm_q8_params *) inst_params;
+ // evict_region_past_l2(p->src0.data, tensor_bytes(&p->src0));
+ rc = mul_mat_Q8_0_entry(p, env);
+ break;
+ }
+ case GGML_ET_UBERKERNEL_KERNEL_MUL_MAT_Q4_0:
+ {
+ struct ggml_et_binary_params * p = (struct ggml_et_binary_params *) inst_params;
+ rc = mul_mat_Q4_0_entry(p, env);
+ break;
+ }
+
+ default:
+ return -1;
+ }
+
+ if (rc != 0) {
+ return rc;
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+//******************************************************************************
+// Unary F32 Kernel
+// Element-wise unary operations: dst[i] = f(src0[i])
+// All ops vectorized using 8-wide ET SIMD (fexp.ps, frcp.ps, flog.ps, etc.)
+//
+// Supports: ABS, SGN, NEG, STEP, TANH, ELU, RELU, SIGMOID, GELU, GELU_QUICK,
+// SILU, HARDSWISH, HARDSIGMOID, EXP, EXPM1, SOFTPLUS, GELU_ERF
+//******************************************************************************
+
+#include "ggml_tensor.h"
+#include "math_fp.h"
+#include "platform.h"
+
+#include <stdint.h>
+
+// Unary kernel parameters structure
+struct ggml_et_unary_params {
+ struct ggml_tensor src0; // F32 input tensor
+ struct ggml_tensor dst; // F32 output tensor
+ int32_t unary_op; // ggml_unary_op enum value
+};
+
+//******************************************************************************
+// Vectorized 8-wide block operations
+// All process exactly 8 floats per call using ET vector instructions.
+// ne0 is guaranteed % 16 == 0, so the inner loop always calls with i0 += 8.
+//******************************************************************************
+
+// NEG: dst = -x (zero - x)
+static inline void vec_neg(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "flw.ps f11, %[x]\n"
+ "fsub.ps f12, f10, f11\n"
+ "fsw.ps f12, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero)
+ : "f10", "f11", "f12");
+ }
+}
+
+// ABS: dst = |x| (negate negative values: abs = x * sgn, or max(x, -x))
+// Uses: negate then fmax.ps
+static inline void vec_abs(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "flw.ps f11, %[x]\n"
+ "fsub.ps f12, f10, f11\n" // f12 = -x
+ "fmax.ps f13, f11, f12\n" // f13 = max(x, -x) = |x|
+ "fsw.ps f13, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero)
+ : "f10", "f11", "f12", "f13");
+ }
+}
+
+// RELU: dst = max(0, x)
+static inline void vec_relu(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "fbc.ps f10, %[z]\n"
+ "flw.ps f11, %[x]\n"
+ "fmax.ps f12, f10, f11\n" // max(0, x)
+ "fsw.ps f12, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero)
+ : "f10", "f11", "f12");
+ }
+}
+
+// STEP: dst = x > 0 ? 1 : 0 (clamp to [0,1] via max then min-ish, or use sign bit)
+// Trick: relu(x) then frcp gives inf for 0 and finite for >0, but simpler:
+// step(x) = min(1, relu(x) * huge) ... too fragile. Scalar is fine for step/sgn.
+static inline void vec_step(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ dst[i] = (src[i] > 0.0f) ? 1.0f : 0.0f;
+ }
+}
+
+// SGN: dst = sign(x) = x>0 ? 1 : (x<0 ? -1 : 0)
+static inline void vec_sgn(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ dst[i] = (src[i] > 0.0f) ? 1.0f : ((src[i] < 0.0f) ? -1.0f : 0.0f);
+ }
+}
+
+// EXP: dst = exp(x)
+// fexp.ps computes 2^x, so feed x * log2(e)
+static inline void vec_exp(float * dst, const float * src, int32_t n) {
+ float log2e = 1.4426950408889634f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f11, %[l2e]\n"
+ "fmul.ps f12, f10, f11\n" // x * log2(e)
+ "fexp.ps f13, f12\n" // 2^(x*log2e) = exp(x)
+ "fsw.ps f13, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [l2e] "m"(log2e)
+ : "f10", "f11", "f12", "f13");
+ }
+}
+
+// EXPM1: dst = exp(x) - 1
+static inline void vec_expm1(float * dst, const float * src, int32_t n) {
+ float log2e = 1.4426950408889634f;
+ float one = 1.0f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f11, %[l2e]\n"
+ "fbc.ps f14, %[one]\n"
+ "fmul.ps f12, f10, f11\n" // x * log2(e)
+ "fexp.ps f13, f12\n" // exp(x)
+ "fsub.ps f13, f13, f14\n" // exp(x) - 1
+ "fsw.ps f13, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [l2e] "m"(log2e), [one] "m"(one)
+ : "f10", "f11", "f12", "f13", "f14");
+ }
+}
+
+// SIGMOID: dst = 1 / (1 + exp(-x))
+// Same pattern as SwiGLU: exp(-x) via fexp.ps, then frcp.ps
+static inline void vec_sigmoid(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ float one = 1.0f;
+ float log2e = 1.4426950408889634f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[z]\n"
+ "fbc.ps f21, %[one]\n"
+ "fbc.ps f22, %[l2e]\n"
+ "fsub.ps f12, f20, f10\n" // -x
+ "fmul.ps f13, f12, f22\n" // -x * log2(e)
+ "fexp.ps f14, f13\n" // exp(-x)
+ "fadd.ps f15, f14, f21\n" // 1 + exp(-x)
+ "frcp.ps f16, f15\n" // 1 / (1 + exp(-x))
+ "fsw.ps f16, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero), [one] "m"(one), [l2e] "m"(log2e)
+ : "f10", "f12", "f13", "f14", "f15", "f16", "f20", "f21", "f22");
+ }
+}
+
+// TANH: dst = (exp(2x) - 1) / (exp(2x) + 1)
+// Rewrite as: 1 - 2/(exp(2x) + 1) to use frcp.ps
+// Or equivalently: 2*sigmoid(2x) - 1
+static inline void vec_tanh(float * dst, const float * src, int32_t n) {
+ float one = 1.0f;
+ float two = 2.0f;
+ float two_log2e = 2.8853900817779268f; // 2 * log2(e)
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[one]\n"
+ "fbc.ps f21, %[two]\n"
+ "fbc.ps f22, %[tl2e]\n"
+ // exp(2x) via fexp.ps: feed 2x * log2(e)
+ "fmul.ps f12, f10, f22\n" // 2x * log2(e)
+ "fexp.ps f13, f12\n" // exp(2x)
+ "fadd.ps f14, f13, f20\n" // exp(2x) + 1
+ "frcp.ps f15, f14\n" // 1 / (exp(2x) + 1)
+ "fmul.ps f16, f21, f15\n" // 2 / (exp(2x) + 1)
+ "fsub.ps f17, f20, f16\n" // 1 - 2/(exp(2x)+1) = tanh(x)
+ "fsw.ps f17, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [one] "m"(one), [two] "m"(two), [tl2e] "m"(two_log2e)
+ : "f10", "f12", "f13", "f14", "f15", "f16", "f17", "f20", "f21", "f22");
+ }
+}
+
+// SILU: dst = x / (1 + exp(-x)) = x * sigmoid(x)
+// Copied from SwiGLU pattern but without the gate multiply
+static inline void vec_silu(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ float one = 1.0f;
+ float log2e = 1.4426950408889634f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[z]\n"
+ "fbc.ps f21, %[one]\n"
+ "fbc.ps f22, %[l2e]\n"
+ "fsub.ps f12, f20, f10\n" // -x
+ "fmul.ps f13, f12, f22\n" // -x * log2(e)
+ "fexp.ps f14, f13\n" // exp(-x)
+ "fadd.ps f15, f14, f21\n" // 1 + exp(-x)
+ "frcp.ps f16, f15\n" // 1 / (1 + exp(-x))
+ "fmul.ps f17, f10, f16\n" // x * sigmoid(x)
+ "fsw.ps f17, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero), [one] "m"(one), [l2e] "m"(log2e)
+ : "f10", "f12", "f13", "f14", "f15", "f16", "f17", "f20", "f21", "f22");
+ }
+}
+
+// ELU: dst = x > 0 ? x : exp(x) - 1
+// Vector: compute exp(x)-1 for all lanes, then fmax(x, exp(x)-1)
+// Works because for x>0: x > exp(x)-1 is not always true...
+// Actually for x>0, exp(x)-1 > x (since exp(x) > x+1 for x>0).
+// So fmax won't work. Use: compute both, blend via comparison.
+// Simpler: exp(x)-1 for all, then for x>0 overwrite with x.
+// Without per-lane masking, do scalar for ELU.
+static inline void vec_elu(float * dst, const float * src, int32_t n) {
+ float log2e = 1.4426950408889634f;
+ float one = 1.0f;
+ // Compute exp(x)-1 vectorized, then fixup positive elements
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f11, %[l2e]\n"
+ "fbc.ps f14, %[one]\n"
+ "fmul.ps f12, f10, f11\n" // x * log2(e)
+ "fexp.ps f13, f12\n" // exp(x)
+ "fsub.ps f13, f13, f14\n" // exp(x) - 1
+ "fsw.ps f13, %[r]\n" // store exp(x)-1
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [l2e] "m"(log2e), [one] "m"(one)
+ : "f10", "f11", "f12", "f13", "f14");
+ // Fixup: for x > 0, dst = x
+ for (int32_t j = 0; j < 8 && (i + j) < n; j++) {
+ if (src[i + j] > 0.0f) {
+ dst[i + j] = src[i + j];
+ }
+ }
+ }
+}
+
+// GELU: 0.5*x*(1 + tanh(sqrt(2/pi) * x * (1 + 0.044715*x^2)))
+// Reformulated as: x * (1 - 1/(exp(2z)+1)) where z = sqrt(2/pi)*x*(1+0.044715*x^2)
+// NaN-safe: avoids inf*0. Copied from GeGLU block pattern.
+static inline void vec_gelu(float * dst, const float * src, int32_t n) {
+ float one = 1.0f;
+ float half = 0.5f;
+ float coef_a = 0.044715f;
+ float sqrt2pi = 0.79788456080286535587989211986876f;
+ float two_log2e = 2.8853900817779268f; // 2 * log2(e)
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[one]\n"
+ "fbc.ps f21, %[half]\n"
+ "fbc.ps f22, %[coef]\n"
+ "fbc.ps f23, %[s2pi]\n"
+ "fbc.ps f24, %[tl2e]\n"
+ // inner = 1 + 0.044715 * x^2
+ "fmul.ps f12, f10, f10\n" // x^2
+ "fmadd.ps f13, f22, f12, f20\n" // 1 + 0.044715*x^2
+ // z = sqrt(2/pi) * x * inner
+ "fmul.ps f14, f23, f10\n" // sqrt(2/pi) * x
+ "fmul.ps f14, f14, f13\n" // z
+ // exp(2z) via fexp.ps
+ "fmul.ps f15, f14, f24\n" // 2z * log2(e)
+ "fexp.ps f15, f15\n" // exp(2z)
+ // gelu(x) = 0.5 * x * (1 + tanh(z))
+ // = 0.5 * x * (1 + 1 - 2/(exp(2z)+1))
+ // = x * (1 - 1/(exp(2z)+1)) ... wait, that's tanh-based
+ // Actually: 0.5*x*(1 + tanh) = 0.5*x*(1 + 1 - 2/(e2z+1)) = x*(1 - 1/(e2z+1))
+ // Hmm: tanh = (e2z-1)/(e2z+1) = 1 - 2/(e2z+1)
+ // So 0.5*(1+tanh) = 0.5*(2 - 2/(e2z+1)) = 1 - 1/(e2z+1)
+ // gelu = x * (1 - 1/(e2z+1)) -- matches GeGLU pattern exactly
+ "fadd.ps f16, f15, f20\n" // exp(2z) + 1
+ "frcp.ps f16, f16\n" // 1/(exp(2z) + 1)
+ "fsub.ps f16, f20, f16\n" // 1 - 1/(exp(2z)+1) = sigmoid(2z)
+ "fmul.ps f17, f10, f16\n" // x * sigmoid(2z) = gelu(x)
+ "fsw.ps f17, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [one] "m"(one), [half] "m"(half), [coef] "m"(coef_a),
+ [s2pi] "m"(sqrt2pi), [tl2e] "m"(two_log2e)
+ : "f10", "f12", "f13", "f14", "f15", "f16", "f17", "f20", "f21", "f22", "f23", "f24");
+ }
+}
+
+// GELU_QUICK: x * sigmoid(1.702 * x) = x / (1 + exp(-1.702*x))
+static inline void vec_gelu_quick(float * dst, const float * src, int32_t n) {
+ float one = 1.0f;
+ // -1.702 * log2(e) precomputed
+ float neg_coef_log2e = -1.702f * 1.4426950408889634f; // ~ -2.4542
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[one]\n"
+ "fbc.ps f21, %[ncl2e]\n"
+ // exp(-1.702*x): feed -1.702*x*log2(e) = x * (-1.702*log2(e))
+ "fmul.ps f12, f10, f21\n" // x * (-1.702*log2(e))
+ "fexp.ps f13, f12\n" // exp(-1.702*x)
+ "fadd.ps f14, f13, f20\n" // 1 + exp(-1.702*x)
+ "frcp.ps f15, f14\n" // sigmoid(1.702*x)
+ "fmul.ps f16, f10, f15\n" // x * sigmoid(1.702*x)
+ "fsw.ps f16, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [one] "m"(one), [ncl2e] "m"(neg_coef_log2e)
+ : "f10", "f12", "f13", "f14", "f15", "f16", "f20", "f21");
+ }
+}
+
+// GELU_ERF: 0.5 * x * (1 + erf(x / sqrt(2)))
+// erf approximation (Abramowitz & Stegun) is hard to vectorize cleanly, keep scalar
+// but use et_expf for the exp(-z^2) part
+static inline void vec_gelu_erf(float * dst, const float * src, int32_t n) {
+ const float SQRT_2_INV = 0.70710678118654752440084436210484f;
+ for (int32_t i = 0; i < n; i++) {
+ float x = src[i];
+ float z = x * SQRT_2_INV;
+ float az = z < 0.0f ? -z : z;
+
+ float t = et_fdiv(1.0f, 1.0f + 0.3275911f * az);
+ float t2 = t * t;
+ float t3 = t2 * t;
+ float t4 = t3 * t;
+ float t5 = t4 * t;
+
+ float poly = 0.254829592f * t - 0.284496736f * t2 + 1.421413741f * t3 - 1.453152027f * t4 + 1.061405429f * t5;
+
+ float erf_pos = 1.0f - poly * et_expf(-(az * az));
+ float erf_val = (z < 0.0f) ? -erf_pos : erf_pos;
+ dst[i] = 0.5f * x * (1.0f + erf_val);
+ }
+}
+
+// HARDSIGMOID: min(1, max(0, (x + 3) / 6))
+// Vector: compute (x+3)/6 via frcp, then clamp with fmax(0) and fmin(1)
+static inline void vec_hardsigmoid(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ float one = 1.0f;
+ float three = 3.0f;
+ float inv6 = 0.16666666666666666f; // 1/6
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[z]\n"
+ "fbc.ps f21, %[one]\n"
+ "fbc.ps f22, %[thr]\n"
+ "fbc.ps f23, %[inv]\n"
+ "fadd.ps f12, f10, f22\n" // x + 3
+ "fmul.ps f13, f12, f23\n" // (x + 3) / 6
+ "fmax.ps f14, f13, f20\n" // max(0, ...)
+ "fmin.ps f15, f14, f21\n" // min(1, ...)
+ "fsw.ps f15, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero), [one] "m"(one), [thr] "m"(three), [inv] "m"(inv6)
+ : "f10", "f12", "f13", "f14", "f15", "f20", "f21", "f22", "f23");
+ }
+}
+
+// HARDSWISH: x * hardsigmoid(x) = x * min(1, max(0, (x+3)/6))
+static inline void vec_hardswish(float * dst, const float * src, int32_t n) {
+ float zero = 0.0f;
+ float one = 1.0f;
+ float three = 3.0f;
+ float inv6 = 0.16666666666666666f;
+ for (int32_t i = 0; i < n; i += 8) {
+ __asm__ volatile(
+ "flw.ps f10, %[x]\n"
+ "fbc.ps f20, %[z]\n"
+ "fbc.ps f21, %[one]\n"
+ "fbc.ps f22, %[thr]\n"
+ "fbc.ps f23, %[inv]\n"
+ "fadd.ps f12, f10, f22\n" // x + 3
+ "fmul.ps f13, f12, f23\n" // (x + 3) / 6
+ "fmax.ps f14, f13, f20\n" // max(0, ...)
+ "fmin.ps f15, f14, f21\n" // min(1, ...)
+ "fmul.ps f16, f10, f15\n" // x * hardsigmoid(x)
+ "fsw.ps f16, %[r]\n"
+ : [r] "=m"(*(float (*)[8]) & dst[i])
+ : [x] "m"(*(const float (*)[8]) & src[i]), [z] "m"(zero), [one] "m"(one), [thr] "m"(three), [inv] "m"(inv6)
+ : "f10", "f12", "f13", "f14", "f15", "f16", "f20", "f21", "f22", "f23");
+ }
+}
+
+// FLOOR: largest integer <= x
+static inline void vec_floor(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ float x = src[i];
+ float t = (float) (int32_t) x;
+ dst[i] = (t > x) ? t - 1.0f : t;
+ }
+}
+
+// CEIL: smallest integer >= x
+static inline void vec_ceil(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ float x = src[i];
+ float t = (float) (int32_t) x;
+ dst[i] = (t < x) ? t + 1.0f : t;
+ }
+}
+
+// TRUNC: round towards zero
+static inline void vec_trunc(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ dst[i] = (float) (int32_t) src[i];
+ }
+}
+
+// ROUND: round to nearest, ties to even (banker's rounding)
+static inline void vec_round(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ float x = src[i];
+ float t = (float) (int32_t) x;
+ float diff = x - t;
+ if (diff > 0.5f || (diff == 0.5f && ((int32_t) t & 1))) {
+ t += 1.0f;
+ } else if (diff < -0.5f || (diff == -0.5f && ((int32_t) t & 1))) {
+ t -= 1.0f;
+ }
+ dst[i] = t;
+ }
+}
+
+// SOFTPLUS: log(1 + exp(x))
+// For large x (>20), softplus(x) ~ x. For moderate x, use fexp + flog.
+// Scalar fallback since flog.ps computes log2, need conversion, and overflow guard
+static inline void vec_softplus(float * dst, const float * src, int32_t n) {
+ for (int32_t i = 0; i < n; i++) {
+ float x = src[i];
+ dst[i] = (x > 20.0f) ? x : et_logf(1.0f + et_expf(x));
+ }
+}
+
+static inline size_t tensor_bytes(const struct ggml_tensor * t) {
+ return (size_t) t->ne[0] * t->ne[1] * t->ne[2] * t->ne[3] * t->nb[0];
+}
+
+//******************************************************************************
+// Main entry point
+//******************************************************************************
+
+int entry_point(struct ggml_et_unary_params * params, void * env) {
+ kernel_environment_t * kernel_env = (kernel_environment_t *) env;
+
+ if (!kernel_env) {
+ return -1;
+ }
+
+ int thread_id = get_relative_thread_id(kernel_env->shire_mask);
+ int num_threads = get_num_threads(kernel_env->shire_mask);
+
+ if (thread_id < 0) {
+ return 0;
+ }
+
+ if (params == 0 || ((uint64_t) params & 0x7) != 0) {
+ return -1;
+ }
+
+ struct ggml_tensor * src0 = ¶ms->src0;
+ struct ggml_tensor * dst = ¶ms->dst;
+
+ // evict_region_past_l2(¶ms->unary_op, sizeof(int32_t));
+ // WAIT_CACHEOPS;
+ // FENCE;
+
+ int32_t unary_op = params->unary_op;
+
+ if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+ return -1;
+ }
+
+ float * src0_data = (float *) src0->data;
+ float * dst_data = (float *) dst->data;
+
+ if (!src0_data || !dst_data) {
+ return -1;
+ }
+
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+ // evict_region_past_l2(dst_data, tensor_bytes(dst));
+ // WAIT_CACHEOPS;
+ // FENCE;
+ // et_barrier(ET_BARRIER_GLOBAL);
+
+ // Tensor layout: src and dst are F32 with at least dim-0 contiguity
+ // - nb[0] == sizeof(float) (rows are dense; SIMD loads stay legal)
+ // - nb[1], nb[2], nb[3] may all be arbitrary strides for 4D views
+ //
+ // We walk rows independently and decompose row index r into (i1,i2,i3),
+ // computing per-row byte offsets via nb[1..3] of each tensor.
+ const int64_t nc = dst->ne[0]; // row width (logical)
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t nr = ne1 * ne2 * dst->ne[3]; // total rows
+ const int64_t total_elements = nr * nc;
+ const size_t s_nb1 = src0->nb[1], s_nb2 = src0->nb[2], s_nb3 = src0->nb[3];
+ const size_t d_nb1 = dst->nb[1], d_nb2 = dst->nb[2], d_nb3 = dst->nb[3];
+
+ // evict_region_past_l2(src0_data, tensor_bytes(src0));
+ // evict_region_past_l2(dst_data, tensor_bytes(dst));
+ // FENCE;
+ // WAIT_CACHEOPS;
+ // et_barrier(ET_BARRIER_GLOBAL);
+ const int64_t elements_per_cacheline = 16; // 64 bytes / 4 bytes per float
+ const int64_t total_cachelines = (total_elements + elements_per_cacheline - 1) / elements_per_cacheline;
+
+ const int64_t cl_per_thread = (total_cachelines + num_threads - 1) / num_threads;
+ const int64_t cl_start = thread_id * cl_per_thread;
+ int64_t cl_end = cl_start + cl_per_thread;
+ if (cl_end > total_cachelines) {
+ cl_end = total_cachelines;
+ }
+
+ if (cl_start >= total_cachelines) {
+ return 0;
+ }
+
+ const int64_t elem_start = cl_start * elements_per_cacheline;
+ int64_t elem_end = cl_end * elements_per_cacheline;
+ if (elem_end > total_elements) {
+ elem_end = total_elements;
+ }
+
+ // Fast path: tensor is fully contiguous (no view), walk it as a flat array.
+ // This preserves perf for the common case and avoids the per-row dispatch loop.
+ const size_t row_bytes = (size_t) nc * sizeof(float);
+ // evict_region_past_l2((src0_data + elem_start), row_bytes);
+ // // evict_region_past_l2((dst_data + elem_start), row_bytes);
+ // FENCE;
+ // WAIT_CACHEOPS;
+ // et_barrier(ET_BARRIER_GLOBAL);
+
+ const int is_flat = s_nb1 == row_bytes && s_nb2 == s_nb1 * (size_t) ne1 && s_nb3 == s_nb2 * (size_t) ne2 &&
+ d_nb1 == row_bytes && d_nb2 == d_nb1 * (size_t) ne1 && d_nb3 == d_nb2 * (size_t) ne2;
+
+ if (is_flat) {
+ float * src_ptr = src0_data + elem_start;
+ // evict_region_past_l2(src_ptr, 1024);
+ float * dst_ptr = dst_data + elem_start;
+ const int32_t count = (int32_t) (elem_end - elem_start);
+ switch (unary_op) {
+ case GGML_UNARY_OP_NEG:
+ vec_neg(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_ABS:
+ vec_abs(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SGN:
+ vec_sgn(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_STEP:
+ vec_step(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_RELU:
+ vec_relu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_EXP:
+ vec_exp(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_EXPM1:
+ vec_expm1(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SIGMOID:
+ vec_sigmoid(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_TANH:
+ vec_tanh(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SILU:
+ vec_silu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_ELU:
+ vec_elu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_GELU:
+ vec_gelu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_GELU_QUICK:
+ vec_gelu_quick(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_GELU_ERF:
+ vec_gelu_erf(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_HARDSWISH:
+ vec_hardswish(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_HARDSIGMOID:
+ vec_hardsigmoid(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SOFTPLUS:
+ vec_softplus(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_FLOOR:
+ vec_floor(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_CEIL:
+ vec_ceil(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_ROUND:
+ vec_round(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_TRUNC:
+ vec_trunc(dst_ptr, src_ptr, count);
+ break;
+ default:
+ return -1;
+ }
+ return 0;
+ }
+
+ // Slow path: arbitrary 4D-strided view. Walk the assigned element range
+ // row-by-row, clipping each segment to a row boundary so we never cross
+ // nb[1]. For each row index r, decompose into (i1,i2,i3) and add the
+ // corresponding nb[*] byte offsets to the base pointers.
+ int64_t e = elem_start;
+ while (e < elem_end) {
+ int64_t row = e / nc;
+ int64_t col = e % nc;
+ int64_t take = nc - col;
+ if (take > elem_end - e) {
+ take = elem_end - e;
+ }
+
+ // Decompose row into (i3,i2,i1) using row-major linearization
+ const int64_t i1 = row % ne1;
+ const int64_t r2 = row / ne1;
+ const int64_t i2 = r2 % ne2;
+ const int64_t i3 = r2 / ne2;
+
+ float * src_ptr = (float *) ((char *) src0_data + i3 * s_nb3 + i2 * s_nb2 + i1 * s_nb1) + col;
+ float * dst_ptr = (float *) ((char *) dst_data + i3 * d_nb3 + i2 * d_nb2 + i1 * d_nb1) + col;
+ const int32_t count = (int32_t) take;
+
+ // evict_region_past_l2(src_ptr, 1024);
+ // FENCE;
+ // et_barrier(ET_BARRIER_GLOBAL);
+
+ switch (unary_op) {
+ case GGML_UNARY_OP_NEG:
+ vec_neg(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_ABS:
+ vec_abs(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SGN:
+ vec_sgn(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_STEP:
+ vec_step(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_RELU:
+ vec_relu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_EXP:
+ vec_exp(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_EXPM1:
+ vec_expm1(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SIGMOID:
+ vec_sigmoid(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_TANH:
+ vec_tanh(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SILU:
+ vec_silu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_ELU:
+ vec_elu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_GELU:
+ vec_gelu(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_GELU_QUICK:
+ vec_gelu_quick(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_GELU_ERF:
+ vec_gelu_erf(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_HARDSWISH:
+ vec_hardswish(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_HARDSIGMOID:
+ vec_hardsigmoid(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_SOFTPLUS:
+ vec_softplus(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_FLOOR:
+ vec_floor(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_CEIL:
+ vec_ceil(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_ROUND:
+ vec_round(dst_ptr, src_ptr, count);
+ break;
+ case GGML_UNARY_OP_TRUNC:
+ vec_trunc(dst_ptr, src_ptr, count);
+ break;
+ default:
+ return -1;
+ }
+
+ e += take;
+ }
+
+ return 0;
+}
--- /dev/null
+#pragma once
+
+#include "ggml-backend-impl.h"
+#include "ggml-et-uberkernel-common.h"
+
+#include <device-layer/IDeviceLayer.h>
+#include <runtime/IProfiler.h>
+#include <runtime/IRuntime.h>
+
+#include <cstdint>
+#include <fstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+std::shared_ptr<rt::IRuntime> ggml_et_runtime();
+
+struct ggml_backend_et_buffer_type_context {
+ int devidx;
+ std::string name;
+};
+
+struct ggml_backend_et_buffer_context {
+ int devidx;
+ void * data; // Device memory pointer
+ size_t size;
+ rt::DeviceId rtid;
+};
+
+struct ggml_backend_et_context {
+ int devidx;
+};
+
+struct ggml_backend_et_device_context;
+
+// One slot in the uberkernel ring. The host vectors back the H2D copy and
+// must outlive the upload; the device buffers feed the kernel that consumes
+// them. pending_event lets us know when both have drained so the slot can
+// be recycled.
+struct ggml_backend_et_uberkernel_slot {
+ std::vector<ggml_et_uberkernel_inst> insts;
+ std::vector<std::byte> params_blob;
+
+ std::byte * device_insts = nullptr;
+ std::byte * device_params = nullptr;
+ size_t device_insts_capacity = 0;
+ size_t device_params_capacity = 0;
+
+ rt::EventId pending_event{};
+ bool has_pending = false;
+};
+
+struct ggml_backend_et_uberkernel_context {
+ bool failed = false;
+ uint64_t shire_mask = 0;
+
+ // Ring of slots. We accumulate into slots[current_slot]; on segment
+ // commit we fire the H2D + launch and rotate to the next slot,
+ // waiting on its previous launch only if it hasn't drained yet.
+ static constexpr size_t SLOT_COUNT = 4;
+ ggml_backend_et_uberkernel_slot slots[SLOT_COUNT];
+ size_t current_slot = 0;
+};
+
+struct ggml_backend_et_device_context {
+ int devidx;
+ rt::DeviceId rtid;
+ std::string name;
+ std::string desc;
+ size_t total_mem;
+ ggml_backend_buffer_type_t buftype;
+
+ // Kernel management - default stream for ordered execution on this device
+ rt::StreamId default_stream;
+ std::unordered_map<std::string, rt::KernelId> loaded_kernels;
+
+ // trace buffer - for printing support
+ std::byte * trace_buffer;
+
+ bool uberkernel_enabled = false;
+ ggml_backend_et_uberkernel_context uberkernel;
+};
+
+struct ggml_backend_et_reg_ctx {
+ std::vector<ggml_backend_dev_t> devices;
+};
--- /dev/null
+#include "ggml-et-cpu-compare.h"
+
+#include "ggml-cpu/ggml-cpu-impl.h"
+#include "ggml-cpu/ops.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+
+bool ggml_et_cpu_compare_init_pre(ggml_et_cpu_compare_ctx * ctx, const ggml_tensor * node, ggml_op op) {
+ if (!ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for CPU compare init\n");
+ return false;
+ }
+
+ // Clear context
+ memset(ctx, 0, sizeof(*ctx));
+
+ // Calculate actual buffer sizes - use backend buffer size for accurate copy
+ auto get_tensor_buffer_size = [](const ggml_tensor * tensor) -> size_t {
+ if (!tensor) {
+ return 0;
+ }
+
+ if (tensor->buffer) {
+ // Get actual backend buffer size
+ size_t buffer_size = ggml_backend_buffer_get_size(tensor->buffer);
+
+ // Use the full buffer size to avoid any truncation issues
+ return buffer_size;
+ } else {
+ // Fallback to logical size if no buffer
+ return ggml_nbytes(tensor);
+ }
+ };
+
+ ctx->src0_size = get_tensor_buffer_size(node->src[0]);
+ ctx->src1_size = get_tensor_buffer_size(node->src[1]);
+ ctx->src2_size = get_tensor_buffer_size(node->src[2]);
+ ctx->dst_size = get_tensor_buffer_size(node);
+
+ // Allocate CPU buffers for all tensors
+ if (ctx->src0_size > 0) {
+ ctx->cpu_src0_data = malloc(ctx->src0_size);
+ if (!ctx->cpu_src0_data) {
+ GGML_LOG_ERROR("ET: Failed to allocate CPU src0 buffer\n");
+ goto cleanup;
+ }
+ }
+
+ if (ctx->src1_size > 0) {
+ ctx->cpu_src1_data = malloc(ctx->src1_size);
+ if (!ctx->cpu_src1_data) {
+ GGML_LOG_ERROR("ET: Failed to allocate CPU src1 buffer\n");
+ goto cleanup;
+ }
+ }
+
+ if (ctx->src2_size > 0) {
+ ctx->cpu_src2_data = malloc(ctx->src2_size);
+ if (!ctx->cpu_src2_data) {
+ GGML_LOG_ERROR("ET: Failed to allocate CPU src2 buffer\n");
+ goto cleanup;
+ }
+ }
+
+ ctx->cpu_dst_data = malloc(ctx->dst_size);
+ if (!ctx->cpu_dst_data) {
+ GGML_LOG_ERROR("ET: Failed to allocate CPU dst buffer\n");
+ goto cleanup;
+ }
+
+ ctx->et_dst_data = malloc(ctx->dst_size);
+ if (!ctx->et_dst_data) {
+ GGML_LOG_ERROR("ET: Failed to allocate ET dst buffer\n");
+ goto cleanup;
+ }
+
+ // Copy data from ET device buffers to CPU host buffers
+ if (ctx->src0_size > 0) {
+ // Copy logical tensor size - ggml_backend_tensor_get handles stride layout internally
+ size_t logical_size = ggml_nbytes(node->src[0]);
+ ggml_backend_tensor_get(node->src[0], ctx->cpu_src0_data, 0, logical_size);
+ }
+ if (ctx->src1_size > 0) {
+ size_t logical_size = ggml_nbytes(node->src[1]);
+ ggml_backend_tensor_get(node->src[1], ctx->cpu_src1_data, 0, logical_size);
+ }
+ if (ctx->src2_size > 0) {
+ size_t logical_size = ggml_nbytes(node->src[2]);
+ ggml_backend_tensor_get(node->src[2], ctx->cpu_src2_data, 0, logical_size);
+ }
+
+ // Copy destination data from device (for operations like SET_ROWS that modify existing data)
+ // Most ops create new tensors so this is unused, but SET_ROWS requires existing dst data
+ {
+ size_t logical_size = ggml_nbytes(node);
+ ggml_backend_tensor_get(node, ctx->cpu_dst_data, 0, logical_size);
+ }
+
+ // Create CPU backend for reference computation
+ GGML_LOG_DEBUG("ET: Creating CPU backend for reference computation\n");
+ ctx->cpu_backend = ggml_backend_cpu_init();
+ if (!ctx->cpu_backend) {
+ GGML_LOG_ERROR("ET: Failed to create CPU backend\n");
+ goto cleanup;
+ }
+
+ // Create GGML context for CPU tensors
+ GGML_LOG_DEBUG("ET: Creating GGML context for CPU computation\n");
+ ggml_init_params ctx_params;
+ ctx_params.mem_size = ggml_tensor_overhead() * 4 + ggml_graph_overhead(); // up to 4 tensors + graph
+ ctx_params.mem_buffer = nullptr;
+ ctx_params.no_alloc = true; // We'll manage data ourselves
+ ctx->ggml_ctx = ggml_init(ctx_params);
+ if (!ctx->ggml_ctx) {
+ GGML_LOG_ERROR("ET: Failed to create GGML context\n");
+ goto cleanup;
+ }
+
+ // Create CPU tensors with proper context
+ if (node->src[0]) {
+ ctx->cpu_src0 = ggml_new_tensor(ctx->ggml_ctx, node->src[0]->type, GGML_MAX_DIMS, node->src[0]->ne);
+ if (!ctx->cpu_src0) {
+ GGML_LOG_ERROR("ET: Failed to create CPU src0 tensor\n");
+ goto cleanup;
+ }
+ ctx->cpu_src0->data = ctx->cpu_src0_data;
+ // Copy stride array (nb) for correct memory layout
+ memcpy(ctx->cpu_src0->nb, node->src[0]->nb, sizeof(node->src[0]->nb));
+ // Copy op_params if present
+ memcpy(ctx->cpu_src0->op_params, node->src[0]->op_params, sizeof(node->src[0]->op_params));
+ }
+
+ if (node->src[1]) {
+ ctx->cpu_src1 = ggml_new_tensor(ctx->ggml_ctx, node->src[1]->type, GGML_MAX_DIMS, node->src[1]->ne);
+ if (!ctx->cpu_src1) {
+ GGML_LOG_ERROR("ET: Failed to create CPU src1 tensor\n");
+ goto cleanup;
+ }
+ ctx->cpu_src1->data = ctx->cpu_src1_data;
+ // Copy stride array (nb) for correct memory layout
+ memcpy(ctx->cpu_src1->nb, node->src[1]->nb, sizeof(node->src[1]->nb));
+ // Copy op_params if present
+ memcpy(ctx->cpu_src1->op_params, node->src[1]->op_params, sizeof(node->src[1]->op_params));
+ }
+
+ if (node->src[2]) {
+ ctx->cpu_src2 = ggml_new_tensor(ctx->ggml_ctx, node->src[2]->type, GGML_MAX_DIMS, node->src[2]->ne);
+ if (!ctx->cpu_src2) {
+ GGML_LOG_ERROR("ET: Failed to create CPU src2 tensor\n");
+ goto cleanup;
+ }
+ ctx->cpu_src2->data = ctx->cpu_src2_data;
+ // Copy stride array (nb) for correct memory layout
+ memcpy(ctx->cpu_src2->nb, node->src[2]->nb, sizeof(node->src[2]->nb));
+ // Copy op_params if present
+ memcpy(ctx->cpu_src2->op_params, node->src[2]->op_params, sizeof(node->src[2]->op_params));
+ }
+
+ return true;
+
+cleanup:
+ ggml_et_cpu_compare_free(ctx);
+ return false;
+}
+
+bool ggml_et_cpu_compare_compute_and_check(ggml_et_cpu_compare_ctx * ctx,
+ const ggml_tensor * node,
+ const ggml_et_cpu_compare_config * config) {
+ if (!ctx || !ctx->cpu_backend || !ctx->ggml_ctx || !node || !config) {
+ GGML_LOG_ERROR("ET: Invalid parameters for CPU compute and check\n");
+ return false;
+ }
+
+ // Create operation-specific CPU destination tensor based on the node's operation
+ ggml_op op = node->op;
+ switch (op) {
+ case GGML_OP_MUL:
+ ctx->cpu_dst = ggml_mul(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1);
+ break;
+ case GGML_OP_ADD:
+ ctx->cpu_dst = ggml_add(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1);
+ break;
+ case GGML_OP_MUL_MAT:
+ ctx->cpu_dst = ggml_mul_mat(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1);
+ break;
+ case GGML_OP_MUL_MAT_ID:
+ // MUL_MAT_ID: Mixture of Experts matrix multiplication
+ // src0 (as): expert weight matrices [K, M, n_expert]
+ // src1 (b): activations [K, n_expert_used, batch]
+ // src2 (ids): expert selection indices [n_expert_used, batch]
+ ctx->cpu_dst = ggml_mul_mat_id(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1, ctx->cpu_src2);
+ break;
+ case GGML_OP_ROPE:
+ {
+ const int32_t * op_params = (const int32_t *) node->op_params;
+ const int32_t n_dims = op_params[1];
+ const int32_t mode = op_params[2];
+ const int32_t n_ctx_orig = op_params[4];
+ const float freq_base = *((const float *) (op_params + 5));
+ const float freq_scale = *((const float *) (op_params + 6));
+ const float ext_factor = *((const float *) (op_params + 7));
+ const float attn_factor = *((const float *) (op_params + 8));
+ const float beta_fast = *((const float *) (op_params + 9));
+ const float beta_slow = *((const float *) (op_params + 10));
+
+ if (mode & GGML_ROPE_TYPE_MROPE) {
+ int sections[GGML_MROPE_SECTIONS];
+ memcpy(sections, op_params + 11, sizeof(sections));
+ ctx->cpu_dst = ggml_rope_multi(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1, ctx->cpu_src2,
+ n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ } else {
+ ctx->cpu_dst = ggml_rope_ext(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1, ctx->cpu_src2,
+ n_dims, mode, n_ctx_orig, freq_base, freq_scale, ext_factor,
+ attn_factor, beta_fast, beta_slow);
+ }
+ }
+ break;
+ case GGML_OP_RMS_NORM:
+ // Extract epsilon parameter from op_params (stored as float)
+ {
+ float eps;
+ memcpy(&eps, node->op_params, sizeof(float));
+ ctx->cpu_dst = ggml_rms_norm(ctx->ggml_ctx, ctx->cpu_src0, eps);
+ }
+ break;
+ case GGML_OP_SQR:
+ ctx->cpu_dst = ggml_sqr(ctx->ggml_ctx, ctx->cpu_src0);
+ break;
+ case GGML_OP_UNARY:
+ {
+ ggml_unary_op uop = (ggml_unary_op) ggml_get_op_params_i32(node, 0);
+ ctx->cpu_dst = ggml_unary(ctx->ggml_ctx, ctx->cpu_src0, uop);
+ }
+ break;
+ case GGML_OP_SUM_ROWS:
+ ctx->cpu_dst = ggml_sum_rows(ctx->ggml_ctx, ctx->cpu_src0);
+ break;
+ case GGML_OP_MEAN:
+ ctx->cpu_dst = ggml_mean(ctx->ggml_ctx, ctx->cpu_src0);
+ break;
+ case GGML_OP_CLAMP:
+ {
+ float clamp_min, clamp_max;
+ memcpy(&clamp_min, (const float *) node->op_params + 0, sizeof(float));
+ memcpy(&clamp_max, (const float *) node->op_params + 1, sizeof(float));
+ ctx->cpu_dst = ggml_clamp(ctx->ggml_ctx, ctx->cpu_src0, clamp_min, clamp_max);
+ }
+ break;
+ case GGML_OP_GLU:
+ // Extract GLU parameters from op_params (split mode only)
+ {
+ int32_t glu_op_type = ggml_get_op_params_i32(node, 0); // GLU variant
+ ggml_glu_op glu_op = (ggml_glu_op) glu_op_type;
+
+ // Only support split tensor mode
+ if (!ctx->cpu_src1) {
+ GGML_LOG_ERROR("ET: GLU CPU comparison requires split tensor mode\n");
+ return false;
+ }
+ ctx->cpu_dst = ggml_glu_split(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1, glu_op);
+ }
+ break;
+ case GGML_OP_SOFT_MAX:
+ {
+ // Extract scale and max_bias from op_params
+ float scale = 1.0f;
+ float max_bias = 0.0f;
+ memcpy(&scale, (const float *) node->op_params + 0, sizeof(float));
+ memcpy(&max_bias, (const float *) node->op_params + 1, sizeof(float));
+
+ if (ctx->cpu_src1 || scale != 1.0f || max_bias != 0.0f) {
+ // Use extended softmax when mask or non-default parameters are present
+ ctx->cpu_dst = ggml_soft_max_ext(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1, scale, max_bias);
+ } else {
+ // Use simple softmax when no mask and default parameters
+ ctx->cpu_dst = ggml_soft_max(ctx->ggml_ctx, ctx->cpu_src0);
+ }
+
+ // Add sinks if present
+ if (ctx->cpu_src2) {
+ ggml_soft_max_add_sinks(ctx->cpu_dst, ctx->cpu_src2);
+ }
+ }
+ break;
+ case GGML_OP_GET_ROWS:
+ ctx->cpu_dst = ggml_get_rows(ctx->ggml_ctx, ctx->cpu_src0, ctx->cpu_src1);
+ break;
+ case GGML_OP_CONT:
+ ctx->cpu_dst = ggml_cont(ctx->ggml_ctx, ctx->cpu_src0);
+ break;
+ case GGML_OP_SET_ROWS:
+ {
+ // SET_ROWS operation scatters src0 rows to dst[src1] positions
+ // Create destination tensor (this is the "view" that SET_ROWS returns)
+ ggml_tensor * cpu_dst_base = ggml_new_tensor(ctx->ggml_ctx, node->type, GGML_MAX_DIMS, node->ne);
+ if (!cpu_dst_base) {
+ GGML_LOG_ERROR("ET: Failed to create CPU destination base tensor for SET_ROWS\n");
+ return false;
+ }
+ cpu_dst_base->data = ctx->cpu_dst_data;
+ memcpy(cpu_dst_base->nb, node->nb, sizeof(node->nb));
+
+ // Note: cpu_dst_data already contains the pre-existing destination data from device
+ // SET_ROWS will update specific rows, leaving others unchanged
+
+ // Perform SET_ROWS operation: returns a view that scatters src0 rows to dst[src1] positions
+ ctx->cpu_dst = ggml_set_rows(ctx->ggml_ctx, cpu_dst_base, ctx->cpu_src0, ctx->cpu_src1);
+ }
+ break;
+ default:
+ GGML_LOG_ERROR("ET: Unsupported operation %s for CPU comparison\n", ggml_op_name(op));
+ return false;
+ }
+
+ if (!ctx->cpu_dst) {
+ GGML_LOG_ERROR("ET: Failed to create CPU destination tensor for operation %s\n", ggml_op_name(op));
+ return false;
+ }
+
+ ctx->cpu_dst->data = ctx->cpu_dst_data;
+ // Copy stride array (nb) for correct memory layout - except for CONT which should keep contiguous strides
+ if (op != GGML_OP_CONT) {
+ memcpy(ctx->cpu_dst->nb, node->nb, sizeof(node->nb));
+ }
+ // For CONT operations, keep the contiguous strides created by ggml_cont()
+
+ // Create minimal computation graph
+ ctx->cpu_graph = ggml_new_graph_custom(ctx->ggml_ctx, 1, false);
+ if (!ctx->cpu_graph) {
+ GGML_LOG_ERROR("ET: Failed to create CPU computation graph\n");
+ return false;
+ }
+ ctx->cpu_graph->nodes[0] = ctx->cpu_dst;
+ ctx->cpu_graph->n_nodes = 1;
+
+ // Log input data for debugging if enabled
+ if (config && config->log_differences) {
+ if (ctx->cpu_src0_data && ctx->src0_size >= 4) {
+ GGML_LOG_DEBUG("ET: CPU src0 first few bytes: %02x %02x %02x %02x\n", ((uint8_t *) ctx->cpu_src0_data)[0],
+ ((uint8_t *) ctx->cpu_src0_data)[1], ((uint8_t *) ctx->cpu_src0_data)[2],
+ ((uint8_t *) ctx->cpu_src0_data)[3]);
+ }
+ if (ctx->cpu_src1_data && ctx->src1_size >= 16) {
+ GGML_LOG_DEBUG("ET: CPU src1 first few floats: %.6f %.6f %.6f %.6f\n", ((float *) ctx->cpu_src1_data)[0],
+ ((float *) ctx->cpu_src1_data)[1], ((float *) ctx->cpu_src1_data)[2],
+ ((float *) ctx->cpu_src1_data)[3]);
+ }
+ }
+
+ // Compute using CPU backend
+ ggml_status cpu_result = ggml_backend_graph_compute(ctx->cpu_backend, ctx->cpu_graph);
+
+ if (cpu_result != GGML_STATUS_SUCCESS) {
+ GGML_LOG_ERROR("ET: CPU reference computation failed with status %d\n", cpu_result);
+ return false;
+ }
+
+ // Log output data for debugging if enabled
+ if (config && config->log_differences && ctx->dst_size >= 16) {
+ GGML_LOG_DEBUG("ET: CPU dst first few floats after computation: %.6f %.6f %.6f %.6f\n",
+ ((float *) ctx->cpu_dst_data)[0], ((float *) ctx->cpu_dst_data)[1],
+ ((float *) ctx->cpu_dst_data)[2], ((float *) ctx->cpu_dst_data)[3]);
+ }
+
+ // Now copy ET device destination to host for comparison
+ size_t dst_logical_size = ggml_nbytes(node);
+ ggml_backend_tensor_get(node, ctx->et_dst_data, 0, dst_logical_size);
+
+ if (config->log_differences) {
+ size_t num_elements = ggml_nelements(node);
+ size_t max_log = std::min(num_elements, config->max_log_elements);
+
+ // Check if this is an elementwise operation that can show src inputs
+ bool is_elementwise = (op == GGML_OP_MUL || op == GGML_OP_ADD || op == GGML_OP_GLU);
+ float * cpu_src0_float = is_elementwise ? (float *) ctx->cpu_src0_data : nullptr;
+ float * cpu_src1_float = is_elementwise ? (float *) ctx->cpu_src1_data : nullptr;
+
+ // Helper to get float value from tensor data (handles f16 and f32)
+ auto get_float = [](const void * data, size_t idx, ggml_type type) -> float {
+ if (type == GGML_TYPE_F16) {
+ const ggml_fp16_t * fp16_data = (const ggml_fp16_t *) data;
+ return ggml_fp16_to_fp32(fp16_data[idx]);
+ }
+
+ const float * float_data = (const float *) data;
+ return float_data[idx];
+ };
+
+ // Compare all elements but log only the first max_log_elements
+ bool matches = true;
+ size_t total_mismatches = 0;
+
+ // First pass: check all elements for mismatches
+ for (size_t i = 0; i < num_elements; i++) {
+ float cpu_val = get_float(ctx->cpu_dst_data, i, node->type);
+ float et_val = get_float(ctx->et_dst_data, i, node->type);
+ float diff = fabsf(cpu_val - et_val);
+ float rel_diff = diff / (fabsf(cpu_val) + 1e-8f);
+
+ if (rel_diff > config->tolerance) {
+ matches = false;
+ total_mismatches++;
+ }
+ }
+
+ // Second pass: log detailed info for first max_log elements only
+ for (size_t i = 0; i < max_log; i++) {
+ float cpu_val = get_float(ctx->cpu_dst_data, i, node->type);
+ float et_val = get_float(ctx->et_dst_data, i, node->type);
+ float diff = fabsf(cpu_val - et_val);
+
+ if (is_elementwise && cpu_src0_float && cpu_src1_float) {
+ GGML_LOG_DEBUG("ET: [%zu] src0=%.6f, src1=%.6f -> CPU=%.6f, ET=%.6f, diff=%.6f\n", i, cpu_src0_float[i],
+ cpu_src1_float[i], cpu_val, et_val, diff);
+ } else if (is_elementwise && cpu_src0_float) {
+ GGML_LOG_DEBUG("ET: [%zu] src0=%.6f -> CPU=%.6f, ET=%.6f, diff=%.6f\n", i, cpu_src0_float[i], cpu_val,
+ et_val, diff);
+ } else {
+ GGML_LOG_DEBUG("ET: [%zu] CPU=%.6f, ET=%.6f, diff=%.6f\n", i, cpu_val, et_val, diff);
+ }
+ }
+
+ // Check some elements from the middle and end for full coverage
+ if (num_elements > max_log) {
+ size_t mid = num_elements / 2;
+ size_t end = num_elements - 1;
+ float cpu_mid = get_float(ctx->cpu_dst_data, mid, node->type);
+ float et_mid = get_float(ctx->et_dst_data, mid, node->type);
+ float cpu_end = get_float(ctx->cpu_dst_data, end, node->type);
+ float et_end = get_float(ctx->et_dst_data, end, node->type);
+
+ GGML_LOG_DEBUG("ET: Middle element [%zu]: CPU=%.6f, ET=%.6f\n", mid, cpu_mid, et_mid);
+ GGML_LOG_DEBUG("ET: Last element [%zu]: CPU=%.6f, ET=%.6f\n", end, cpu_end, et_end);
+ }
+
+ GGML_LOG_DEBUG("ET: Results %s (%zu/%zu elements match within tolerance %.6f)\n", matches ? "MATCH" : "DIFFER",
+ num_elements - total_mismatches, num_elements, config->tolerance);
+ }
+
+ // Copy CPU result to device if flag is set
+ if (config->use_cpu_result) {
+ GGML_LOG_DEBUG("ET: Overwriting ET device result with CPU result for correct inference\n");
+ size_t dst_logical_size = ggml_nbytes(node);
+ ggml_backend_tensor_set(const_cast<ggml_tensor *>(node), ctx->cpu_dst_data, 0, dst_logical_size);
+ GGML_LOG_DEBUG("ET: CPU result copied to ET device buffer\n");
+ }
+
+ return true;
+}
+
+void ggml_et_cpu_compare_free(ggml_et_cpu_compare_ctx * ctx) {
+ if (!ctx) {
+ return;
+ }
+
+ if (ctx->cpu_src0_data) {
+ free(ctx->cpu_src0_data);
+ ctx->cpu_src0_data = nullptr;
+ }
+ if (ctx->cpu_src1_data) {
+ free(ctx->cpu_src1_data);
+ ctx->cpu_src1_data = nullptr;
+ }
+ if (ctx->cpu_src2_data) {
+ free(ctx->cpu_src2_data);
+ ctx->cpu_src2_data = nullptr;
+ }
+ if (ctx->cpu_dst_data) {
+ free(ctx->cpu_dst_data);
+ ctx->cpu_dst_data = nullptr;
+ }
+ if (ctx->et_dst_data) {
+ free(ctx->et_dst_data);
+ ctx->et_dst_data = nullptr;
+ }
+
+ if (ctx->ggml_ctx) {
+ ggml_free(ctx->ggml_ctx);
+ ctx->ggml_ctx = nullptr;
+ }
+
+ if (ctx->cpu_backend) {
+ ggml_backend_free(ctx->cpu_backend);
+ ctx->cpu_backend = nullptr;
+ }
+
+ // Clear pointers
+ ctx->cpu_src0 = nullptr;
+ ctx->cpu_src1 = nullptr;
+ ctx->cpu_src2 = nullptr;
+ ctx->cpu_dst = nullptr;
+ ctx->cpu_graph = nullptr;
+}
--- /dev/null
+#pragma once
+
+#include "ggml-cpu.h"
+#include "ggml-et-common.h"
+#include "ggml-impl.h"
+
+// Configuration for CPU comparison
+struct ggml_et_cpu_compare_config {
+ bool enabled; // Whether to enable CPU comparison
+ bool use_cpu_result; // Whether to replace ET result with CPU result
+ bool log_differences; // Whether to log detailed element differences
+ float tolerance; // Relative tolerance for comparison (default: 1e-5f)
+ size_t max_log_elements; // Maximum number of elements to log (default: 10)
+};
+
+// Default configuration
+static const ggml_et_cpu_compare_config ggml_et_cpu_compare_default_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 10
+};
+
+// CPU comparison context for a single operation
+struct ggml_et_cpu_compare_ctx {
+ ggml_backend_t cpu_backend;
+ ggml_context * ggml_ctx;
+ ggml_tensor * cpu_src0;
+ ggml_tensor * cpu_src1;
+ ggml_tensor * cpu_src2;
+ ggml_tensor * cpu_dst;
+ ggml_cgraph * cpu_graph;
+ void * cpu_src0_data;
+ void * cpu_src1_data;
+ void * cpu_src2_data;
+ void * cpu_dst_data;
+ void * et_dst_data;
+ size_t src0_size;
+ size_t src1_size;
+ size_t src2_size;
+ size_t dst_size;
+};
+
+// Phase 1: Initialize CPU comparison context and copy source buffers (call before ET kernel)
+bool ggml_et_cpu_compare_init_pre(ggml_et_cpu_compare_ctx * ctx, const ggml_tensor * node, ggml_op op);
+
+// Phase 2: Execute CPU computation and compare with ET result (call after ET kernel)
+bool ggml_et_cpu_compare_compute_and_check(ggml_et_cpu_compare_ctx * ctx,
+ const ggml_tensor * node,
+ const ggml_et_cpu_compare_config * config);
+
+// Free CPU comparison context resources
+void ggml_et_cpu_compare_free(ggml_et_cpu_compare_ctx * ctx);
--- /dev/null
+#include "ggml-et-kernels.h"
+
+#include "ggml-et-kernels-embed.hpp"
+#include "ggml-et-uberkernel-kernel-map.h"
+#include "ggml-impl.h"
+
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+
+#define ET_TRACE_DECODER_IMPL
+#include <et-trace/decoder.h>
+#include <et-trace/layout.h>
+
+static constexpr size_t GGML_ET_UBERKERNEL_PARAM_ALIGN = 64;
+
+static size_t ggml_et_align_up(size_t value, size_t alignment) {
+ return (value + alignment - 1) & ~(alignment - 1);
+}
+
+static size_t ggml_et_next_capacity(size_t current_capacity, size_t required_capacity) {
+ if (current_capacity == 0) {
+ return required_capacity;
+ }
+
+ size_t next_capacity = current_capacity;
+ while (next_capacity < required_capacity) {
+ next_capacity *= 2;
+ }
+
+ return next_capacity;
+}
+
+static ggml_backend_et_uberkernel_slot & ggml_et_uberkernel_current_slot(ggml_backend_et_uberkernel_context * uk_ctx) {
+ return uk_ctx->slots[uk_ctx->current_slot];
+}
+
+// Wait for any in-flight launch that previously used this slot to finish,
+// so the host vectors and device buffers are safe to mutate / free.
+static void ggml_et_uberkernel_slot_wait(ggml_backend_et_uberkernel_slot & slot,
+ const std::shared_ptr<rt::IRuntime> & runtime) {
+ if (!slot.has_pending || !runtime) {
+ return;
+ }
+ runtime->waitForEvent(slot.pending_event);
+ slot.has_pending = false;
+}
+
+static void ggml_et_uberkernel_reset_segment(ggml_backend_et_uberkernel_context * uk_ctx) {
+ if (!uk_ctx) {
+ return;
+ }
+
+ uk_ctx->shire_mask = 0;
+ auto & slot = ggml_et_uberkernel_current_slot(uk_ctx);
+ // Drain any prior launch on this slot before clearing its host buffers.
+ // begin_graph and abort_graph both come through here; in either case we
+ // must not yank the source memory out from under an in-flight DMA.
+ ggml_et_uberkernel_slot_wait(slot, ggml_et_runtime());
+ slot.insts.clear();
+ slot.params_blob.clear();
+}
+
+static bool ggml_et_uberkernel_ensure_slot_capacity(ggml_backend_et_uberkernel_slot & slot,
+ ggml_backend_et_device_context * dev_ctx,
+ size_t insts_size,
+ size_t params_size) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!dev_ctx || !runtime) {
+ return false;
+ }
+
+ try {
+ if (slot.device_insts == nullptr || insts_size > slot.device_insts_capacity) {
+ const size_t new_capacity = ggml_et_next_capacity(slot.device_insts_capacity, insts_size);
+ if (slot.device_insts) {
+ runtime->freeDevice(dev_ctx->rtid, slot.device_insts);
+ }
+ slot.device_insts = runtime->mallocDevice(dev_ctx->rtid, new_capacity);
+ slot.device_insts_capacity = slot.device_insts ? new_capacity : 0;
+ }
+
+ if (slot.device_params == nullptr || params_size > slot.device_params_capacity) {
+ const size_t new_capacity = ggml_et_next_capacity(slot.device_params_capacity, params_size);
+ if (slot.device_params) {
+ runtime->freeDevice(dev_ctx->rtid, slot.device_params);
+ }
+ slot.device_params = runtime->mallocDevice(dev_ctx->rtid, new_capacity);
+ slot.device_params_capacity = slot.device_params ? new_capacity : 0;
+ }
+ } catch (const std::exception & e) {
+ GGML_LOG_ERROR("ET: Failed to resize uberkernel buffers: %s\n", e.what());
+ return false;
+ }
+
+ return slot.device_insts != nullptr && slot.device_params != nullptr;
+}
+
+// Get embedded kernel data by name
+static std::vector<std::byte> ggml_et_get_embedded_kernel(const std::string & kernel_name) {
+ auto it = ggml_et_embedded_kernels.find(kernel_name);
+ if (it == ggml_et_embedded_kernels.end()) {
+ GGML_LOG_ERROR("ET: Unknown embedded kernel: %s\n", kernel_name.c_str());
+ return {};
+ }
+
+ const unsigned char * data = it->second.first;
+ uint64_t size = it->second.second;
+
+ std::vector<std::byte> buffer(size);
+ std::memcpy(buffer.data(), data, size);
+
+ return buffer;
+}
+
+// Read kernel from file (for development/override)
+static std::vector<std::byte> ggml_et_read_kernel_file(const std::string & kernel_path) {
+ std::ifstream file(kernel_path, std::ios::binary | std::ios::ate);
+ if (!file) {
+ return {};
+ }
+
+ auto size = file.tellg();
+ file.seekg(0, std::ios::beg);
+
+ std::vector<std::byte> buffer(size);
+ file.read(reinterpret_cast<char *>(buffer.data()), size);
+
+ return buffer;
+}
+
+// Load kernel from file or embedded data
+bool ggml_et_load_kernel(ggml_backend_et_device_context * dev_ctx, const std::string & kernel_name) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ GGML_LOG_ERROR("ET: Runtime not available for kernel loading\n");
+ return false;
+ }
+
+ // Check if kernel already loaded
+ if (dev_ctx->loaded_kernels.find(kernel_name) != dev_ctx->loaded_kernels.end()) {
+ GGML_LOG_DEBUG("ET: Kernel %s already loaded on device %d\n", kernel_name.c_str(), dev_ctx->devidx);
+ return true;
+ }
+
+ std::vector<std::byte> kernel_data;
+ const char * kernels_path = getenv("GGML_ET_KERNELS_PATH");
+
+ // If GGML_ET_KERNELS_PATH is set, try to load from file first
+ if (kernels_path) {
+ std::string kernel_file = std::string(kernels_path) + "/" + kernel_name + ".elf";
+ kernel_data = ggml_et_read_kernel_file(kernel_file);
+
+ if (!kernel_data.empty()) {
+ GGML_LOG_INFO("ET: Loading kernel %s from file: %s\n", kernel_name.c_str(), kernel_file.c_str());
+ } else {
+ GGML_LOG_INFO("ET: Kernel file not found: %s, falling back to embedded\n", kernel_file.c_str());
+ }
+ }
+
+ // If no file data, use embedded kernel
+ if (kernel_data.empty()) {
+ kernel_data = ggml_et_get_embedded_kernel(kernel_name);
+ if (kernel_data.empty()) {
+ GGML_LOG_ERROR("ET: Failed to get kernel data for %s\n", kernel_name.c_str());
+ return false;
+ }
+ }
+
+ try {
+ // Load kernel code using device's default stream
+ auto load_result = runtime->loadCode(dev_ctx->default_stream, kernel_data.data(), kernel_data.size());
+ runtime->waitForEvent(load_result.event_);
+
+ // Store kernel handle
+ dev_ctx->loaded_kernels[kernel_name] = load_result.kernel_;
+ return true;
+
+ } catch (const std::exception & e) {
+ GGML_LOG_ERROR("ET: Failed to load kernel %s: %s\n", kernel_name.c_str(), e.what());
+ return false;
+ }
+}
+
+static bool ggml_et_launch_kernel_internal(ggml_backend_et_device_context * dev_ctx,
+ const std::string & kernel_name,
+ void * params,
+ size_t params_size,
+ uint64_t shire_mask,
+ bool enable_print,
+ bool sync_error_check,
+ rt::EventId * out_event = nullptr) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ GGML_LOG_ERROR("ET: Runtime not available for kernel launch\n");
+ return false;
+ }
+
+ // Lazy loading: check if kernel is loaded, load if needed
+ auto kernel_it = dev_ctx->loaded_kernels.find(kernel_name);
+ if (kernel_it == dev_ctx->loaded_kernels.end()) {
+ // Kernel not loaded - load it
+ if (!ggml_et_load_kernel(dev_ctx, kernel_name)) {
+ GGML_LOG_ERROR("ET: Failed to lazy-load kernel %s\n", kernel_name.c_str());
+ return false;
+ }
+
+ // Update iterator after successful load
+ kernel_it = dev_ctx->loaded_kernels.find(kernel_name);
+ if (kernel_it == dev_ctx->loaded_kernels.end()) {
+ GGML_LOG_ERROR("ET: Kernel %s not found after loading\n", kernel_name.c_str());
+ return false;
+ }
+ }
+
+ rt::KernelId kernel_id = kernel_it->second;
+
+ try {
+ // Setup kernel launch options
+ rt::KernelLaunchOptions k_opts;
+ k_opts.setShireMask(shire_mask); // Default: all shires (0xFFFFFFFF)
+ k_opts.setBarrier(true); // Wait for completion
+ k_opts.setFlushL3(false); // No L3 flush needed
+ if (enable_print) {
+ k_opts.setUserTracing(reinterpret_cast<uint64_t>(dev_ctx->trace_buffer),
+ static_cast<uint32_t>(ET_TRACE_BUFFER_SIZE),
+ 0, // threshold
+ shire_mask, // shire mask
+ 0xFFFFFFFFFFFFFFFFULL, // threadMask - all threads
+ 0xFFFFFFFFU, // eventMask - all events
+ 0xFFFFFFFFU // filterMask - all levels
+ );
+ }
+
+ if (sync_error_check) {
+ runtime->waitForStream(dev_ctx->default_stream);
+ auto errors = runtime->retrieveStreamErrors(dev_ctx->default_stream);
+ if (!errors.empty()) {
+ GGML_LOG_ERROR("ET: Errors detected before kernel \"%s\" launch\n", kernel_name.c_str());
+ for (const auto & error : errors) {
+ GGML_LOG_ERROR("ET: Error code: %d\n", (int) error.errorCode_);
+ }
+ abort();
+ }
+ }
+
+ rt::EventId launch_event = runtime->kernelLaunch(dev_ctx->default_stream, kernel_id,
+ reinterpret_cast<std::byte *>(params), params_size, k_opts);
+ if (out_event) {
+ *out_event = launch_event;
+ }
+
+ if (enable_print) {
+ std::vector<std::byte> host_trace_buf(ET_TRACE_BUFFER_SIZE);
+ runtime->memcpyDeviceToHost(dev_ctx->default_stream, dev_ctx->trace_buffer, host_trace_buf.data(),
+ ET_TRACE_BUFFER_SIZE);
+ runtime->waitForStream(dev_ctx->default_stream);
+ const auto * trace_header = reinterpret_cast<const trace_buffer_std_header_t *>(host_trace_buf.data());
+ const trace_entry_header_t * entry = nullptr;
+ while ((entry = Trace_Decode(trace_header, entry))) {
+ if (entry->type != TRACE_TYPE_STRING) {
+ continue;
+ }
+ const auto * str_entry = reinterpret_cast<const trace_string_t *>(entry);
+ printf("[hart %d] %s", entry->hart_id, str_entry->string);
+ }
+ }
+
+ if (sync_error_check) {
+ // Already triggered. No need to retrigger
+ if (!enable_print) {
+ runtime->waitForStream(dev_ctx->default_stream);
+ }
+ auto errors = runtime->retrieveStreamErrors(dev_ctx->default_stream);
+ if (!errors.empty()) {
+ GGML_LOG_ERROR("ET: Errors detected during kernel \"%s\" execution\n", kernel_name.c_str());
+ for (const auto & error : errors) {
+ GGML_LOG_ERROR("ET: Error code: %d\n", (int) error.errorCode_);
+ }
+ abort();
+ }
+ }
+
+ return true;
+ } catch (const std::exception & e) {
+ GGML_LOG_ERROR("ET: Failed to launch kernel %s: %s\n", kernel_name.c_str(), e.what());
+ return false;
+ }
+}
+
+void ggml_et_uberkernel_begin_graph(ggml_backend_et_uberkernel_context * uk_ctx) {
+ if (!uk_ctx) {
+ return;
+ }
+
+ uk_ctx->failed = false;
+ ggml_et_uberkernel_reset_segment(uk_ctx);
+}
+
+static bool ggml_et_launch_uberkernel_segment(ggml_backend_et_device_context * dev_ctx,
+ ggml_backend_et_uberkernel_context * uk_ctx) {
+ if (!uk_ctx || !dev_ctx) {
+ return false;
+ }
+
+ auto & slot = ggml_et_uberkernel_current_slot(uk_ctx);
+ if (slot.insts.empty()) {
+ return true;
+ }
+
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ GGML_LOG_ERROR("ET: Runtime not available for uberkernel commit\n");
+ uk_ctx->failed = true;
+ return false;
+ }
+
+ const size_t insts_size = slot.insts.size() * sizeof(ggml_et_uberkernel_inst);
+ const size_t params_size = slot.params_blob.size();
+ const uint64_t shire_mask = uk_ctx->shire_mask;
+ bool ok = false;
+
+ try {
+ if (!ggml_et_uberkernel_ensure_slot_capacity(slot, dev_ctx, insts_size, params_size)) {
+ GGML_LOG_ERROR("ET: Failed to allocate uberkernel device buffers\n");
+ uk_ctx->failed = true;
+ // Drop this segment but keep the slot drained so we don't leak
+ // host vectors into the next graph.
+ slot.insts.clear();
+ slot.params_blob.clear();
+ uk_ctx->shire_mask = 0;
+ return false;
+ }
+
+ // Fire-and-forget H2D + launch on default_stream. In-stream FIFO
+ // ordering guarantees the kernel sees fully-uploaded buffers; the
+ // host source bytes (slot.insts / slot.params_blob) stay alive
+ // because we won't touch this slot again until pending_event fires.
+ runtime->memcpyHostToDevice(dev_ctx->default_stream, reinterpret_cast<const std::byte *>(slot.insts.data()),
+ slot.device_insts, insts_size, true);
+ runtime->memcpyHostToDevice(dev_ctx->default_stream, slot.params_blob.data(), slot.device_params, params_size,
+ true);
+
+ ggml_et_uberkernel_params params = {
+ static_cast<uint32_t>(slot.insts.size()),
+ static_cast<uint32_t>(sizeof(ggml_et_uberkernel_inst)),
+ reinterpret_cast<uint64_t>(slot.device_insts),
+ reinterpret_cast<uint64_t>(slot.device_params),
+ };
+
+ rt::EventId launch_event{};
+ ok = ggml_et_launch_kernel_internal(dev_ctx, "uberkernel", ¶ms, sizeof(params), shire_mask, false, false,
+ &launch_event);
+ if (ok) {
+ // The kernelLaunch above is the last thing on default_stream
+ // that touches this slot's device buffers. Recording its event
+ // lets the next reuse of this slot wait on that one event
+ // instead of the whole stream.
+ slot.pending_event = launch_event;
+ slot.has_pending = true;
+ }
+ } catch (const std::exception & e) {
+ GGML_LOG_ERROR("ET: Failed to commit uberkernel segment: %s\n", e.what());
+ }
+ uk_ctx->failed = !ok;
+
+ if (ok) {
+ uk_ctx->current_slot = (uk_ctx->current_slot + 1) % ggml_backend_et_uberkernel_context::SLOT_COUNT;
+ auto & next = ggml_et_uberkernel_current_slot(uk_ctx);
+ ggml_et_uberkernel_slot_wait(next, runtime);
+ next.insts.clear();
+ next.params_blob.clear();
+ } else {
+ slot.insts.clear();
+ slot.params_blob.clear();
+ }
+ uk_ctx->shire_mask = 0;
+ return ok;
+}
+
+void ggml_et_uberkernel_abort_graph(ggml_backend_et_uberkernel_context * uk_ctx) {
+ if (!uk_ctx) {
+ return;
+ }
+
+ uk_ctx->failed = false;
+ ggml_et_uberkernel_reset_segment(uk_ctx);
+}
+
+bool ggml_et_uberkernel_failed(const ggml_backend_et_uberkernel_context * uk_ctx) {
+ return uk_ctx && uk_ctx->failed;
+}
+
+static bool ggml_et_launch_uberkernel(ggml_backend_et_device_context * dev_ctx,
+ const std::string & kernel_name,
+ void * params,
+ size_t params_size,
+ uint64_t shire_mask,
+ bool enable_print,
+ bool sync_error_check) {
+ if (!dev_ctx) {
+ return false;
+ }
+
+ ggml_backend_et_uberkernel_context * uk_ctx = &dev_ctx->uberkernel;
+ const uint16_t uberkernel_id = ggml_et_uberkernel_kernel_id_from_name(kernel_name.c_str());
+ if (uberkernel_id == GGML_ET_UBERKERNEL_KERNEL_INVALID) {
+ if (!ggml_et_launch_uberkernel_segment(dev_ctx, uk_ctx)) {
+ return false;
+ }
+ return ggml_et_launch_kernel_internal(dev_ctx, kernel_name, params, params_size, shire_mask, enable_print,
+ sync_error_check);
+ }
+
+ auto & slot = ggml_et_uberkernel_current_slot(uk_ctx);
+ const size_t params_offset = ggml_et_align_up(slot.params_blob.size(), GGML_ET_UBERKERNEL_PARAM_ALIGN);
+ if (params_offset > slot.params_blob.size()) {
+ slot.params_blob.resize(params_offset);
+ }
+
+ const std::byte * params_bytes = reinterpret_cast<const std::byte *>(params);
+ slot.params_blob.insert(slot.params_blob.end(), params_bytes, params_bytes + params_size);
+
+ ggml_et_uberkernel_inst inst = {
+ uberkernel_id,
+ 0,
+ static_cast<uint32_t>(params_offset),
+ static_cast<uint32_t>(params_size),
+ };
+ slot.insts.push_back(inst);
+
+ if (slot.insts.size() == 1) {
+ uk_ctx->shire_mask = shire_mask;
+ }
+
+ return true;
+}
+
+bool ggml_et_uberkernel_end_graph(ggml_backend_et_device_context * dev_ctx) {
+ if (!dev_ctx || !dev_ctx->uberkernel_enabled) {
+ return true;
+ }
+
+ return ggml_et_launch_uberkernel_segment(dev_ctx, &dev_ctx->uberkernel);
+}
+
+bool ggml_et_launch_kernel(ggml_backend_et_device_context * dev_ctx,
+ const std::string & kernel_name,
+ void * params,
+ size_t params_size,
+ uint64_t shire_mask,
+ bool enable_print,
+ bool sync_error_check) {
+ if (!dev_ctx) {
+ return false;
+ }
+
+ if (!dev_ctx->uberkernel_enabled) {
+ return ggml_et_launch_kernel_internal(dev_ctx, kernel_name, params, params_size, shire_mask, enable_print,
+ sync_error_check);
+ }
+
+ return ggml_et_launch_uberkernel(dev_ctx, kernel_name, params, params_size, shire_mask, enable_print,
+ sync_error_check);
+}
+
+void ggml_et_unload_kernel(ggml_backend_et_device_context * dev_ctx, const std::string & kernel_name) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ return;
+ }
+
+ auto kernel_it = dev_ctx->loaded_kernels.find(kernel_name);
+ if (kernel_it != dev_ctx->loaded_kernels.end()) {
+ try {
+ runtime->unloadCode(kernel_it->second);
+ dev_ctx->loaded_kernels.erase(kernel_it);
+ } catch (const std::exception & e) {
+ GGML_LOG_ERROR("ET: Failed to unload kernel %s: %s\n", kernel_name.c_str(), e.what());
+ }
+ }
+}
+
+void ggml_et_unload_all_kernels(ggml_backend_et_device_context * dev_ctx) {
+ if (!dev_ctx) {
+ return;
+ }
+
+ // Make a copy of kernel names since ggml_et_unload_kernel modifies the map
+ std::vector<std::string> kernel_names;
+ kernel_names.reserve(dev_ctx->loaded_kernels.size());
+ for (const auto & kernel_pair : dev_ctx->loaded_kernels) {
+ kernel_names.push_back(kernel_pair.first);
+ }
+
+ for (const auto & kernel_name : kernel_names) {
+ ggml_et_unload_kernel(dev_ctx, kernel_name);
+ }
+}
+
+std::vector<std::pair<std::string, rt::KernelId>> ggml_et_get_loaded_kernels(ggml_backend_et_device_context * dev_ctx) {
+ std::vector<std::pair<std::string, rt::KernelId>> loaded_kernels;
+ loaded_kernels.reserve(dev_ctx->loaded_kernels.size());
+ for (const auto & kernel_pair : dev_ctx->loaded_kernels) {
+ loaded_kernels.push_back(kernel_pair);
+ }
+ return loaded_kernels;
+}
--- /dev/null
+#pragma once
+
+#include "ggml-et-common.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#define ET_TRACE_BUFFER_SIZE (1024 * 1024 * 8UL)
+
+// Load kernel from file or embedded data and store handle in device context
+// Returns true on success, false on failure
+//
+// Loading strategy:
+// - If GGML_ET_KERNELS_PATH env var is set: tries to load from ${GGML_ET_KERNELS_PATH}/${kernel_name}.elf
+// - If file not found or env var not set: falls back to embedded kernel data
+// - Returns false if kernel cannot be loaded from either source
+//
+// Kernel is loaded using the device's default stream
+bool ggml_et_load_kernel(ggml_backend_et_device_context * dev_ctx, const std::string & kernel_name);
+
+// Launch kernel with parameters on device's default stream
+// Performs lazy loading: automatically loads kernel if not already loaded
+// Kernel path: ${GGML_ET_KERNELS_PATH}/${kernel_name}.elf (default: /opt/et/ggml/kernels/)
+// Returns true on success, false on failure
+// Execution is synchronous - waits for completion
+bool ggml_et_launch_kernel(ggml_backend_et_device_context * dev_ctx,
+ const std::string & kernel_name,
+ void * params,
+ size_t params_size,
+ uint64_t shire_mask = 0xFFFFFFFF,
+ bool enable_print = false,
+ bool sync_error_check = false);
+
+void ggml_et_uberkernel_begin_graph(ggml_backend_et_uberkernel_context * uk_ctx);
+bool ggml_et_uberkernel_end_graph(ggml_backend_et_device_context * dev_ctx);
+void ggml_et_uberkernel_abort_graph(ggml_backend_et_uberkernel_context * uk_ctx);
+bool ggml_et_uberkernel_failed(const ggml_backend_et_uberkernel_context * uk_ctx);
+
+// Unload kernel from device and free resources
+// Safe to call even if kernel not loaded
+void ggml_et_unload_kernel(ggml_backend_et_device_context * dev_ctx, const std::string & kernel_name);
+
+// Unload all kernels from device context
+// Called during device cleanup
+void ggml_et_unload_all_kernels(ggml_backend_et_device_context * dev_ctx);
+
+std::vector<std::pair<std::string, rt::KernelId>> ggml_et_get_loaded_kernels(ggml_backend_et_device_context * dev_ctx);
--- /dev/null
+#include "ggml-et-memops.h"
+
+#include "ggml-et-kernels.h"
+#include "ggml-impl.h"
+
+// Kernel parameter structure for memset operation
+struct memset_params {
+ uint32_t op_type; // GGML_ET_MEMOP_MEMSET
+ uint32_t value; // Value to set (extended to uint32_t for alignment)
+ void * dst_ptr; // Destination device pointer
+ size_t size; // Number of bytes to set
+};
+
+bool ggml_et_memset(ggml_backend_et_device_context * dev_ctx, void * dst_ptr, uint8_t value, size_t size) {
+ if (!dev_ctx || !dst_ptr || size == 0) {
+ GGML_LOG_ERROR("ET: Invalid memset parameters\n");
+ return false;
+ }
+
+ // Prepare kernel parameters
+ memset_params params;
+ params.op_type = GGML_ET_MEMOP_MEMSET;
+ params.value = value;
+ params.dst_ptr = dst_ptr;
+ params.size = size;
+
+ // Launch memops kernel (will lazy-load if not already loaded)
+ bool success = ggml_et_launch_kernel(dev_ctx, "memops", ¶ms, sizeof(params));
+
+ if (!success) {
+ GGML_LOG_ERROR("ET: memset kernel launch failed\n");
+ return false;
+ }
+
+ return true;
+}
--- /dev/null
+#pragma once
+
+#include "ggml-et-common.h"
+
+#include <cstddef>
+#include <cstdint>
+
+// Memory operations using device kernel (memops.elf)
+// Single kernel handles multiple operations via operation identifier
+
+// Operation identifiers for memops kernel
+enum ggml_et_memop_type : uint32_t {
+ GGML_ET_MEMOP_MEMSET = 0,
+};
+
+// Memset operation: fill device memory with a value
+// Returns true on success, false on failure
+bool ggml_et_memset(ggml_backend_et_device_context * dev_ctx, void * dst_ptr, uint8_t value, size_t size);
--- /dev/null
+#include "ggml-et-ops.h"
+
+#include "ggml-et-cpu-compare.h"
+#include "ggml-et-kernels.h"
+#include "ggml-impl.h"
+
+#include <stdio.h>
+
+#include <cstdint>
+
+// CPU comparison configuration - can be enabled for debugging
+static ggml_et_cpu_compare_config rope_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false, // Replace ET result with CPU result
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config rms_norm_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config norm_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config l2_norm_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config group_norm_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config im2col_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config unary_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-4f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config sum_rows_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config clamp_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config mean_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config sqr_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config elmap_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config glu_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config mul_mat_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 0.01,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config mul_mat_id_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 0.01,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config softmax_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-5f,
+ /* .max_log_elements = */ 1024
+};
+
+static ggml_et_cpu_compare_config get_rows_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 2048
+};
+
+static ggml_et_cpu_compare_config pad_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config cont_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config concat_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config cumsum_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config repeat_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config ssm_conv_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config rwkv_wkv6_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-4f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config rwkv_wkv7_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-4f,
+ /* .max_log_elements = */ 4096
+};
+
+static ggml_et_cpu_compare_config set_rows_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-6f,
+ /* .max_log_elements = */ 2048
+};
+
+bool ggml_et_op_rms_norm_mul(ggml_backend_et_device_context * dev_ctx,
+ const ggml_tensor * rms_norm_node,
+ const ggml_tensor * mul_node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !rms_norm_node || !mul_node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for fused RMS_NORM_MUL operation\n");
+ return false;
+ }
+
+ if (!rms_norm_node->src[0]) {
+ GGML_LOG_ERROR("ET: Fused RMS_NORM_MUL missing required input\n");
+ return false;
+ }
+
+ // Extract weights: the MUL operand that isn't the rms_norm output
+ const ggml_tensor * weights = (mul_node->src[0] == rms_norm_node) ? mul_node->src[1] : mul_node->src[0];
+
+ if (!weights) {
+ GGML_LOG_ERROR("ET: Fused RMS_NORM_MUL missing weights tensor\n");
+ return false;
+ }
+
+ float eps;
+ memcpy(&eps, rms_norm_node->op_params, sizeof(float));
+
+ ggml_et_rms_norm_mul_params params;
+ params.src0 = *rms_norm_node->src[0]; // input to normalize
+ params.src1 = *weights; // normalization weights
+ params.dst = *mul_node; // final output
+ params.eps = eps;
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "rms_norm_mul_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END_EXT("RMS_NORM_MUL", "rms_norm_mul_f32", mul_node, "eps=%.6f", (double) eps);
+ return kernel_result;
+}
+
+bool ggml_et_op_scale(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SCALE operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: SCALE operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SCALE operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ float scale, bias;
+ memcpy(&scale, (const float *) node->op_params + 0, sizeof(float));
+ memcpy(&bias, (const float *) node->op_params + 1, sizeof(float));
+
+ ggml_et_scale_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+ params.scale = scale;
+ params.bias = bias;
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "scale_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END_EXT("SCALE", "scale_f32", node, "scale=%.6f|bias=%.6f", (double) scale, (double) bias);
+ return kernel_result;
+}
+
+bool ggml_et_op_sqr(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SQR operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: SQR operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SQR operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ ggml_et_sqr_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ params.dst = *node; // F32 output tensor
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (sqr_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_SQR)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for SQR operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "sqr_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &sqr_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for SQR operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("SQR", "sqr_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_sum_rows(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SUM_ROWS operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: SUM_ROWS operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SUM_ROWS operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ ggml_et_sum_rows_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+
+ // Phase 1: Initialize CPU comparison context
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (sum_rows_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_SUM_ROWS)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for SUM_ROWS operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "sum_rows_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &sum_rows_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for SUM_ROWS operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("SUM_ROWS", "sum_rows_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_mean(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for MEAN operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: MEAN operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: MEAN operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ ggml_et_mean_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (mean_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_MEAN)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for MEAN operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "mean_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &mean_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for MEAN operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("MEAN", "mean_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_clamp(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for CLAMP operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: CLAMP operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: CLAMP operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ ggml_et_clamp_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+ // op_params layout per ggml.c::ggml_clamp: { min, max } as floats
+ memcpy(¶ms.min_val, (const float *) node->op_params + 0, sizeof(float));
+ memcpy(¶ms.max_val, (const float *) node->op_params + 1, sizeof(float));
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (clamp_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_CLAMP)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for CLAMP operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "clamp_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &clamp_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for CLAMP operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("CLAMP", "clamp_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_unary(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for UNARY operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: UNARY operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: UNARY operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ const ggml_unary_op uop = ggml_get_unary_op(node);
+ const char * op_name = ggml_unary_op_name(uop);
+
+ ggml_et_unary_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ params.dst = *node; // F32 output tensor
+ params.unary_op = (int32_t) uop;
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (unary_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_UNARY)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for UNARY/%s operation\n", op_name);
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "unary_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &unary_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for UNARY/%s operation\n", op_name);
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("UNARY", "unary_f32", node, "op=%s", op_name);
+ return kernel_result;
+}
+
+bool ggml_et_op_mul(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ // Delegate to generic element map operation
+ return ggml_et_op_elmap(dev_ctx, node);
+}
+
+bool ggml_et_op_add(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ // Delegate to generic element map operation
+ return ggml_et_op_elmap(dev_ctx, node);
+}
+
+bool ggml_et_op_sub(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ // Delegate to generic element map operation
+ return ggml_et_op_elmap(dev_ctx, node);
+}
+
+bool ggml_et_op_elmap(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for element map operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: Element map operation missing required inputs\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32 || node->src[1]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: Element map operation with unsupported types: dst=%s src0=%s src1=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ const char * op_name = ggml_op_name(node->op);
+
+ ggml_et_elmap_params params;
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.dst = *node; // F32 output tensor (op type stored in dst.op)
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (elmap_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, node->op)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for %s operation\n", op_name);
+ }
+ }
+
+ // fprintf(stderr, "ET: el_map s0 [%ld, %ld, %ld, %ld] s1 [%ld, %ld, %ld, %ld]\n",
+ // node->src[0]->ne[0], node->src[0]->ne[1], node->src[0]->ne[2], node->src[0]->ne[3],
+ // node->src[1]->ne[0], node->src[1]->ne[1], node->src[1]->ne[2], node->src[1]->ne[3]);
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "el_map_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &elmap_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for %s operation\n", op_name);
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END(op_name, "el_map_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_glu(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ // Validate inputs
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for GLU operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: GLU operation missing required input\n");
+ return false;
+ }
+
+ const bool is_split_mode = node->src[1] != nullptr;
+
+ // Only support F32 (as validated by supports_op)
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32 ||
+ (is_split_mode && node->src[1]->type != GGML_TYPE_F32)) {
+ return false;
+ }
+
+ // Extract GLU operation parameters from op_params
+ int32_t glu_op_type = ggml_get_op_params_i32(node, 0); // GLU variant (REGLU, GEGLU, SWIGLU, etc.)
+ int32_t swapped = ggml_get_op_params_i32(node, 1); // Whether gate/value are swapped
+
+ // Supported variants
+ switch (glu_op_type) {
+ case GGML_GLU_OP_REGLU:
+ case GGML_GLU_OP_GEGLU:
+ case GGML_GLU_OP_SWIGLU:
+ case GGML_GLU_OP_SWIGLU_OAI:
+ case GGML_GLU_OP_GEGLU_ERF:
+ case GGML_GLU_OP_GEGLU_QUICK:
+ break;
+ default:
+ GGML_LOG_ERROR("ET: GLU operation with unsupported variant: %s\n",
+ ggml_glu_op_name((ggml_glu_op) glu_op_type));
+ return false;
+ }
+
+ // Get GLU operation name for logging
+ const char * glu_op_name = ggml_glu_op_name((ggml_glu_op) glu_op_type);
+
+ // Pack parameters. Single-tensor mode is encoded by zeroing src1.
+ ggml_et_glu_params params = {};
+ params.src0 = *node->src[0];
+ if (is_split_mode) {
+ params.src1 = *node->src[1];
+ }
+ params.dst = *node;
+ params.glu_op_type = glu_op_type;
+ params.swapped = swapped;
+ params.alpha = 0.0f;
+ params.limit = 0.0f;
+ if (glu_op_type == GGML_GLU_OP_SWIGLU_OAI) {
+ params.alpha = ggml_get_op_params_f32(node, 2);
+ params.limit = ggml_get_op_params_f32(node, 3);
+ }
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (glu_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_GLU)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for %s operation\n", glu_op_name);
+ }
+ }
+
+ // Launch ET kernel
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "glu_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &glu_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for %s operation\n", glu_op_name);
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("GLU", "glu_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_mul_mat(ggml_backend_et_device_context * dev_ctx,
+ const ggml_tensor * node,
+ const ggml_tensor * add_node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for MUL_MAT operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: MUL_MAT operation missing required inputs\n");
+ return false;
+ }
+
+ // Fused MM+ADD: when add_node is non-NULL the caller has already validated
+ // (Q8_0 weights, F32 acts, exact-shape ADD with stride parity to dst) via
+ // ggml_et_can_fuse({MUL_MAT, ADD}). The kernel writes dst = mm + bias and
+ // the ADD's output replaces MM's as the actual dst.
+ const ggml_tensor * fused_dst = add_node ? add_node : node;
+ const ggml_tensor * bias_tensor = nullptr;
+ if (add_node) {
+ bias_tensor = (add_node->src[0] == node) ? add_node->src[1] : add_node->src[0];
+ }
+
+ const char * kernel_name;
+ const char * src0_type_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_Q4_0 && node->src[1]->type == GGML_TYPE_F32 &&
+ node->src[1]->ne[1] >= 53 && // N >= 53
+ node->src[0]->ne[1] % 16 == 0 && // M % TILE_M
+ node->src[0]->ne[0] % 32 == 0) { // K % BLOCK_K (Q4_0 block)
+
+ // Matrix engine for N >= 53; partial N (via n_cur-1) and errata padding are handled in-kernel.
+ kernel_name = "mul_mat_Q4_0_matrix_engine";
+ src0_type_name = "Q4_0";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_Q4_0 &&
+ node->src[1]->type == GGML_TYPE_F32) {
+ kernel_name = "mul_mat_Q4_0"; // N < 53, or M % 16 != 0 or K % 32 != 0
+ src0_type_name = "Q4_0";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_Q8_0 &&
+ node->src[1]->type == GGML_TYPE_F32) {
+ kernel_name = "mul_mat_Q8_0";
+ src0_type_name = "Q8_0";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F16 &&
+ node->src[1]->type == GGML_TYPE_F16 && node->ne[0] % 16 == 0 && node->src[0]->ne[0] % 16 == 0 &&
+ node->src[0]->ne[1] % 16 == 0 && node->src[1]->ne[0] != 1) {
+ kernel_name = "mul_mat_f16_matrix_engine";
+ src0_type_name = "F16";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F16 &&
+ (node->src[1]->type == GGML_TYPE_F16 || node->src[1]->type == GGML_TYPE_F32)) {
+ kernel_name = "mul_mat_f16";
+ src0_type_name = "F16";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 &&
+ node->src[1]->type == GGML_TYPE_F32 && node->ne[0] % 16 == 0 && node->src[0]->ne[0] % 16 == 0 &&
+ node->src[0]->ne[1] % 16 == 0 && node->src[1]->ne[0] != 1) { // GEMV is faster with the generic path
+
+ kernel_name = "mul_mat_f32_matrix_engine";
+ src0_type_name = "F32";
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 &&
+ (node->src[1]->type == GGML_TYPE_F16 || node->src[1]->type == GGML_TYPE_F32)) {
+ kernel_name = "mul_mat_f32";
+ src0_type_name = "F32";
+ } else {
+ GGML_LOG_ERROR("ET: MUL_MAT operation with unsupported types: dst=%s src0=%s src1=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ ggml_et_binary_params params;
+ params.src0 = *node->src[0]; // weight matrix
+ params.src1 = *node->src[1]; // activation matrix
+ params.dst = *fused_dst; // output (= add_node when fused, else node)
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (mul_mat_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, fused_dst, GGML_OP_MUL_MAT)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for MUL_MAT operation\n");
+ }
+ }
+
+ bool kernel_result;
+ if (node->src[0]->type == GGML_TYPE_Q8_0) {
+ // Q8_0 kernel always takes the extended struct. bias.data is non-NULL
+ // only on the fused path; otherwise the kernel skips the add entirely.
+ ggml_et_mm_q8_params q8_params = {};
+ q8_params.src0 = params.src0;
+ q8_params.src1 = params.src1;
+ q8_params.dst = params.dst;
+ if (bias_tensor) {
+ q8_params.bias = *bias_tensor;
+ }
+ kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, &q8_params, sizeof(q8_params), 0xFFFFFFFF);
+ } else {
+ // Non-Q8 MM kernels don't yet support fused-add; the graph fuse check
+ // already rejects non-Q8 pairs, so add_node is always nullptr here.
+ kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+ }
+
+ // printf("Tensor error:");
+ // if (params.src0.data != NULL)
+ // {
+ // printf("Ptr OK\n");
+ // printf("node->data ptr = %p\n", node->data);
+ // // if (once < 100){
+ // // // uint64_t * host_data = (uint64_t *) node->data;
+ // // // printf("Tensor error: %lu\n", host_data[0]);
+
+ // // // printf("Tensor error:");
+ // // once++;
+ // // }
+ // }
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, fused_dst, &mul_mat_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for MUL_MAT operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ {
+ // Calculate actual FLOPs including batch/sequence dimensions
+ // dst shape: [M, N, ne2, ne3] where M=ne[1], N=ne[0]
+ int64_t m = node->ne[1];
+ int64_t n = node->ne[0];
+ int64_t k = node->src[0]->ne[0];
+ int64_t ne2 = node->ne[2];
+ int64_t ne3 = node->ne[3];
+
+ // Total FLOPs = (batch_size) * M * N * (2*K - 1)
+ // Each MxN matrix-matrix multiply does M*N*(2*K-1) FLOPs
+ // Broadcasting is handled by repeating computation, so count actual operations
+ int64_t batch_size = ne2 * ne3;
+ int64_t total_flops = batch_size * m * n * (2 * k - 1);
+
+ char kernel_variant[64];
+ snprintf(kernel_variant, sizeof(kernel_variant), "%s_%sx%s", kernel_name, src0_type_name,
+ ggml_type_name(node->src[1]->type));
+ ET_PERF_END_EXT("MUL_MAT", kernel_variant, node, "flops=%" PRId64, total_flops);
+ }
+ return kernel_result;
+}
+
+bool ggml_et_op_mul_mat_id(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for MUL_MAT_ID operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1] || !node->src[2]) {
+ GGML_LOG_ERROR("ET: MUL_MAT_ID operation missing required inputs\n");
+ return false;
+ }
+
+ const char * kernel_name;
+ const char * src0_type_name;
+
+ // Support Q8_0/Q4_0/F16/F32 x F32 -> F32 matrix multiplication with expert selection
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_Q8_0 && node->src[1]->type == GGML_TYPE_F32 &&
+ node->src[2]->type == GGML_TYPE_I32) {
+ kernel_name = "mul_mat_id_Q8_0";
+ src0_type_name = "Q8_0";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_Q4_0 &&
+ node->src[1]->type == GGML_TYPE_F32 && node->src[2]->type == GGML_TYPE_I32) {
+ kernel_name = "mul_mat_id_Q4_0";
+ src0_type_name = "Q4_0";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F16 &&
+ node->src[1]->type == GGML_TYPE_F32 && node->src[2]->type == GGML_TYPE_I32) {
+ kernel_name = "mul_mat_id_f32";
+ src0_type_name = "F16";
+
+ } else if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 &&
+ node->src[1]->type == GGML_TYPE_F32 && node->src[2]->type == GGML_TYPE_I32) {
+ kernel_name = "mul_mat_id_f32";
+ src0_type_name = "F32";
+
+ } else {
+ GGML_LOG_ERROR("ET: MUL_MAT_ID operation with unsupported types: dst=%s src0=%s src1=%s src2=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type), ggml_type_name(node->src[2]->type));
+ return false;
+ }
+
+ // Pack parameters - copy full tensor structures
+ ggml_et_mul_mat_id_params params;
+ params.src0 = *node->src[0]; // Expert weight matrices (Q8_0/F16/F32)
+ params.src1 = *node->src[1]; // Activation matrix (F32)
+ params.src2 = *node->src[2]; // Expert indices (I32)
+ params.dst = *node; // Output matrix (F32)
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (mul_mat_id_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_MUL_MAT_ID)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for MUL_MAT_ID operation\n");
+ }
+ }
+
+ // Launch ET kernel
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &mul_mat_id_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for MUL_MAT_ID operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ // Calculate FLOPs (approximate - similar to MUL_MAT but with expert routing overhead)
+ // Each expert computation is similar to a MUL_MAT, but we only compute for selected experts
+ int64_t K = node->src[0]->ne[0];
+ int64_t M = node->src[0]->ne[1];
+ int64_t n_expert_used = node->src[2]->ne[0];
+ int64_t batch = node->src[2]->ne[1];
+
+ int64_t total_flops = batch * n_expert_used * M * (2 * K - 1);
+
+ char kernel_variant[64];
+ snprintf(kernel_variant, sizeof(kernel_variant), "%s_%sx%s", kernel_name, src0_type_name,
+ ggml_type_name(node->src[1]->type));
+ ET_PERF_END_EXT("MUL_MAT_ID", kernel_variant, node, "flops=%" PRId64 "|n_expert=%lld|n_expert_used=%lld",
+ total_flops, (long long) node->src[0]->ne[2], (long long) n_expert_used);
+
+ return kernel_result;
+}
+
+bool ggml_et_op_rope(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for ROPE operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: ROPE operation missing required inputs\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_I32) {
+ kernel_name = "rope_f32";
+ } else {
+ return false;
+ }
+
+ // Pack parameters - copy full tensor structures and op_params
+ ggml_et_rope_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ params.src1 = *node->src[1]; // I32 position tensor
+ if (node->src[2]) {
+ params.src2 = *node->src[2]; // F32 frequency factors (optional)
+ } else {
+ memset(¶ms.src2, 0, sizeof(params.src2)); // Zero if not provided
+ }
+ params.dst = *node; // F32 output tensor
+
+ params.rope_params.n_past = ((const int32_t *) node->op_params)[0];
+ params.rope_params.n_dims = ((const int32_t *) node->op_params)[1];
+ params.rope_params.mode = ((const int32_t *) node->op_params)[2];
+ params.rope_params.n_ctx = ((const int32_t *) node->op_params)[3];
+ params.rope_params.n_ctx_orig = ((const int32_t *) node->op_params)[4];
+ memcpy(¶ms.rope_params.freq_base, (const int32_t *) node->op_params + 5, sizeof(float));
+ memcpy(¶ms.rope_params.freq_scale, (const int32_t *) node->op_params + 6, sizeof(float));
+ memcpy(¶ms.rope_params.ext_factor, (const int32_t *) node->op_params + 7, sizeof(float));
+ memcpy(¶ms.rope_params.attn_factor, (const int32_t *) node->op_params + 8, sizeof(float));
+ memcpy(¶ms.rope_params.beta_fast, (const int32_t *) node->op_params + 9, sizeof(float));
+ memcpy(¶ms.rope_params.beta_slow, (const int32_t *) node->op_params + 10, sizeof(float));
+ if (params.rope_params.mode & GGML_ROPE_TYPE_MROPE) {
+ memcpy(params.rope_params.sections, (const int32_t *) node->op_params + 11, sizeof(int32_t) * 4);
+ } else {
+ memset(params.rope_params.sections, 0, sizeof(params.rope_params.sections));
+ }
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (rope_cpu_compare_config.enabled) {
+ GGML_LOG_DEBUG("ET: Initializing CPU comparison for ROPE operation\n");
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_ROPE)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for ROPE operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &rope_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for ROPE operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("ROPE", kernel_name, node, "mode=0x%x|n_dims=%d|freq_base=%.2f|freq_scale=%.2f",
+ params.rope_params.mode, params.rope_params.n_dims, (double) params.rope_params.freq_base,
+ (double) params.rope_params.freq_scale);
+ return kernel_result;
+}
+
+bool ggml_et_op_rms_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for RMS_NORM operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: RMS_NORM operation missing required input\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32) {
+ kernel_name = "rms_norm_f32";
+
+ } else {
+ GGML_LOG_ERROR("ET: RMS_NORM operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ float eps;
+ memcpy(&eps, node->op_params, sizeof(float));
+
+ ggml_et_rms_norm_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ params.dst = *node; // F32 output tensor
+ params.eps = eps; // Epsilon parameter for numerical stability
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (rms_norm_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_RMS_NORM)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for RMS_NORM operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &rms_norm_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for RMS_NORM operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("RMS_NORM", kernel_name, node, "eps=%.6f", (double) eps);
+ return kernel_result;
+}
+
+bool ggml_et_op_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for NORM operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: NORM operation missing required input\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32) {
+ kernel_name = "norm_f32";
+
+ } else {
+ GGML_LOG_ERROR("ET: NORM operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ float eps;
+ memcpy(&eps, node->op_params, sizeof(float));
+
+ ggml_et_norm_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ params.dst = *node; // F32 output tensor
+ params.eps = eps; // Epsilon parameter for numerical stability
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (norm_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_NORM)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for NORM operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &norm_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for NORM operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("NORM", kernel_name, node, "eps=%.6f", (double) eps);
+ return kernel_result;
+}
+
+bool ggml_et_op_l2_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for L2_NORM operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: L2_NORM operation missing required input\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32) {
+ kernel_name = "l2_norm_f32";
+
+ } else {
+ GGML_LOG_ERROR("ET: L2_NORM operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ float eps;
+ memcpy(&eps, node->op_params, sizeof(float));
+
+ ggml_et_l2_norm_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ params.dst = *node; // F32 output tensor
+ params.eps = eps; // Epsilon parameter for numerical stability
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (l2_norm_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_L2_NORM)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for L2_NORM operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &l2_norm_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for L2_NORM operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("L2_NORM", kernel_name, node, "eps=%.6f", (double) eps);
+ return kernel_result;
+}
+
+bool ggml_et_op_group_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for GROUP_NORM operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: GROUP_NORM operation missing required input\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: GROUP_NORM operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ const int32_t n_groups = ggml_get_op_params_i32(node, 0);
+ float eps;
+ memcpy(&eps, (const float *) node->op_params + 1, sizeof(float));
+
+ ggml_et_group_norm_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+ params.n_groups = n_groups;
+ params.eps = eps;
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (group_norm_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_GROUP_NORM)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for GROUP_NORM operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "group_norm_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &group_norm_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for GROUP_NORM operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("GROUP_NORM", "group_norm_f32", node, "eps=%.6f|n_groups=%d", (double) eps, n_groups);
+ return kernel_result;
+}
+
+bool ggml_et_op_im2col(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for IM2COL operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: IM2COL operation missing required inputs\n");
+ return false;
+ }
+
+ const bool supported_types =
+ (node->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_F32) ||
+ (node->type == GGML_TYPE_F16 && (node->src[1]->type == GGML_TYPE_F16 || node->src[1]->type == GGML_TYPE_F32));
+
+ if (!supported_types) {
+ GGML_LOG_ERROR("ET: IM2COL operation with unsupported types: dst=%s src1=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ ggml_et_im2col_params params;
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.dst = *node;
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (im2col_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_IM2COL)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for IM2COL operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "im2col", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &im2col_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for IM2COL operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("IM2COL", "im2col", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_conv_2d(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ return false;
+ }
+ if (!node->src[0] || !node->src[1]) {
+ return false;
+ }
+ if (!node->data || !node->src[0]->data || !node->src[1]->data) {
+ return false;
+ }
+
+ // Kernel constraints (mirror supports_op; recheck here as a guard).
+ const ggml_tensor * flt = node->src[0]; // [Kw, Kh, Cin, Cout]
+ const ggml_tensor * in = node->src[1]; // [W, H, Cin, N]
+ if (node->type != GGML_TYPE_F32 || flt->type != GGML_TYPE_F32 || in->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ const int32_t s0 = ggml_get_op_params_i32(node, 0);
+ const int32_t s1 = ggml_get_op_params_i32(node, 1);
+ const int32_t p0 = ggml_get_op_params_i32(node, 2);
+ const int32_t p1 = ggml_get_op_params_i32(node, 3);
+ const int32_t d0 = ggml_get_op_params_i32(node, 4);
+ const int32_t d1 = ggml_get_op_params_i32(node, 5);
+
+ if (s0 < 1 || s1 < 1) {
+ return false;
+ }
+ if (d0 != 1 || d1 != 1) {
+ return false;
+ }
+ if (flt->ne[2] % 16 != 0 || flt->ne[3] % 16 != 0) {
+ return false;
+ }
+ if (in->ne[3] != 1) {
+ return false;
+ }
+ if (node->ne[0] <= 0) {
+ return false; // OW > 0 (any width OK; staging path handles non-16)
+ }
+ (void) p0;
+ (void) p1;
+
+ ggml_et_binary_params params;
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.dst = *node;
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "conv_2d_f32_me", ¶ms, sizeof(params), 0xFFFFFFFFu);
+
+ ET_PERF_END("CONV_2D", "conv_2d_f32_me", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_softmax(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SOFTMAX operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation missing required input\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32) {
+ kernel_name = "softmax_f32";
+
+ } else {
+ GGML_LOG_ERROR("ET: SOFTMAX operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ // Validate contiguity requirements
+ if (!ggml_is_contiguous(node)) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation requires contiguous destination tensor\n");
+ return false;
+ }
+
+ if (!ggml_is_contiguous(node->src[0])) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation requires contiguous source tensor\n");
+ return false;
+ }
+
+ // Check optional mask tensor
+ if (node->src[1]) {
+ if (node->src[1]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation with unsupported mask type: %s (F32 required)\n",
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+ if (!ggml_is_contiguous(node->src[1])) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation requires contiguous mask tensor\n");
+ return false;
+ }
+ }
+
+ // Check optional sinks tensor
+ if (node->src[2]) {
+ if (node->src[2]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation with unsupported sinks type: %s (F32 required)\n",
+ ggml_type_name(node->src[2]->type));
+ return false;
+ }
+ if (!ggml_is_contiguous(node->src[2])) {
+ GGML_LOG_ERROR("ET: SOFTMAX operation requires contiguous sinks tensor\n");
+ return false;
+ }
+ }
+
+ // Extract scale and max_bias from op_params
+ float scale = 1.0f;
+ float max_bias = 0.0f;
+ if (node->op_params) {
+ memcpy(&scale, (const float *) node->op_params + 0, sizeof(float));
+ memcpy(&max_bias, (const float *) node->op_params + 1, sizeof(float));
+ }
+
+ ggml_et_softmax_params params;
+ params.src0 = *node->src[0]; // F32 input tensor
+ if (node->src[1]) {
+ params.src1 = *node->src[1]; // F32 mask tensor
+ } else {
+ memset(¶ms.src1, 0, sizeof(params.src1)); // Zero if no mask
+ }
+ if (node->src[2]) {
+ params.src2 = *node->src[2]; // F32 sinks tensor
+ } else {
+ memset(¶ms.src2, 0, sizeof(params.src2)); // Zero if no sinks
+ }
+ params.dst = *node; // F32 output tensor
+ params.scale = scale; // Scale factor
+ params.max_bias = max_bias; // ALiBi bias
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (softmax_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_SOFT_MAX)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for SOFTMAX operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &softmax_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for SOFTMAX operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("SOFTMAX", kernel_name, node, "scale=%.6f|max_bias=%.6f|has_mask=%s", (double) scale,
+ (double) max_bias, node->src[1] ? "yes" : "no");
+ return kernel_result;
+}
+
+bool ggml_et_op_flash_attn_ext(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for FLASH_ATTN_EXT operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1] || !node->src[2]) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT operation missing required inputs\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT requires F32 Q and dst, got dst=%s q=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ // K and V can be F16 or F32
+ if ((node->src[1]->type != GGML_TYPE_F32 && node->src[1]->type != GGML_TYPE_F16) ||
+ (node->src[2]->type != GGML_TYPE_F32 && node->src[2]->type != GGML_TYPE_F16)) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT K/V must be F16 or F32, got k=%s v=%s\n", ggml_type_name(node->src[1]->type),
+ ggml_type_name(node->src[2]->type));
+ return false;
+ }
+
+ if (node->src[4] != nullptr) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel does not support sinks\n");
+ return false;
+ }
+
+ // Mask is optional; if present must be F16 or F32
+ if (node->src[3] != nullptr && node->src[3]->type != GGML_TYPE_F32 && node->src[3]->type != GGML_TYPE_F16) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT mask must be F16 or F32, got %s\n", ggml_type_name(node->src[3]->type));
+ return false;
+ }
+
+ // Q and dst must be row-contiguous F32
+ if (!ggml_is_contiguous_rows(node) || !ggml_is_contiguous_rows(node->src[0])) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT requires row-contiguous Q and dst\n");
+ return false;
+ }
+
+ if (node->nb[0] != sizeof(float) || node->src[0]->nb[0] != sizeof(float)) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT requires contiguous F32 rows for Q and dst\n");
+ return false;
+ }
+
+ // K/V must have element-sized stride in dim 0
+ const size_t k_elem = node->src[1]->type == GGML_TYPE_F16 ? 2 : 4;
+ const size_t v_elem = node->src[2]->type == GGML_TYPE_F16 ? 2 : 4;
+ if (node->src[1]->nb[0] != k_elem || node->src[2]->nb[0] != v_elem) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT K/V must have element-sized stride in dim 0\n");
+ return false;
+ }
+
+ float scale = 1.0f;
+ float max_bias = 0.0f;
+ float logit_softcap = 0.0f;
+ memcpy(&scale, (const float *) node->op_params + 0, sizeof(scale));
+ memcpy(&max_bias, (const float *) node->op_params + 1, sizeof(max_bias));
+ memcpy(&logit_softcap, (const float *) node->op_params + 2, sizeof(logit_softcap));
+
+ if (max_bias != 0.0f || logit_softcap != 0.0f) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel does not support max_bias or logit_softcap\n");
+ return false;
+ }
+
+ const ggml_prec prec = ggml_flash_attn_ext_get_prec(node);
+ if (prec != GGML_PREC_F32 && prec != GGML_PREC_DEFAULT) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel only supports F32 precision\n");
+ return false;
+ }
+
+ // dk must match between Q and K; dv must match between V and dst
+ if (node->src[0]->ne[0] != node->src[1]->ne[0]) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT dk mismatch: Q=%lld K=%lld\n", (long long) node->src[0]->ne[0],
+ (long long) node->src[1]->ne[0]);
+ return false;
+ }
+
+ if (node->src[2]->ne[0] != node->ne[0]) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT dv mismatch: V=%lld dst=%lld\n", (long long) node->src[2]->ne[0],
+ (long long) node->ne[0]);
+ return false;
+ }
+
+ if (node->src[2]->ne[0] > 512) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT dv=%lld exceeds maximum 512\n", (long long) node->src[2]->ne[0]);
+ return false;
+ }
+
+ if (node->src[0]->ne[0] > 512) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT dk=%lld exceeds maximum 512\n", (long long) node->src[0]->ne[0]);
+ return false;
+ }
+
+ // GQA: n_head_q must be a multiple of n_head_kv
+ const int64_t nhq = node->src[0]->ne[2];
+ const int64_t nhk = node->src[1]->ne[2];
+ if (nhq % nhk != 0) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT n_head_q (%lld) not divisible by n_head_kv (%lld)\n", (long long) nhq,
+ (long long) nhk);
+ return false;
+ }
+
+ // K and V must have matching sequence length, heads, and batch dims
+ if (node->src[1]->ne[1] != node->src[2]->ne[1] || node->src[1]->ne[2] != node->src[2]->ne[2] ||
+ node->src[1]->ne[3] != node->src[2]->ne[3]) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT K/V shape mismatch\n");
+ return false;
+ }
+
+ // dst layout checks: [dv, nhq, nq, no]
+ if (node->src[0]->ne[1] != node->ne[2] || node->src[0]->ne[2] != node->ne[1] ||
+ node->src[0]->ne[3] != node->ne[3]) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT dst shape mismatch\n");
+ return false;
+ }
+
+ // Batch dims: Q batch must match K batch
+ if (node->src[0]->ne[3] != node->src[1]->ne[3]) {
+ GGML_LOG_ERROR("ET: FLASH_ATTN_EXT batch dimension mismatch\n");
+ return false;
+ }
+
+ ggml_et_flash_attn_ext_params params;
+ memset(¶ms, 0, sizeof(params));
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.src2 = *node->src[2];
+ if (node->src[3] != nullptr) {
+ params.mask = *node->src[3];
+ params.has_mask = 1;
+ }
+ params.dst = *node;
+ params.scale = scale;
+
+ // Use matrix engine kernel when K/V are F16 and dk is a multiple of 32
+ const char * kernel_name;
+ if (node->src[1]->type == GGML_TYPE_F16 && node->src[2]->type == GGML_TYPE_F16 && (node->src[0]->ne[0] % 32) == 0) {
+ kernel_name = "flash_attn_ext_f16_me";
+ } else {
+ kernel_name = "flash_attn_ext_f32";
+ }
+
+ const bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END_EXT("FLASH_ATTN_EXT", kernel_name, node, "scale=%.6f", (double) scale);
+ return kernel_result;
+}
+
+bool ggml_et_op_get_rows(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for GET_ROWS operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: GET_ROWS operation missing required inputs\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_I32 &&
+ (node->src[0]->type == GGML_TYPE_F32 || node->src[0]->type == GGML_TYPE_F16 ||
+ node->src[0]->type == GGML_TYPE_Q4_0 || node->src[0]->type == GGML_TYPE_Q8_0 ||
+ node->src[0]->type == GGML_TYPE_Q4_K)) {
+ kernel_name = "get_rows_f32";
+
+ } else {
+ GGML_LOG_ERROR("ET: GET_ROWS operation with unsupported types: dst=%s src0=%s src1=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ // Validate contiguity requirements
+ if (!ggml_is_contiguous(node)) {
+ GGML_LOG_ERROR("ET: GET_ROWS operation requires contiguous destination tensor\n");
+ return false;
+ }
+
+ if (!ggml_is_contiguous(node->src[0])) {
+ GGML_LOG_ERROR("ET: GET_ROWS operation requires contiguous data tensor\n");
+ return false;
+ }
+
+ if (!ggml_is_contiguous(node->src[1])) {
+ GGML_LOG_ERROR("ET: GET_ROWS operation requires contiguous indices tensor\n");
+ return false;
+ }
+
+ // Validate dimension constraints from ggml implementation
+ if (node->src[0]->ne[2] != node->src[1]->ne[1] || node->src[1]->ne[3] != 1) {
+ GGML_LOG_ERROR(
+ "ET: GET_ROWS operation dimension constraint failed: src0.ne[2]=%lld != src1.ne[1]=%lld or src1.ne[3]=%lld "
+ "!= 1\n",
+ (long long) node->src[0]->ne[2], (long long) node->src[1]->ne[1], (long long) node->src[1]->ne[3]);
+ return false;
+ }
+
+ ggml_et_get_rows_params params;
+ params.src0 = *node->src[0]; // Data tensor (F32 or Q8_0)
+ params.src1 = *node->src[1]; // Indices tensor (I32)
+ params.dst = *node; // Output tensor (F32)
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (get_rows_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_GET_ROWS)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for GET_ROWS operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &get_rows_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for GET_ROWS operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("GET_ROWS", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_cont(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ // Validate source tensor exists
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: CONT operation missing source tensor\n");
+ return false;
+ }
+
+ // Validate types match (input and output must be same type)
+ if (node->type != node->src[0]->type) {
+ GGML_LOG_ERROR("ET: CONT operation type mismatch: src=%s dst=%s\n", ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->type));
+ return false;
+ }
+
+ // Validate supported types
+ if (node->type != GGML_TYPE_F32 && node->type != GGML_TYPE_F16) {
+ GGML_LOG_ERROR("ET: CONT operation unsupported type: %s (only F32 and F16 supported)\n",
+ ggml_type_name(node->type));
+ return false;
+ }
+
+ // Validate contiguity - output must be contiguous, input can be non-contiguous
+ if (!ggml_is_contiguous(node)) {
+ GGML_LOG_ERROR("ET: CONT operation requires contiguous output tensor\n");
+ return false;
+ }
+
+ // Select kernel based on type
+ const char * kernel_name;
+ if (node->type == GGML_TYPE_F32) {
+ kernel_name = "cont_f32";
+ } else if (node->type == GGML_TYPE_F16) {
+ kernel_name = "cont_f16";
+ } else {
+ GGML_LOG_ERROR("ET: CONT operation with unsupported type: %s\n", ggml_type_name(node->type));
+ return false;
+ }
+
+ ggml_et_cont_params params;
+ params.src0 = *node->src[0]; // Input tensor (potentially non-contiguous)
+ params.dst = *node; // Output tensor (contiguous)
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (cont_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_CONT)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for CONT operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &cont_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for CONT operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("CONT", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_cumsum(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node || !node->src[0]) {
+ GGML_LOG_ERROR("ET: Invalid parameters for CUMSUM operation\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: CUMSUM operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ const char * kernel_name = "cumsum_f32";
+
+ ggml_et_cumsum_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (cumsum_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_CUMSUM)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for CUMSUM operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &cumsum_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for CUMSUM operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("CUMSUM", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_cpy(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ // CPY copies data from src[0] into the layout of dst (which matches src[1])
+ // For same-type with contiguous dst, this is identical to CONT
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: CPY operation missing source tensor\n");
+ return false;
+ }
+
+ // Scalar / zero-element special path: if any dimension is 0, nothing to copy
+ const int64_t nelements = node->ne[0] * node->ne[1] * node->ne[2] * node->ne[3];
+ if (nelements == 0) {
+ GGML_LOG_DEBUG("ET: CPY no-op (zero elements): ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 "]\n",
+ node->ne[0], node->ne[1], node->ne[2], node->ne[3]);
+ ET_PERF_END("CPY", "noop", node);
+ return true;
+ }
+
+ // Only F32 and F16 supported for dst
+ if (node->type != GGML_TYPE_F32 && node->type != GGML_TYPE_F16) {
+ GGML_LOG_ERROR("ET: CPY unsupported dst type: %s\n", ggml_type_name(node->type));
+ return false;
+ }
+
+ // Select kernel based on src/dst type combination
+ const char * kernel_name;
+ if (node->src[0]->type == GGML_TYPE_F32 && node->type == GGML_TYPE_F32) {
+ kernel_name = "cont_f32";
+ } else if (node->src[0]->type == GGML_TYPE_F16 && node->type == GGML_TYPE_F16) {
+ kernel_name = "cont_f16";
+ } else if (node->src[0]->type == GGML_TYPE_F32 && node->type == GGML_TYPE_F16) {
+ kernel_name = "cpy_f32_f16";
+ } else {
+ GGML_LOG_ERROR("ET: CPY unsupported type combination: src=%s dst=%s\n", ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->type));
+ return false;
+ }
+
+ ggml_et_cont_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+
+ // CPU comparison for debugging
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (cont_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_CPY)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for CPY operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &cont_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for CPY operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("CPY", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_concat(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for CONCAT operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: CONCAT operation missing required inputs\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_F32) {
+ kernel_name = "concat_f32";
+
+ } else {
+ GGML_LOG_ERROR("ET: CONCAT operation with unsupported types: dst=%s src0=%s src1=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ int32_t dim;
+ memcpy(&dim, node->op_params, sizeof(int32_t));
+
+ ggml_et_concat_params params;
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.dst = *node;
+ params.dim = dim;
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (concat_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_CONCAT)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for CONCAT operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &concat_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for CONCAT operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("CONCAT", kernel_name, node, "dim=%d", dim);
+ return kernel_result;
+}
+
+bool ggml_et_op_repeat(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for REPEAT operation\n");
+ return false;
+ }
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: REPEAT operation missing required input\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ if (node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32) {
+ // No-op REPEAT (every repeat factor is 1): the output is just a copy
+ // of the input. Route to cont_f32, whose contiguous fast path handles
+ // arbitrary sizes (including those rejected by repeat_f32's gate,
+ // e.g. ne[0]=1).
+ if (ggml_are_same_shape(node->src[0], node)) {
+ kernel_name = "cont_f32";
+ } else {
+ kernel_name = "repeat_f32";
+ }
+
+ } else {
+ GGML_LOG_ERROR("ET: REPEAT operation with unsupported types: dst=%s src0=%s\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type));
+ return false;
+ }
+
+ // ggml_et_cont_params and ggml_et_repeat_params have identical layouts
+ // (just src0 + dst), so the same payload works for either kernel.
+ ggml_et_repeat_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (repeat_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_REPEAT)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for REPEAT operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &repeat_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for REPEAT operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("REPEAT", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_ssm_conv(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node || !node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SSM_CONV operation\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32 || node->src[1]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SSM_CONV operation with unsupported types: dst=%s src0=%s src1=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ const char * kernel_name = "ssm_conv_f32";
+
+ ggml_et_ssm_conv_params params;
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.dst = *node;
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (ssm_conv_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_SSM_CONV)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for SSM_CONV operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &ssm_conv_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for SSM_CONV operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("SSM_CONV", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SSM_SCAN operation\n");
+ return false;
+ }
+
+ for (int i = 0; i < 7; ++i) {
+ if (!node->src[i]) {
+ GGML_LOG_ERROR("ET: SSM_SCAN missing required input %d\n", i);
+ return false;
+ }
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32 || node->src[1]->type != GGML_TYPE_F32 ||
+ node->src[2]->type != GGML_TYPE_F32 || node->src[3]->type != GGML_TYPE_F32 ||
+ node->src[4]->type != GGML_TYPE_F32 || node->src[5]->type != GGML_TYPE_F32 ||
+ node->src[6]->type != GGML_TYPE_I32) {
+ GGML_LOG_ERROR("ET: SSM_SCAN operation with unsupported types\n");
+ return false;
+ }
+
+ ggml_et_ssm_scan_params params;
+ params.src0 = *node->src[0];
+ params.src1 = *node->src[1];
+ params.src2 = *node->src[2];
+ params.src3 = *node->src[3];
+ params.src4 = *node->src[4];
+ params.src5 = *node->src[5];
+ params.src6 = *node->src[6];
+ params.dst = *node;
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "ssm_scan_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END("SSM_SCAN", "ssm_scan_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_rwkv_wkv6(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for RWKV_WKV6 operation\n");
+ return false;
+ }
+
+ // Validate all 6 source tensors exist
+ for (int i = 0; i <= 5; i++) {
+ if (!node->src[i]) {
+ GGML_LOG_ERROR("ET: RWKV_WKV6 operation missing src[%d]\n", i);
+ return false;
+ }
+ }
+
+ if (node->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: RWKV_WKV6 only supports F32, got %s\n", ggml_type_name(node->type));
+ return false;
+ }
+
+ const char * kernel_name = "rwkv_wkv6_f32";
+
+ const int64_t S = node->src[0]->ne[0]; // head_size
+ const int64_t H = node->src[0]->ne[1]; // num heads
+ const int64_t T = node->src[1]->ne[2]; // num tokens
+ const int64_t n_seqs = node->src[5]->ne[1]; // num sequences
+ const int64_t C = S * H;
+
+ ggml_et_rwkv_wkv6_params params;
+ params.k = (float *) node->src[0]->data;
+ params.v = (float *) node->src[1]->data;
+ params.r = (float *) node->src[2]->data;
+ params.tf = (float *) node->src[3]->data;
+ params.td = (float *) node->src[4]->data;
+ params.state_in = (float *) node->src[5]->data;
+ params.dst = (float *) node->data;
+ params.C = (int32_t) C;
+ params.H = (int32_t) H;
+ params.S = (int32_t) S;
+ params.T = (int32_t) T;
+ params.n_seqs = (int32_t) n_seqs;
+
+ // Phase 1: Initialize CPU comparison context
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (rwkv_wkv6_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_RWKV_WKV6)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for RWKV_WKV6 operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &rwkv_wkv6_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for RWKV_WKV6 operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("RWKV_WKV6", kernel_name, node, "S=%d H=%d T=%d n_seqs=%d", (int) S, (int) H, (int) T,
+ (int) n_seqs);
+ return kernel_result;
+}
+
+bool ggml_et_op_rwkv_wkv7(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for RWKV_WKV7 operation\n");
+ return false;
+ }
+
+ // Validate all 7 source tensors exist
+ for (int i = 0; i <= 6; i++) {
+ if (!node->src[i]) {
+ GGML_LOG_ERROR("ET: RWKV_WKV7 operation missing src[%d]\n", i);
+ return false;
+ }
+ }
+
+ if (node->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: RWKV_WKV7 only supports F32, got %s\n", ggml_type_name(node->type));
+ return false;
+ }
+
+ const char * kernel_name = "rwkv_wkv7_f32";
+
+ const int64_t S = node->src[2]->ne[0]; // head_size
+ const int64_t H = node->src[2]->ne[1]; // num heads
+ const int64_t T = node->src[1]->ne[2]; // num tokens
+ const int64_t n_seqs = node->src[6]->ne[1]; // num sequences
+ const int64_t C = S * H;
+
+ ggml_et_rwkv_wkv7_params params;
+ params.r = (float *) node->src[0]->data;
+ params.w = (float *) node->src[1]->data;
+ params.k = (float *) node->src[2]->data;
+ params.v = (float *) node->src[3]->data;
+ params.a = (float *) node->src[4]->data;
+ params.b = (float *) node->src[5]->data;
+ params.state_in = (float *) node->src[6]->data;
+ params.dst = (float *) node->data;
+ params.C = (int32_t) C;
+ params.H = (int32_t) H;
+ params.S = (int32_t) S;
+ params.T = (int32_t) T;
+ params.n_seqs = (int32_t) n_seqs;
+
+ // Phase 1: Initialize CPU comparison context
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (rwkv_wkv7_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_RWKV_WKV7)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for RWKV_WKV7 operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &rwkv_wkv7_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for RWKV_WKV7 operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("RWKV_WKV7", kernel_name, node, "S=%d H=%d T=%d n_seqs=%d", (int) S, (int) H, (int) T,
+ (int) n_seqs);
+ return kernel_result;
+}
+
+static ggml_et_cpu_compare_config gated_delta_net_cpu_compare_config = {
+ /* .enabled = */ false,
+ /* .use_cpu_result = */ false,
+ /* .log_differences = */ true,
+ /* .tolerance = */ 1e-4f,
+ /* .max_log_elements = */ 4096
+};
+
+bool ggml_et_op_gated_delta_net(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for GATED_DELTA_NET operation\n");
+ return false;
+ }
+
+ // Validate all 6 source tensors exist
+ for (int i = 0; i <= 5; i++) {
+ if (!node->src[i]) {
+ GGML_LOG_ERROR("ET: GATED_DELTA_NET operation missing src[%d]\n", i);
+ return false;
+ }
+ }
+
+ if (node->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: GATED_DELTA_NET only supports F32, got %s\n", ggml_type_name(node->type));
+ return false;
+ }
+
+ const char * kernel_name = "gated_delta_net_f32";
+
+ const ggml_tensor * src_q = node->src[0];
+ const ggml_tensor * src_k = node->src[1];
+ const ggml_tensor * src_v = node->src[2];
+ const ggml_tensor * src_g = node->src[3];
+ const ggml_tensor * src_beta = node->src[4];
+ const ggml_tensor * src_state = node->src[5];
+
+ const int64_t S_v = src_v->ne[0];
+ const int64_t H = src_v->ne[1];
+ const int64_t n_tokens = src_v->ne[2];
+ const int64_t n_seqs = src_v->ne[3];
+ const int64_t H_q = src_q->ne[1];
+ const int64_t H_k = src_k->ne[1];
+ const int64_t n_seqs_q = src_q->ne[3];
+ const int64_t n_seqs_k = src_k->ne[3];
+
+ ggml_et_gated_delta_net_params params;
+ params.q = *src_q;
+ params.k = *src_k;
+ params.v = *src_v;
+ params.g = *src_g;
+ params.beta = *src_beta;
+ params.state_in = *src_state;
+ params.dst = *node;
+ params.S_v = (int32_t) S_v;
+ params.H = (int32_t) H;
+ params.H_q = (int32_t) H_q;
+ params.H_k = (int32_t) H_k;
+ params.n_tokens = (int32_t) n_tokens;
+ params.n_seqs = (int32_t) n_seqs;
+ params.n_seqs_q = (int32_t) n_seqs_q;
+ params.n_seqs_k = (int32_t) n_seqs_k;
+ params.kda = (src_g->ne[0] == S_v) ? 1 : 0;
+ params.K = ggml_get_op_params_i32(node, 0);
+ params.scale = 1.0f / sqrtf((float) S_v);
+
+ // CPU comparison for debugging
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (gated_delta_net_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_GATED_DELTA_NET)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for GATED_DELTA_NET operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &gated_delta_net_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for GATED_DELTA_NET operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END_EXT("GATED_DELTA_NET", kernel_name, node, "S_v=%d H=%d n_tokens=%d n_seqs=%d kda=%d", (int) S_v,
+ (int) H, (int) n_tokens, (int) n_seqs, params.kda);
+ return kernel_result;
+}
+
+bool ggml_et_op_set_rows(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!dev_ctx || !node) {
+ GGML_LOG_ERROR("ET: Invalid parameters for SET_ROWS operation\n");
+ return false;
+ }
+
+ if (!node->src[0] || !node->src[1] || !node->src[2]) {
+ GGML_LOG_ERROR(
+ "ET: SET_ROWS operation missing required inputs (needs src[0]=base, src[1]=indices, src[2]=data)\n");
+ return false;
+ }
+
+ const char * kernel_name;
+
+ // Support F32 data with I64 indices -> F32/F16 output (scatter operation)
+ if (node->src[0]->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_I64 &&
+ (node->type == GGML_TYPE_F32 || node->type == GGML_TYPE_F16)) {
+ if (node->type == GGML_TYPE_F32 || node->type == GGML_TYPE_F16) {
+ kernel_name = "set_rows_f32";
+ } else {
+ GGML_LOG_ERROR("ET: SET_ROWS unsupported output type: %s\n", ggml_type_name(node->type));
+ return false;
+ }
+
+ } else {
+ GGML_LOG_ERROR("ET: SET_ROWS operation with unsupported types: dst=%s src0=%s src1=%s\n",
+ ggml_type_name(node->type), ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ // Validate contiguity requirements
+ if (!ggml_is_contiguous_rows(node)) {
+ GGML_LOG_ERROR("ET: SET_ROWS operation requires contiguous-rows destination tensor\n");
+ return false;
+ }
+
+ if (!ggml_is_contiguous_rows(node->src[0])) {
+ GGML_LOG_ERROR("ET: SET_ROWS operation requires contiguous-rows source tensor\n");
+ return false;
+ }
+
+ if (!ggml_is_contiguous(node->src[1])) {
+ GGML_LOG_ERROR("ET: SET_ROWS operation requires contiguous indices tensor\n");
+ return false;
+ }
+
+ // Validate dimension constraints from ggml implementation
+ if (!(node->ne[0] == node->src[0]->ne[0] && // same number of columns
+ node->ne[2] == node->src[0]->ne[2] && // same batch size
+ node->ne[3] == node->src[0]->ne[3] && // same outer dimension
+ node->src[0]->ne[1] == node->src[1]->ne[0] && // src rows = index count
+ node->src[0]->ne[2] % node->src[1]->ne[1] == 0 && // batch constraint
+ node->src[0]->ne[3] % node->src[1]->ne[2] == 0 && // outer constraint
+ node->src[1]->ne[3] == 1)) { // indices constraint
+ GGML_LOG_ERROR("ET: SET_ROWS operation dimension constraint failed\n");
+ return false;
+ }
+
+ ggml_et_set_rows_params params;
+ params.src0 = *node->src[0]; // F32 source data tensor
+ params.src1 = *node->src[1]; // I64 indices tensor
+ params.dst = *node; // F32/F16 destination tensor
+
+ // Phase 1: Initialize CPU comparison context and copy source buffers (before ET kernel)
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (set_rows_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_SET_ROWS)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for SET_ROWS operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, kernel_name, ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ // Phase 2: Execute CPU computation and compare with ET result (after ET kernel)
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &set_rows_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for SET_ROWS operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("SET_ROWS", kernel_name, node);
+ return kernel_result;
+}
+
+bool ggml_et_op_fill(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ ggml_et_fill_params params;
+ params.dst = *node;
+ memcpy(¶ms.c, node->op_params, sizeof(float));
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "fill_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END("FILL", "fill_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_diag(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: DIAG operation missing source tensor\n");
+ return false;
+ }
+
+ ggml_et_diag_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "diag_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END("DIAG", "diag_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_tri(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: TRI operation missing source tensor\n");
+ return false;
+ }
+
+ ggml_et_tri_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+ memcpy(¶ms.tri_type, node->op_params, sizeof(int32_t));
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "tri_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END("TRI", "tri_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_solve_tri(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: SOLVE_TRI operation missing source tensor(s)\n");
+ return false;
+ }
+
+ ggml_et_solve_tri_params params;
+ params.src0 = *node->src[0]; // A (lower-triangular)
+ params.src1 = *node->src[1]; // B (RHS)
+ params.dst = *node; // X (solution)
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "solve_tri_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END("SOLVE_TRI", "solve_tri_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_set(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!node->src[0] || !node->src[1]) {
+ GGML_LOG_ERROR("ET: SET operation missing source tensor(s)\n");
+ return false;
+ }
+
+ const bool inplace = (bool) ((const int32_t *) node->op_params)[4];
+ const size_t offset = ((const int32_t *) node->op_params)[3];
+ const size_t nb1 = ((const int32_t *) node->op_params)[0];
+ const size_t nb2 = ((const int32_t *) node->op_params)[1];
+ const size_t nb3 = ((const int32_t *) node->op_params)[2];
+
+ if (!inplace) {
+ GGML_LOG_ERROR("ET: SET only supports inplace (inplace=%d)\n", inplace);
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32 || node->src[1]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: SET only supports F32 (dst=%s src0=%s src1=%s)\n", ggml_type_name(node->type),
+ ggml_type_name(node->src[0]->type), ggml_type_name(node->src[1]->type));
+ return false;
+ }
+
+ if (!ggml_are_same_shape(node, node->src[0])) {
+ GGML_LOG_ERROR("ET: SET requires same-shape src0 and dst\n");
+ return false;
+ }
+
+ if (!ggml_is_contiguous(node) || !ggml_is_contiguous(node->src[0]) || !ggml_is_contiguous(node->src[1])) {
+ GGML_LOG_ERROR("ET: SET requires contiguous dst, src0, and src1\n");
+ return false;
+ }
+
+ ggml_et_set_params params;
+ params.src1 = *node->src[1];
+ params.dst = *node;
+ params.nb1 = (int32_t) nb1;
+ params.nb2 = (int32_t) nb2;
+ params.nb3 = (int32_t) nb3;
+ params.offset = (int32_t) offset;
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "set_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ ET_PERF_END("SET", "set_f32", node);
+ return kernel_result;
+}
+
+bool ggml_et_op_pad(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node) {
+ ET_PERF_START();
+
+ if (!node->src[0]) {
+ GGML_LOG_ERROR("ET: PAD operation missing source tensor\n");
+ return false;
+ }
+
+ if (node->type != GGML_TYPE_F32 || node->src[0]->type != GGML_TYPE_F32) {
+ GGML_LOG_ERROR("ET: PAD only supports F32 (src=%s dst=%s)\n", ggml_type_name(node->src[0]->type),
+ ggml_type_name(node->type));
+ return false;
+ }
+
+ if (!ggml_is_contiguous(node)) {
+ GGML_LOG_ERROR("ET: PAD requires contiguous output tensor\n");
+ return false;
+ }
+
+ if (node->src[0]->nb[0] != sizeof(float)) {
+ GGML_LOG_ERROR("ET: PAD requires element-contiguous src dim0 (nb[0]=%zu)\n", (size_t) node->src[0]->nb[0]);
+ return false;
+ }
+
+ // Extract padding parameters from op_params
+ const int32_t * op_params = (const int32_t *) node->op_params;
+
+ ggml_et_pad_params params;
+ params.src0 = *node->src[0];
+ params.dst = *node;
+ params.lp[0] = op_params[0];
+ params.rp[0] = op_params[1];
+ params.lp[1] = op_params[2];
+ params.rp[1] = op_params[3];
+ params.lp[2] = op_params[4];
+ params.rp[2] = op_params[5];
+ params.lp[3] = op_params[6];
+ params.rp[3] = op_params[7];
+
+ // v1: no dim0 padding
+ if (params.lp[0] != 0 || params.rp[0] != 0) {
+ GGML_LOG_ERROR("ET: PAD dim0 padding not supported (lp0=%d rp0=%d)\n", params.lp[0], params.rp[0]);
+ return false;
+ }
+
+ ggml_et_cpu_compare_ctx cpu_cmp_ctx;
+ bool cpu_comparison_active = false;
+ if (pad_cpu_compare_config.enabled) {
+ if (ggml_et_cpu_compare_init_pre(&cpu_cmp_ctx, node, GGML_OP_PAD)) {
+ cpu_comparison_active = true;
+ } else {
+ GGML_LOG_WARN("ET: Failed to initialize CPU comparison for PAD operation\n");
+ }
+ }
+
+ bool kernel_result = ggml_et_launch_kernel(dev_ctx, "pad_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
+
+ if (cpu_comparison_active) {
+ if (!ggml_et_cpu_compare_compute_and_check(&cpu_cmp_ctx, node, &pad_cpu_compare_config)) {
+ GGML_LOG_WARN("ET: CPU comparison failed for PAD operation\n");
+ }
+ ggml_et_cpu_compare_free(&cpu_cmp_ctx);
+ }
+
+ ET_PERF_END("PAD", "pad_f32", node);
+ return kernel_result;
+}
--- /dev/null
+#pragma once
+
+#include "ggml-et-common.h"
+#include "ggml.h"
+
+#include <inttypes.h>
+
+// Performance logging macros for ET ops
+// Logs in machine-parseable pipe-delimited format: ET_PERF|field=value|...
+#ifdef ET_PERF_RECORD
+# define ET_PERF_START() int64_t _et_perf_start = ggml_time_us()
+
+# define ET_PERF_END(op_name, kernel_name, node) \
+ do { \
+ int64_t _et_perf_end = ggml_time_us(); \
+ int64_t _et_perf_duration = _et_perf_end - _et_perf_start; \
+ GGML_LOG_DEBUG("ET_PERF|op=%s|kernel=%s|duration_us=%" PRId64 "|tensor=%s|shape=[%" PRId64 ",%" PRId64 \
+ ",%" PRId64 ",%" PRId64 "]|start_us=%" PRId64 "|end_us=%" PRId64 "\n", \
+ op_name, kernel_name, _et_perf_duration, (node)->name, (node)->ne[0], (node)->ne[1], \
+ (node)->ne[2], (node)->ne[3], _et_perf_start, _et_perf_end); \
+ } while (0)
+
+# define ET_PERF_END_EXT(op_name, kernel_name, node, fmt, ...) \
+ do { \
+ int64_t _et_perf_end = ggml_time_us(); \
+ int64_t _et_perf_duration = _et_perf_end - _et_perf_start; \
+ GGML_LOG_DEBUG("ET_PERF|op=%s|kernel=%s|duration_us=%" PRId64 "|tensor=%s|shape=[%" PRId64 ",%" PRId64 \
+ ",%" PRId64 ",%" PRId64 "]|start_us=%" PRId64 "|end_us=%" PRId64 "|" fmt "\n", \
+ op_name, kernel_name, _et_perf_duration, (node)->name, (node)->ne[0], (node)->ne[1], \
+ (node)->ne[2], (node)->ne[3], _et_perf_start, _et_perf_end, ##__VA_ARGS__); \
+ } while (0)
+#else
+
+# define ET_PERF_START() \
+ do { \
+ } while (0)
+# define ET_PERF_END_EXT(op_name, kernel_name, node, fmt, ...) \
+ do { \
+ (void) (node); \
+ } while (0)
+# define ET_PERF_END(op_name, kernel_name, node) \
+ do { \
+ (void) (node); \
+ } while (0)
+
+#endif // ET_PERF_RECORD
+
+struct ggml_et_binary_params {
+ ggml_tensor src0;
+ ggml_tensor src1;
+ ggml_tensor dst;
+};
+
+// Q8_0 mul_mat with optional residual bias.
+// bias.data == NULL means "no bias" - kernel skips the add.
+// When non-NULL, bias must have the same shape and strides as dst.
+struct ggml_et_mm_q8_params {
+ ggml_tensor src0;
+ ggml_tensor src1;
+ ggml_tensor dst;
+ ggml_tensor bias;
+};
+
+struct ggml_et_im2col_params {
+ ggml_tensor src0;
+ ggml_tensor src1;
+ ggml_tensor dst;
+};
+
+// Element map parameters for embarrassingly parallel binary operations (MUL, ADD, etc.)
+// Operation type is determined by dst->op (GGML_OP_MUL, GGML_OP_ADD, etc.)
+struct ggml_et_elmap_params {
+ ggml_tensor src0;
+ ggml_tensor src1;
+ ggml_tensor dst;
+};
+
+struct ggml_et_rope_settings {
+ int32_t n_past;
+ int32_t n_dims; // Number of dimensions to apply ROPE to (must be even)
+ int32_t mode; // ROPE mode, GGML_ROPE_TYPE_*
+ int32_t n_ctx;
+ int32_t n_ctx_orig;
+ float freq_base; // Base frequency (usually 10000.0f)
+ float freq_scale; // Frequency scaling factor
+ float ext_factor; // Extension factor for YaRN
+ float attn_factor; // Attention factor for YaRN
+ float beta_fast; // Fast beta for YaRN
+ float beta_slow; // Slow beta for YaRN
+ int32_t sections[4]; // Sections for multi-modal ROPE
+};
+
+struct ggml_et_rope_params {
+ ggml_tensor src0;
+ ggml_tensor src1;
+ ggml_tensor src2;
+ ggml_tensor dst;
+ ggml_et_rope_settings rope_params;
+};
+
+struct ggml_et_rms_norm_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon parameter for numerical stability
+};
+
+struct ggml_et_norm_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon parameter for numerical stability
+};
+
+struct ggml_et_l2_norm_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon parameter for numerical stability
+};
+
+struct ggml_et_group_norm_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ int32_t n_groups; // Number of channel groups
+ float eps; // Epsilon parameter for numerical stability
+};
+
+struct ggml_et_glu_params {
+ ggml_tensor src0; // F32 input tensor A (or combined tensor if src1 is null)
+ ggml_tensor src1; // F32 input tensor B (null for single tensor mode)
+ ggml_tensor dst; // F32 output tensor (n/2 columns)
+ int32_t glu_op_type; // GLU operation type (REGLU=0, GEGLU=1, SWIGLU=2, etc.)
+ int32_t swapped; // Whether gate and value are swapped
+ float alpha; // SWIGLU_OAI: sigmoid scaling factor (unused for other variants)
+ float limit; // SWIGLU_OAI: clamp limit (unused for other variants)
+};
+
+struct ggml_et_softmax_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor src1; // F32 mask tensor (optional, may be zeroed if not used)
+ ggml_tensor src2; // F32 sinks tensor (optional, may be zeroed if not used)
+ ggml_tensor dst; // F32 output tensor
+ float scale; // Scale factor
+ float max_bias; // Max bias for ALiBi (0.0f if not used)
+};
+
+struct ggml_et_flash_attn_ext_params {
+ ggml_tensor src0; // Q tensor (F32)
+ ggml_tensor src1; // K tensor (F32)
+ ggml_tensor src2; // V tensor (F32)
+ ggml_tensor mask; // mask tensor (F16 or F32), zeroed when absent
+ ggml_tensor dst; // Output tensor (F32)
+ float scale; // Scale factor applied to QK
+ int32_t has_mask; // nonzero if mask is present
+};
+
+struct ggml_et_get_rows_params {
+ ggml_tensor src0; // Data tensor (F32 or Q8_0)
+ ggml_tensor src1; // Row indices tensor (I32)
+ ggml_tensor dst; // Output tensor (F32)
+};
+
+struct ggml_et_cont_params {
+ ggml_tensor src0; // F32 input tensor (non-contiguous)
+ ggml_tensor dst; // F32 output tensor (contiguous)
+};
+
+struct ggml_et_concat_params {
+ ggml_tensor src0; // F32 input tensor 0
+ ggml_tensor src1; // F32 input tensor 1
+ ggml_tensor dst; // F32 output tensor
+ int32_t dim; // Concatenation dimension
+};
+
+struct ggml_et_repeat_params {
+ ggml_tensor src0; // F32 input tensor (tile)
+ ggml_tensor dst; // F32 output tensor (tiled result)
+};
+
+struct ggml_et_fill_params {
+ ggml_tensor dst; // F32 output tensor (contiguous)
+ float c; // Constant value to fill
+};
+
+struct ggml_et_tri_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ int32_t tri_type; // ggml_tri_type enum value
+};
+
+struct ggml_et_solve_tri_params {
+ ggml_tensor src0; // A: lower-triangular [n, n, B1, B2]
+ ggml_tensor src1; // B: RHS [k, n, B1, B2]
+ ggml_tensor dst; // X: solution [k, n, B1, B2]
+};
+
+struct ggml_et_pad_params {
+ ggml_tensor src0; // F32 input (may be non-contiguous, nb[0] must == 4)
+ ggml_tensor dst; // F32 output (contiguous, ne[0] % 16 == 0)
+ int32_t lp[4]; // left padding per dimension
+ int32_t rp[4]; // right padding per dimension
+};
+
+struct ggml_et_diag_params {
+ ggml_tensor src0; // F32 input vector
+ ggml_tensor dst; // F32 output diagonal matrix
+};
+
+struct ggml_et_ssm_conv_params {
+ ggml_tensor src0; // conv_x: [d_conv - 1 + n_t, d_inner, n_seqs]
+ ggml_tensor src1; // conv1d.weight: [d_conv, d_inner]
+ ggml_tensor dst; // output: [d_inner, n_t, n_seqs]
+};
+
+struct ggml_et_ssm_scan_params {
+ ggml_tensor src0; // s: [d_state, head_dim, n_head, n_seqs]
+ ggml_tensor src1; // x: [head_dim, n_head, n_seq_tokens, n_seqs]
+ ggml_tensor src2; // dt: [n_head, n_seq_tokens, n_seqs]
+ ggml_tensor src3; // A: [d_state, n_head] or [1, n_head]
+ ggml_tensor src4; // B: [d_state, n_group, n_seq_tokens, n_seqs]
+ ggml_tensor src5; // C: [d_state, n_group, n_seq_tokens, n_seqs]
+ ggml_tensor src6; // ids: [n_seqs] i32
+ ggml_tensor dst; // [y, final_state] packed output from ggml_ssm_scan()
+};
+
+struct ggml_et_rwkv_wkv6_params {
+ float * k; // src[0]: [S, H, T] key
+ float * v; // src[1]: [S, H, T] value
+ float * r; // src[2]: [S, H, T] receptance
+ float * tf; // src[3]: [S, H] time_faaaa (per-head)
+ float * td; // src[4]: [S, H, T] time_decay
+ float * state_in; // src[5]: [S*S*H, n_seqs] initial state
+ float * dst; // [C, T + S*n_seqs] output + state_out
+ int32_t C; // total channels (S * H)
+ int32_t H; // number of heads
+ int32_t S; // head size
+ int32_t T; // number of tokens
+ int32_t n_seqs; // number of sequences
+};
+
+struct ggml_et_rwkv_wkv7_params {
+ float * r; // [S, H, T] receptance
+ float * w; // [S, H, T] decay
+ float * k; // [S, H, T] key
+ float * v; // [S, H, T] value
+ float * a; // [S, H, T] bonus gate
+ float * b; // [S, H, T] bonus key
+ float * state_in; // [S*S*H, n_seqs] initial state
+ float * dst; // [C, T + S*n_seqs] output + state_out
+ int32_t C; // total channels (S * H)
+ int32_t H; // number of heads
+ int32_t S; // head size
+ int32_t T; // number of tokens
+ int32_t n_seqs; // number of sequences
+};
+
+struct ggml_et_gated_delta_net_params {
+ ggml_tensor q; // [S_v, H_q, n_tokens, n_seqs_q]
+ ggml_tensor k; // [S_v, H_k, n_tokens, n_seqs_k]
+ ggml_tensor v; // [S_v, H, n_tokens, n_seqs]
+ ggml_tensor g; // [1 or S_v, H, n_tokens, n_seqs]
+ ggml_tensor beta; // [1, H, n_tokens, n_seqs]
+ ggml_tensor state_in; // [S_v*S_v*H, K, n_seqs]
+ ggml_tensor dst; // [S_v*H, n_tokens*n_seqs + S_v*n_seqs*K]
+ int32_t S_v; // head dimension (value size)
+ int32_t H; // number of value heads
+ int32_t H_q; // number of Q heads
+ int32_t H_k; // number of K heads
+ int32_t n_tokens; // total tokens
+ int32_t n_seqs; // number of sequences (from V)
+ int32_t n_seqs_q; // Q sequence count
+ int32_t n_seqs_k; // K sequence count
+ int32_t kda; // 1 if per-element gate (g_ne0 == S_v), 0 if scalar
+ int32_t K; // snapshot slot count
+ float scale; // 1/sqrt(S_v)
+};
+
+struct ggml_et_set_rows_params {
+ ggml_tensor src0; // F32 source data tensor
+ ggml_tensor src1; // I64 row indices tensor
+ ggml_tensor dst; // F32/F16 destination tensor
+};
+
+struct ggml_et_set_params {
+ ggml_tensor src1; // F32 source view to write into dst
+ ggml_tensor dst; // F32 destination/base tensor
+ int32_t nb1; // destination view stride for dim 1
+ int32_t nb2; // destination view stride for dim 2
+ int32_t nb3; // destination view stride for dim 3
+ int32_t offset; // byte offset into destination
+};
+
+struct ggml_et_rms_norm_mul_params {
+ ggml_tensor src0; // F32 input tensor (to be normalized)
+ ggml_tensor src1; // F32 weights tensor (element-wise multiply)
+ ggml_tensor dst; // F32 output tensor
+ float eps; // Epsilon for numerical stability
+};
+
+struct ggml_et_mul_mat_id_params {
+ ggml_tensor src0; // Expert weight matrices (Q8_0/F16/F32) [K, M, n_expert]
+ ggml_tensor src1; // Activations (F32) [K, n_expert_used, batch]
+ ggml_tensor src2; // Expert indices (I32) [n_expert_used, batch]
+ ggml_tensor dst; // Output (F32) [M, n_expert_used, batch, 1]
+};
+
+struct ggml_et_sqr_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+};
+
+struct ggml_et_unary_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ int32_t unary_op; // ggml_unary_op enum value
+};
+
+struct ggml_et_sum_rows_params {
+ ggml_tensor src0; // F32 input tensor [ne00, ne01, ne02, ne03]
+ ggml_tensor dst; // F32 output tensor [1, ne01, ne02, ne03]
+};
+
+struct ggml_et_mean_params {
+ ggml_tensor src0; // F32 input tensor [ne00, ne01, ne02, ne03]
+ ggml_tensor dst; // F32 output tensor [1, ne01, ne02, ne03]
+};
+
+struct ggml_et_clamp_params {
+ ggml_tensor src0; // F32 input tensor (contiguous)
+ ggml_tensor dst; // F32 output tensor (contiguous; may alias src0)
+ float min_val;
+ float max_val;
+};
+
+struct ggml_et_cumsum_params {
+ ggml_tensor src0; // F32 input tensor [ne00, ne01, ne02, ne03]
+ ggml_tensor dst; // F32 output tensor [ne00, ne01, ne02, ne03]
+};
+
+struct ggml_et_scale_params {
+ ggml_tensor src0; // F32 input tensor
+ ggml_tensor dst; // F32 output tensor
+ float scale; // Scale factor
+ float bias; // Bias (additive offset)
+};
+
+bool ggml_et_op_cumsum(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_sqr(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_unary(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_sum_rows(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_mean(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_clamp(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_scale(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_mul(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_add(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_sub(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+// add_node is optional: when non-NULL and the pair (node, add_node) was
+// validated by ggml_et_can_fuse({MUL_MAT, ADD}), the Q8_0 path writes
+// dst = mm(...) + add_node's "other" operand (the bias) in one launch.
+bool ggml_et_op_mul_mat(ggml_backend_et_device_context * dev_ctx,
+ const ggml_tensor * node,
+ const ggml_tensor * add_node = nullptr);
+bool ggml_et_op_mul_mat_id(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_rope(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_rms_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_l2_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_group_norm(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_glu(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_softmax(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_im2col(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_conv_2d(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_flash_attn_ext(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_get_rows(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_set_rows(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_cont(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_concat(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_repeat(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_rwkv_wkv6(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_rwkv_wkv7(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_cpy(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_gated_delta_net(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_elmap(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_fill(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_diag(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_tri(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_solve_tri(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_pad(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_set(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_ssm_conv(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_tensor * node);
+bool ggml_et_op_rms_norm_mul(ggml_backend_et_device_context * dev_ctx,
+ const ggml_tensor * rms_norm_node,
+ const ggml_tensor * mul_node);
--- /dev/null
+#pragma once
+
+#include <stdint.h>
+
+struct ggml_et_uberkernel_inst {
+ uint16_t kernel_id;
+ uint16_t flags;
+ uint32_t params_offset;
+ uint32_t params_size;
+};
+
+struct ggml_et_uberkernel_params {
+ uint32_t num_insts;
+ uint32_t inst_stride;
+ uint64_t insts;
+ uint64_t params_blob;
+};
--- /dev/null
+#include "ggml-et.h"
+
+#include "ggml-backend-impl.h"
+#include "ggml-backend.h"
+#include "ggml-et-common.h"
+#include "ggml-et-kernels.h"
+#include "ggml-et-memops.h"
+#include "ggml-et-ops.h"
+#include "ggml-impl.h"
+#include "ggml.h"
+
+#include <stdarg.h>
+
+#include <cstdarg>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+#if __has_include(<filesystem>)
+# include <filesystem>
+namespace fs = std::filesystem;
+#elif __has_include(<experimental/filesystem>)
+# include <experimental/filesystem>
+namespace fs = std::experimental::filesystem;
+#else
+# error "cannot include the filesystem library"
+#endif
+
+/*
+ * ggml_et_dump_tensor_metadata
+ * @brief prints the metadata of a single tensorf
+ */
+static void ggml_et_dump_tensor_metadata(const ggml_tensor * ggtensor, size_t indent_level, const char * title) {
+ char * spaces = (char *) alloca(indent_level + 1);
+ memset(spaces, ' ', indent_level);
+ spaces[indent_level] = '\0';
+ fprintf(stderr,
+ "%s%s: %s\n"
+ "%s type: %s\n"
+ "%s ne: %lld %lld %lld %lld\n"
+ "%s nb: %zu %zu %zu %zu\n"
+ "%s op: %s\n"
+ "%s data: %p\n"
+ "%s src0: %p\n",
+ spaces, title, ggtensor->name, spaces, ggml_type_name(ggtensor->type), spaces, (long long) ggtensor->ne[0],
+ (long long) ggtensor->ne[1], (long long) ggtensor->ne[2], (long long) ggtensor->ne[3], spaces,
+ ggtensor->nb[0], ggtensor->nb[1], ggtensor->nb[2], ggtensor->nb[3], spaces, ggml_op_name(ggtensor->op),
+ spaces, ggtensor->data, spaces, (void *) ggtensor->src[0]);
+}
+
+/*
+ * ggml_et_dump_operator_metadata
+ * @brief prints the metadata of a single tensor (or operator) including it's input and views
+ */
+static void ggml_et_dump_operator_metadata(const ggml_tensor * ggtensor) {
+ GGML_ASSERT(ggtensor != NULL);
+ ggml_et_dump_tensor_metadata(ggtensor, 0, "GGML tensor");
+ for (int i = 0; i < GGML_MAX_SRC && ggtensor->src[i]; i++) {
+ char arr[16];
+ int n = snprintf(arr, sizeof(arr), "src[%i]->name", i);
+ GGML_ASSERT((unsigned) n < sizeof(arr) && "printed too much data to stack buffer");
+ ggml_et_dump_tensor_metadata(ggtensor->src[i], 2, arr);
+ }
+ if (ggtensor->view_src) {
+ ggml_et_dump_tensor_metadata(ggtensor, 2, "view_src");
+ }
+}
+
+static struct ggml_et_driver {
+ std::shared_ptr<dev::IDeviceLayer> device_layer;
+ std::shared_ptr<rt::IRuntime> runtime;
+ std::unique_ptr<std::ofstream> profile_stream;
+ std::unique_ptr<std::ofstream> kernel_id_stream;
+ std::vector<std::pair<std::string, rt::KernelId>> kernel_map;
+ bool profiling_enabled = false;
+} _drv;
+
+// Check at runtime environment variables for paths likely holding ET toolchain with sysemu elf files
+static std::string ggml_et_get_default_et_path() {
+ // List of environment variables to check in order of preference
+ const char * const env_vars[] = { "ET_TOOLCHAIN", "TOOLCHAIN_ROOT" };
+
+ for (const char * var : env_vars) {
+ if (const char * et_path = std::getenv(var)) {
+ if (et_path && *et_path != '\0') {
+ return fs::path(et_path).string();
+ }
+ }
+ }
+
+ // Otherwise assume default
+ return fs::path("/opt/et").string();
+}
+
+// config when using sysemu instead of PCIe hardware device
+// adapted from `ainekko/et-platform/esperanto-tools-libs/tools/src/bench.cpp`
+static inline auto ggml_et_get_default_sysemu_options() {
+ constexpr uint64_t kSysEmuMaxCycles = std::numeric_limits<uint64_t>::max();
+ constexpr uint64_t kSysEmuMinionShiresMask = 0x1FFFFFFFFu;
+ const std::string et_path = ggml_et_get_default_et_path() + "/";
+
+ emu::SysEmuOptions sysEmuOptions;
+
+ // Construct all paths
+ sysEmuOptions.bootromTrampolineToBL2ElfPath =
+ et_path + "lib/esperanto-fw/BootromTrampolineToBL2/BootromTrampolineToBL2.elf";
+ sysEmuOptions.spBL2ElfPath =
+ et_path + "lib/esperanto-fw/ServiceProcessorBL2/fast-boot/ServiceProcessorBL2_fast-boot.elf";
+ sysEmuOptions.machineMinionElfPath = et_path + "lib/esperanto-fw/MachineMinion/MachineMinion.elf";
+ sysEmuOptions.masterMinionElfPath = et_path + "lib/esperanto-fw/MasterMinion/MasterMinion.elf";
+ sysEmuOptions.workerMinionElfPath = et_path + "lib/esperanto-fw/WorkerMinion/WorkerMinion.elf";
+ sysEmuOptions.executablePath = et_path + "bin/sys_emu";
+
+ // Check that each path has a valid existing non-zero file otherwise emulator just silently hangs
+ const std::vector<std::string> required_files = {
+ sysEmuOptions.bootromTrampolineToBL2ElfPath, sysEmuOptions.spBL2ElfPath,
+ sysEmuOptions.machineMinionElfPath, sysEmuOptions.masterMinionElfPath,
+ sysEmuOptions.workerMinionElfPath, sysEmuOptions.executablePath,
+ };
+
+ for (const auto & file : required_files) {
+ if (!fs::exists(file) || fs::file_size(file) == 0) {
+ // Check that each path has a valid existing non-zero file otherwise emulator just silently hangs
+ GGML_LOG_ERROR("ET: Unable to find required sysemu file: %s", file.c_str());
+ GGML_LOG_ERROR("ET: Confirm et-platform is correctly installed at configured path.");
+ abort();
+ }
+ }
+
+ sysEmuOptions.runDir = (fs::current_path().string() + "/");
+ sysEmuOptions.maxCycles = kSysEmuMaxCycles;
+ sysEmuOptions.minionShiresMask = kSysEmuMinionShiresMask;
+ sysEmuOptions.puUart0Path = sysEmuOptions.runDir + "pu_uart0_tx.log";
+ sysEmuOptions.puUart1Path = sysEmuOptions.runDir + "pu_uart1_tx.log";
+ sysEmuOptions.spUart0Path = sysEmuOptions.runDir + "spio_uart0_tx.log";
+ sysEmuOptions.spUart1Path = sysEmuOptions.runDir + "spio_uart1_tx.log";
+ sysEmuOptions.startGdb = false;
+ sysEmuOptions.memcheck = false;
+
+ return sysEmuOptions;
+}
+
+// Forward declaration
+static void ggml_et_driver_cleanup();
+
+static bool ggml_et_driver_init() {
+ if (_drv.runtime != nullptr) {
+ assert(_drv.device_layer != nullptr);
+ } else {
+ try {
+#if defined GGML_ET_SYSEMU && GGML_ET_SYSEMU
+ // For emulator device using sysEmuOptions provided by function above enabled compiling with `-DGGML_ET_SYSEMU=ON`
+ _drv.device_layer = dev::IDeviceLayer::createSysEmuDeviceLayer(ggml_et_get_default_sysemu_options());
+#else
+ // For physical PCIe device
+ _drv.device_layer = dev::IDeviceLayer::createPcieDeviceLayer();
+#endif // GGML_ET_SYSEMU
+
+ _drv.runtime = rt::IRuntime::create(_drv.device_layer);
+
+ // Initialize profiler if requested via environment variable
+ const char * profile_path = getenv("GGML_ET_PROFILE");
+ if (profile_path) {
+ std::string output_path = std::string(profile_path) + "/et_runtime_trace.json";
+ std::string kernel_id_path = std::string(profile_path) + "/kernel_id.json";
+
+ _drv.profile_stream = std::make_unique<std::ofstream>(output_path);
+ _drv.kernel_id_stream = std::make_unique<std::ofstream>(kernel_id_path);
+ if (!_drv.profile_stream->is_open()) {
+ GGML_LOG_ERROR("ET: Failed to open profiling output file: %s", output_path.c_str());
+ abort();
+ }
+ if (!_drv.kernel_id_stream->is_open()) {
+ GGML_LOG_ERROR("ET: Failed to open profiling kernel map: %s", kernel_id_path.c_str());
+ abort();
+ }
+
+ auto * profiler = _drv.runtime->getProfiler();
+ profiler->start(*_drv.profile_stream, rt::IProfiler::OutputType::Json);
+ _drv.profiling_enabled = true;
+ GGML_LOG_INFO("ET: Runtime profiler started (JSON format)");
+
+ // Register cleanup at program exit
+ std::atexit(ggml_et_driver_cleanup);
+ }
+ } catch (const std::exception & e) {
+ GGML_LOG_ERROR("ggml_et: %s", e.what());
+ if (_drv.device_layer != nullptr) {
+ _drv.device_layer.reset();
+ }
+ if (_drv.runtime != nullptr) {
+ _drv.runtime.reset();
+ }
+ return false;
+ }
+ }
+ return true;
+}
+
+static std::shared_ptr<dev::IDeviceLayer> ggml_et_devicelayer() {
+ return _drv.device_layer;
+}
+
+std::shared_ptr<rt::IRuntime> ggml_et_runtime() {
+ return _drv.runtime;
+}
+
+static void ggml_et_driver_cleanup() {
+ if (_drv.profiling_enabled && _drv.runtime) {
+ GGML_LOG_INFO("ET: Stopping runtime profiler");
+ auto * profiler = _drv.runtime->getProfiler();
+ profiler->stop();
+ _drv.profiling_enabled = false;
+
+ if (_drv.profile_stream) {
+ _drv.profile_stream->close();
+ _drv.profile_stream.reset();
+ }
+
+ // Save kernel map
+ if (_drv.kernel_id_stream && !_drv.kernel_map.empty()) {
+ auto & os = *_drv.kernel_id_stream;
+ // XXX: Manual JSON construction. Not pretty but removes dependency
+ os << "{\n";
+ for (size_t i = 0; i < _drv.kernel_map.size(); i++) {
+ os << " \"" << _drv.kernel_map[i].first << "\": " << (int) _drv.kernel_map[i].second;
+ if (i + 1 < _drv.kernel_map.size()) {
+ os << ",";
+ }
+ os << "\n";
+ }
+ os << "}\n";
+ _drv.kernel_id_stream->close();
+ _drv.kernel_id_stream.reset();
+ }
+ }
+}
+
+static ggml_backend_dev_t ggml_backend_et_reg_get_device(ggml_backend_reg_t reg, size_t devidx);
+
+static void ggml_backend_et_buffer_free_buffer(ggml_backend_buffer_t buffer) {
+ ggml_backend_et_buffer_context * ctx = (ggml_backend_et_buffer_context *) buffer->context;
+ if (ctx->data != nullptr) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (runtime) {
+ runtime->freeDevice(ctx->rtid, static_cast<std::byte *>(ctx->data));
+ }
+ }
+ delete ctx;
+}
+
+static void * ggml_backend_et_buffer_get_base(ggml_backend_buffer_t buffer) {
+ ggml_backend_et_buffer_context * ctx = (ggml_backend_et_buffer_context *) buffer->context;
+ return ctx->data;
+}
+
+static ggml_status ggml_backend_et_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {
+ // View tensors share buffer with their view_src, no additional initialization needed
+ if (tensor->view_src != NULL) {
+ return GGML_STATUS_SUCCESS;
+ }
+
+ const size_t original_size = ggml_nbytes(tensor);
+ const size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor);
+
+ // Clear padding bytes to avoid NaN values
+ // XXX: Martin - do we need this?
+ if (padded_size > original_size) {
+ const size_t padding_size = padded_size - original_size;
+
+ // Get device context to access memops kernel
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) buffer->buft->device->context;
+ if (!dev_ctx) {
+ GGML_LOG_ERROR("ET: Failed to get device context for padding clear");
+ return GGML_STATUS_FAILED;
+ }
+
+ // Use device-side memset kernel for efficient padding clear
+ std::byte * padding_ptr = static_cast<std::byte *>(tensor->data) + original_size;
+ if (!ggml_et_memset(dev_ctx, padding_ptr, 0, padding_size)) {
+ GGML_LOG_ERROR("ET: Failed to clear padding using memset kernel for tensor %s", tensor->name);
+ return GGML_STATUS_FAILED;
+ }
+ }
+
+ return GGML_STATUS_SUCCESS;
+}
+
+static void ggml_backend_et_buffer_set_tensor(ggml_backend_buffer_t buffer,
+ ggml_tensor * tensor,
+ const void * data,
+ size_t offset,
+ size_t size) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ return;
+ }
+
+ // Create short-lived stream for this transfer
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) buffer->buft->device->context;
+ rt::StreamId stream = dev_ctx->default_stream;
+
+ std::byte * dst_ptr = static_cast<std::byte *>(tensor->data) + offset;
+ const std::byte * src_ptr = static_cast<const std::byte *>(data);
+
+ rt::EventId event = runtime->memcpyHostToDevice(stream, src_ptr, dst_ptr, size, true /*barrier*/);
+
+ runtime->waitForEvent(event);
+}
+
+static void ggml_backend_et_buffer_get_tensor(ggml_backend_buffer_t buffer,
+ const ggml_tensor * tensor,
+ void * data,
+ size_t offset,
+ size_t size) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ return;
+ }
+
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) buffer->buft->device->context;
+ rt::StreamId stream = dev_ctx->default_stream;
+
+ const std::byte * src_ptr = static_cast<const std::byte *>(tensor->data) + offset;
+ std::byte * dst_ptr = static_cast<std::byte *>(data);
+
+ rt::EventId event = runtime->memcpyDeviceToHost(stream, src_ptr, dst_ptr, size, true /*barrier*/);
+
+ runtime->waitForEvent(event);
+}
+
+static bool ggml_backend_et_buffer_cpy_tensor(ggml_backend_buffer_t buffer,
+ const ggml_tensor * src,
+ ggml_tensor * dst) {
+ GGML_UNUSED(buffer);
+ GGML_UNUSED(src);
+ GGML_UNUSED(dst);
+ return false;
+}
+
+static void ggml_backend_et_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
+ ggml_backend_et_buffer_context * ctx = (ggml_backend_et_buffer_context *) buffer->context;
+
+ if (ctx->size == 0 || ctx->data == nullptr) {
+ return;
+ }
+
+ // Get device context to access memops kernel
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) buffer->buft->device->context;
+ if (!dev_ctx) {
+ GGML_LOG_ERROR("ET: Failed to get device context for buffer clear");
+ return;
+ }
+
+ // Use device-side memset kernel for efficient clearing
+ if (!ggml_et_memset(dev_ctx, ctx->data, value, ctx->size)) {
+ GGML_LOG_ERROR("ET: buffer_clear failed using memset kernel");
+ return;
+ }
+
+ GGML_LOG_DEBUG("ET: Buffer cleared successfully using memops kernel");
+}
+
+static const struct ggml_backend_buffer_i ggml_backend_et_buffer_i = {
+ /* .free_buffer = */ ggml_backend_et_buffer_free_buffer,
+ /* .get_base = */ ggml_backend_et_buffer_get_base,
+ /* .init_tensor = */ ggml_backend_et_buffer_init_tensor,
+ /* .memset_tensor = */ NULL,
+ /* .set_tensor = */ ggml_backend_et_buffer_set_tensor,
+ /* .get_tensor = */ ggml_backend_et_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
+ /* .cpy_tensor = */ ggml_backend_et_buffer_cpy_tensor,
+ /* .clear = */ ggml_backend_et_buffer_clear,
+ /* .reset = */ NULL,
+};
+
+static const char * ggml_backend_et_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
+ GGML_UNUSED(buft);
+ return GGML_ET_NAME;
+}
+
+static ggml_backend_buffer_t ggml_backend_et_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
+ ggml_backend_et_buffer_type_context * btctx = (ggml_backend_et_buffer_type_context *) buft->context;
+
+ ggml_backend_et_buffer_context * ctx = new ggml_backend_et_buffer_context;
+ ctx->devidx = btctx->devidx;
+ ctx->size = size;
+
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ delete ctx;
+ return nullptr;
+ }
+
+ std::vector<rt::DeviceId> rtids = runtime->getDevices();
+ if (static_cast<size_t>(btctx->devidx) >= rtids.size()) {
+ delete ctx;
+ return nullptr;
+ }
+ ctx->rtid = rtids[btctx->devidx];
+
+ ctx->data = runtime->mallocDevice(ctx->rtid, size);
+ if (ctx->data == nullptr) {
+ delete ctx;
+ return nullptr;
+ }
+
+ return ggml_backend_buffer_init(buft, ggml_backend_et_buffer_i, ctx, size);
+}
+
+static size_t ggml_backend_et_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime || !buft->device) {
+ return GGML_MEM_ALIGN;
+ }
+
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) buft->device->context;
+ rt::DeviceProperties prop = runtime->getDeviceProperties(dev_ctx->rtid);
+ return prop.cacheLineSize_;
+}
+
+static size_t ggml_backend_et_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
+ if (buft->device) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) buft->device->context;
+ return dev_ctx->total_mem;
+ }
+ return SIZE_MAX;
+}
+
+static size_t ggml_backend_et_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
+ GGML_UNUSED(buft);
+ return ggml_nbytes_pad(tensor);
+}
+
+static bool ggml_backend_et_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
+ GGML_UNUSED(buft);
+ return false;
+}
+
+static const struct ggml_backend_buffer_type_i ggml_backend_et_buffer_type_i = {
+ /* .get_name = */ ggml_backend_et_buffer_type_get_name,
+ /* .alloc_buffer = */ ggml_backend_et_buffer_type_alloc_buffer,
+ /* .get_alignment = */ ggml_backend_et_buffer_type_get_alignment,
+ /* .get_max_size = */ ggml_backend_et_buffer_type_get_max_size,
+ /* .get_alloc_size = */ ggml_backend_et_buffer_type_get_alloc_size,
+ /* .is_host = */ ggml_backend_et_buffer_type_is_host,
+};
+
+static const char * ggml_backend_et_get_name(ggml_backend_t backend) {
+ GGML_UNUSED(backend);
+ return GGML_ET_NAME;
+}
+
+static void ggml_backend_et_free(ggml_backend_t backend) {
+ ggml_backend_et_context * et_ctx = (ggml_backend_et_context *) backend->context;
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+
+ // Clean up kernels on this device before freeing backend
+ ggml_backend_dev_t dev = ggml_backend_et_reg_get_device(ggml_backend_et_reg(), et_ctx->devidx);
+ if (dev && dev->context) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+
+ if (_drv.profiling_enabled) {
+ auto kernels = ggml_et_get_loaded_kernels(dev_ctx);
+ _drv.kernel_map.insert(_drv.kernel_map.end(), kernels.begin(), kernels.end());
+ }
+
+ ggml_et_unload_all_kernels(dev_ctx);
+
+ if (runtime) {
+ if (dev_ctx->trace_buffer) {
+ runtime->freeDevice(dev_ctx->rtid, dev_ctx->trace_buffer);
+ dev_ctx->trace_buffer = nullptr;
+ }
+ // Drain any in-flight uberkernel launches before freeing the
+ // device buffers they read from.
+ runtime->waitForStream(dev_ctx->default_stream);
+ for (auto & slot : dev_ctx->uberkernel.slots) {
+ if (slot.device_insts) {
+ runtime->freeDevice(dev_ctx->rtid, slot.device_insts);
+ slot.device_insts = nullptr;
+ }
+ if (slot.device_params) {
+ runtime->freeDevice(dev_ctx->rtid, slot.device_params);
+ slot.device_params = nullptr;
+ }
+ slot.has_pending = false;
+ }
+ }
+ }
+
+ delete et_ctx;
+ delete backend;
+}
+
+static ggml_backend_buffer_type_t ggml_backend_et_get_default_buffer_type(ggml_backend_t backend) {
+ ggml_backend_et_context * et_ctx = (ggml_backend_et_context *) backend->context;
+
+ return ggml_backend_et_buffer_type(et_ctx->devidx);
+}
+
+static void ggml_backend_et_set_tensor_async(ggml_backend_t backend,
+ ggml_tensor * tensor,
+ const void * data,
+ size_t offset,
+ size_t size) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ return;
+ }
+
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) backend->device->context;
+ rt::StreamId stream = dev_ctx->default_stream;
+
+ std::byte * dst_ptr = static_cast<std::byte *>(tensor->data) + offset;
+ const std::byte * src_ptr = static_cast<const std::byte *>(data);
+
+ runtime->memcpyHostToDevice(stream, src_ptr, dst_ptr, size, true /*barrier*/);
+}
+
+static void ggml_backend_et_get_tensor_async(ggml_backend_t backend,
+ const ggml_tensor * tensor,
+ void * data,
+ size_t offset,
+ size_t size) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ return;
+ }
+
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) backend->device->context;
+ rt::StreamId stream = dev_ctx->default_stream;
+
+ const std::byte * src_ptr = static_cast<const std::byte *>(tensor->data) + offset;
+ std::byte * dst_ptr = static_cast<std::byte *>(data);
+
+ runtime->memcpyDeviceToHost(stream, src_ptr, dst_ptr, size, true /*barrier*/);
+}
+
+static bool ggml_backend_et_cpy_tensor_async(ggml_backend_t backend_src,
+ ggml_backend_t backend_dst,
+ const ggml_tensor * src,
+ ggml_tensor * dst) {
+ GGML_UNUSED(backend_src);
+ GGML_UNUSED(backend_dst);
+ GGML_UNUSED(src);
+ GGML_UNUSED(dst);
+ return false;
+}
+
+static void ggml_backend_et_synchronize(ggml_backend_t backend) {
+ std::shared_ptr<rt::IRuntime> runtime = ggml_et_runtime();
+ if (!runtime) {
+ return;
+ }
+
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) backend->device->context;
+ runtime->waitForStream(dev_ctx->default_stream);
+
+ auto errors = runtime->retrieveStreamErrors(dev_ctx->default_stream);
+ if (errors.empty()) {
+ return;
+ }
+ for (const auto & err : errors) {
+ GGML_LOG_ERROR("ET: stream error detected at synchronization point. Code: %d,Type: %d\n", (int) err.errorCode_,
+ (int) err.errorContext_.value()[0].type_);
+ }
+ abort();
+}
+
+static bool ggml_et_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list<ggml_op> ops) {
+ if (!ggml_can_fuse(cgraph, node_idx, ops)) {
+ return false;
+ }
+
+ if (ops.size() == 2 && ops.begin()[0] == GGML_OP_MUL_MAT && ops.begin()[1] == GGML_OP_ADD) {
+ const ggml_tensor * mm = cgraph->nodes[node_idx];
+ const ggml_tensor * add = cgraph->nodes[node_idx + 1];
+
+ // Only Q8_0 weights x F32 activations -> F32 (the kernel that has
+ // the bias path). Other MM variants must wait for their own kernel
+ // bias support.
+ if (mm->type != GGML_TYPE_F32 || mm->src[0]->type != GGML_TYPE_Q8_0 || mm->src[1]->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ // ADD must be F32 and one of its operands must be the MM output.
+ if (add->type != GGML_TYPE_F32) {
+ return false;
+ }
+ if (add->src[0] != mm && add->src[1] != mm) {
+ return false;
+ }
+
+ const ggml_tensor * bias = (add->src[0] == mm) ? add->src[1] : add->src[0];
+
+ if (bias->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ // No broadcasting: bias shape must equal MM output shape.
+ for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+ if (bias->ne[i] != mm->ne[i]) {
+ return false;
+ }
+ }
+
+ // Bias and dst must be contiguous and have identical strides - the
+ // kernel uses dst-style offset arithmetic against bias's nb[].
+ if (!ggml_is_contiguous(bias) || !ggml_is_contiguous(mm)) {
+ return false;
+ }
+ for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+ if ((int64_t) bias->nb[i] != (int64_t) add->nb[i]) {
+ return false;
+ }
+ }
+ }
+
+ if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) {
+ const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
+ const ggml_tensor * mul = cgraph->nodes[node_idx + 1];
+
+ // ET only supports F32
+ if (rms_norm->src[0]->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ // Identify the weights tensor (the MUL operand that isn't rms_norm output)
+ const ggml_tensor * weights = (mul->src[0] == rms_norm) ? mul->src[1] : mul->src[0];
+
+ if (weights->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ // Both inputs must be contiguous (ET hardware requirement)
+ if (!ggml_is_contiguous(rms_norm->src[0]) || !ggml_is_contiguous_rows(weights)) {
+ return false;
+ }
+
+ // ET requires cache-aligned rows (ne[0] % 16 == 0)
+ if (rms_norm->src[0]->ne[0] % 16 != 0 || weights->ne[0] % 16 != 0) {
+ return false;
+ }
+
+ // Fused kernel doesn't handle dim-0 broadcasting
+ if (weights->ne[0] != rms_norm->src[0]->ne[0]) {
+ return false;
+ }
+ }
+
+ return true;
+}
+
+static ggml_status ggml_backend_et_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) backend->device->context;
+ ggml_et_uberkernel_begin_graph(&dev_ctx->uberkernel);
+
+ for (int i = 0; i < cgraph->n_nodes; i++) {
+ ggml_tensor * node = cgraph->nodes[i];
+
+ if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+ node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+ continue;
+ }
+
+ // --- Fusion checks (before regular dispatch) ---
+ if (ggml_et_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) {
+ ggml_et_op_rms_norm_mul(dev_ctx, node, cgraph->nodes[i + 1]);
+ i++; // skip the MUL node
+ continue;
+ }
+ if (ggml_et_can_fuse(cgraph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) {
+ ggml_et_op_mul_mat(dev_ctx, node, cgraph->nodes[i + 1]);
+ i++; // skip the ADD node
+ continue;
+ }
+
+ switch (node->op) {
+ case GGML_OP_SQR:
+ ggml_et_op_sqr(dev_ctx, node);
+ break;
+
+ case GGML_OP_UNARY:
+ ggml_et_op_unary(dev_ctx, node);
+ break;
+
+ case GGML_OP_SUM_ROWS:
+ ggml_et_op_sum_rows(dev_ctx, node);
+ break;
+
+ case GGML_OP_MEAN:
+ ggml_et_op_mean(dev_ctx, node);
+ break;
+
+ case GGML_OP_CLAMP:
+ ggml_et_op_clamp(dev_ctx, node);
+ break;
+
+ case GGML_OP_MUL:
+ ggml_et_op_mul(dev_ctx, node);
+ break;
+
+ case GGML_OP_ADD:
+ ggml_et_op_add(dev_ctx, node);
+ break;
+
+ case GGML_OP_SUB:
+ ggml_et_op_sub(dev_ctx, node);
+ break;
+
+ case GGML_OP_CUMSUM:
+ ggml_et_op_cumsum(dev_ctx, node);
+ break;
+
+ case GGML_OP_MUL_MAT:
+ ggml_et_op_mul_mat(dev_ctx, node);
+ break;
+
+ case GGML_OP_MUL_MAT_ID:
+ ggml_et_op_mul_mat_id(dev_ctx, node);
+ break;
+
+ case GGML_OP_ROPE:
+ ggml_et_op_rope(dev_ctx, node);
+ break;
+
+ case GGML_OP_RMS_NORM:
+ ggml_et_op_rms_norm(dev_ctx, node);
+ break;
+
+ case GGML_OP_NORM:
+ ggml_et_op_norm(dev_ctx, node);
+ break;
+
+ case GGML_OP_L2_NORM:
+ ggml_et_op_l2_norm(dev_ctx, node);
+ break;
+
+ case GGML_OP_GROUP_NORM:
+ ggml_et_op_group_norm(dev_ctx, node);
+ break;
+
+ case GGML_OP_SCALE:
+ ggml_et_op_scale(dev_ctx, node);
+ break;
+
+ case GGML_OP_GLU:
+ ggml_et_op_glu(dev_ctx, node);
+ break;
+
+ case GGML_OP_SOFT_MAX:
+ ggml_et_op_softmax(dev_ctx, node);
+ break;
+
+ case GGML_OP_IM2COL:
+ ggml_et_op_im2col(dev_ctx, node);
+ break;
+
+ case GGML_OP_CONV_2D:
+ ggml_et_op_conv_2d(dev_ctx, node);
+ break;
+
+ case GGML_OP_FLASH_ATTN_EXT:
+ ggml_et_op_flash_attn_ext(dev_ctx, node);
+ break;
+
+ case GGML_OP_GET_ROWS:
+ ggml_et_op_get_rows(dev_ctx, node);
+ break;
+
+ case GGML_OP_CONT:
+ ggml_et_op_cont(dev_ctx, node);
+ break;
+
+ case GGML_OP_CPY:
+ ggml_et_op_cpy(dev_ctx, node);
+ break;
+
+ case GGML_OP_CONCAT:
+ ggml_et_op_concat(dev_ctx, node);
+ break;
+
+ case GGML_OP_REPEAT:
+ ggml_et_op_repeat(dev_ctx, node);
+ break;
+
+ case GGML_OP_SSM_CONV:
+ ggml_et_op_ssm_conv(dev_ctx, node);
+ break;
+
+ case GGML_OP_SSM_SCAN:
+ ggml_et_op_ssm_scan(dev_ctx, node);
+ break;
+
+ case GGML_OP_PAD:
+ ggml_et_op_pad(dev_ctx, node);
+ break;
+
+ case GGML_OP_SET_ROWS:
+ ggml_et_op_set_rows(dev_ctx, node);
+ break;
+
+ case GGML_OP_FILL:
+ ggml_et_op_fill(dev_ctx, node);
+ break;
+
+ case GGML_OP_DIAG:
+ ggml_et_op_diag(dev_ctx, node);
+ break;
+
+ case GGML_OP_TRI:
+ ggml_et_op_tri(dev_ctx, node);
+ break;
+
+ case GGML_OP_SOLVE_TRI:
+ ggml_et_op_solve_tri(dev_ctx, node);
+ break;
+
+ case GGML_OP_SET:
+ ggml_et_op_set(dev_ctx, node);
+ break;
+
+ case GGML_OP_RWKV_WKV6:
+ ggml_et_op_rwkv_wkv6(dev_ctx, node);
+ break;
+
+ case GGML_OP_RWKV_WKV7:
+ ggml_et_op_rwkv_wkv7(dev_ctx, node);
+ break;
+
+ case GGML_OP_GATED_DELTA_NET:
+ ggml_et_op_gated_delta_net(dev_ctx, node);
+ break;
+
+ default:
+ ggml_et_uberkernel_abort_graph(&dev_ctx->uberkernel);
+ GGML_LOG_ERROR("ET: Unsupported operation in graph: %s", ggml_op_name(node->op));
+ return GGML_STATUS_FAILED;
+ }
+
+ if (ggml_et_uberkernel_failed(&dev_ctx->uberkernel)) {
+ ggml_et_uberkernel_abort_graph(&dev_ctx->uberkernel);
+ return GGML_STATUS_FAILED;
+ }
+ }
+
+ if (!ggml_et_uberkernel_end_graph(dev_ctx)) {
+ ggml_et_uberkernel_abort_graph(&dev_ctx->uberkernel);
+ return GGML_STATUS_FAILED;
+ }
+
+ return GGML_STATUS_SUCCESS;
+}
+
+// Check that elements within each row are contiguous (nb[0] == type_size).
+// Higher-dim strides can be arbitrary - kernels navigate them via byte offsets.
+static bool et_ggml_is_row_contiguous(const ggml_tensor * t) {
+ return t->nb[0] == ggml_type_size(t->type);
+}
+
+static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
+ GGML_UNUSED(dev);
+
+ bool supported = false;
+ switch (op->op) {
+ case GGML_OP_CUMSUM:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->src[0]->nb[0] == sizeof(float) && ggml_is_contiguous(op);
+ break;
+ case GGML_OP_SQR:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->ne[0] % 16 == 0 && ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]);
+ break;
+ case GGML_OP_SUM_ROWS:
+ // dst has ne[0]=1, src0 row length must be cache-aligned
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->src[0]->ne[0] % 16 == 0 && ggml_is_contiguous(op->src[0]);
+ break;
+ case GGML_OP_MEAN:
+ // Kernel handles arbitrary ne00 (per-row alignment guard with
+ // scalar tail), so no row-length divisibility constraint here.
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op->src[0]);
+ break;
+ case GGML_OP_CLAMP:
+ // Element-wise; kernel distributes by cache lines and handles a
+ // scalar tail, so any contiguous F32 size is fine - including the
+ // 1x1x1x1 scalar case.
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]);
+ break;
+ case GGML_OP_UNARY:
+ // Only require dim-0 contiguity (nb[0] == sizeof(float)). Higher
+ // dims may be arbitrarily strided views; the kernel walks per-row
+ // using all four nb[] values. See unary_f32.c entry_point.
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_nelements(op) % 16 == 0 && op->nb[0] == sizeof(float) && op->src[0]->nb[0] == sizeof(float)) {
+ switch (ggml_get_unary_op(op)) {
+ case GGML_UNARY_OP_ABS:
+ case GGML_UNARY_OP_SGN:
+ case GGML_UNARY_OP_NEG:
+ case GGML_UNARY_OP_STEP:
+ case GGML_UNARY_OP_TANH:
+ case GGML_UNARY_OP_ELU:
+ case GGML_UNARY_OP_RELU:
+ case GGML_UNARY_OP_SIGMOID:
+ case GGML_UNARY_OP_GELU:
+ case GGML_UNARY_OP_GELU_QUICK:
+ case GGML_UNARY_OP_SILU:
+ case GGML_UNARY_OP_HARDSWISH:
+ case GGML_UNARY_OP_HARDSIGMOID:
+ case GGML_UNARY_OP_EXP:
+ case GGML_UNARY_OP_EXPM1:
+ case GGML_UNARY_OP_SOFTPLUS:
+ case GGML_UNARY_OP_GELU_ERF:
+ case GGML_UNARY_OP_FLOOR:
+ case GGML_UNARY_OP_CEIL:
+ case GGML_UNARY_OP_ROUND:
+ case GGML_UNARY_OP_TRUNC:
+ supported = true;
+ break;
+ default:
+ break;
+ }
+ }
+ break;
+ case GGML_OP_MUL:
+ case GGML_OP_ADD:
+ case GGML_OP_SUB:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && op->nb[0] == sizeof(float) &&
+ op->src[0]->nb[0] == sizeof(float) &&
+ (op->src[1]->nb[0] == sizeof(float) || op->src[1]->ne[0] == 1) &&
+ op->nb[1] == op->ne[0] * sizeof(float);
+ break;
+ case GGML_OP_MUL_MAT:
+ // Support Q8_0 x F32 -> F32, F16 x F32 -> F32, F16 x F16 -> F32, and F32 x F32 -> F32 matrix multiplication
+ // Stride requirements: first dimension must be contiguous for all tensors
+ if (op->type == GGML_TYPE_F32 &&
+ ((op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32) ||
+ (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F16)) &&
+ op->ne[0] % 16 == 0 && // dst row length for tensor-store path
+ op->src[0]->ne[1] % 16 == 0 && // m
+ op->src[0]->ne[0] % 16 == 0 && // k
+ ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) {
+ // Special path for the FP32 TensorFMA kernel
+ // Limitation - generic kernels can tolerate non-cache-aligned dst rows
+ // because they publish each output element atomically. The matrix
+ // engine path still uses tiled tensor stores, so keep dst rows aligned.
+ // The m edge is difficult to do because of the 4 conseqtive load hardware limitation
+ // And the k edge is impossible because that is encoded as `stride & 0xFFFFFFFFFFC0ULL` which becomes 0 for stride 16 (4x FP32) :(
+ // FIXME: Right now this overwrites the mul_mat_f32 kernel - whatever. Fix later. Demo code
+ supported = true;
+ } else if (op->type == GGML_TYPE_F32 && op->src[0] &&
+ (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32) && op->src[1] &&
+ (op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32)) {
+ // Check first dimension contiguity requirements
+ bool src0_first_dim_contiguous = (op->src[0]->nb[0] == ggml_type_size(op->src[0]->type));
+ bool src1_first_dim_contiguous = (op->src[1]->nb[0] == ggml_type_size(op->src[1]->type));
+ bool dst_first_dim_contiguous = (op->nb[0] == sizeof(float));
+
+ // Check destination stride ordering (only for dimensions with ne > 1)
+ bool dst_properly_ordered = true;
+ for (int d = 0; d < 3; d++) {
+ if (op->ne[d] > 1 && op->ne[d + 1] > 1 && op->nb[d] > op->nb[d + 1]) {
+ dst_properly_ordered = false;
+ }
+ }
+
+ supported = src0_first_dim_contiguous && src1_first_dim_contiguous && dst_first_dim_contiguous &&
+ dst_properly_ordered;
+ } else if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_Q8_0 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32) {
+ // Keep the existing quantized path constraints separate from the
+ // relaxed non-quant generic fallback.
+ bool src0_first_dim_contiguous = (op->src[0]->nb[0] == ggml_type_size(op->src[0]->type));
+ bool src1_first_dim_contiguous = (op->src[1]->nb[0] == ggml_type_size(op->src[1]->type));
+ bool dst_first_dim_contiguous = (op->nb[0] == sizeof(float));
+
+ bool dst_properly_ordered = true;
+ for (int d = 0; d < 3; d++) {
+ if (op->ne[d] > 1 && op->ne[d + 1] > 1 && op->nb[d] > op->nb[d + 1]) {
+ dst_properly_ordered = false;
+ }
+ }
+
+ supported = src0_first_dim_contiguous && src1_first_dim_contiguous && dst_first_dim_contiguous &&
+ dst_properly_ordered;
+
+ } else if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_Q4_0 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32) {
+ // Keep the existing quantized path constraints separate from the
+ // relaxed non-quant generic fallback.
+ bool src0_first_dim_contiguous = (op->src[0]->nb[0] == ggml_type_size(op->src[0]->type));
+ bool src1_first_dim_contiguous = (op->src[1]->nb[0] == ggml_type_size(op->src[1]->type));
+ bool dst_first_dim_contiguous = (op->nb[0] == sizeof(float));
+
+ bool dst_properly_ordered = true;
+ for (int d = 0; d < 3; d++) {
+ if (op->ne[d] > 1 && op->ne[d + 1] > 1 && op->nb[d] > op->nb[d + 1]) {
+ dst_properly_ordered = false;
+ }
+ }
+
+ supported = src0_first_dim_contiguous && src1_first_dim_contiguous && dst_first_dim_contiguous &&
+ dst_properly_ordered;
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_MUL_MAT_ID:
+ // Support MUL_MAT_ID for Mixture of Experts: (Q8_0/Q4_0/F16/F32) x F32 -> F32 with I32 expert indices
+ // src0 (as): [K, M, n_expert] - expert weight matrices (can be quantized)
+ // src1 (b): [K, n_expert_used, batch] - activations (F32)
+ // src2 (ids): [n_expert_used, batch] - expert selection indices (I32)
+ // dst: [M, n_expert_used, batch, 1] - output (F32)
+ if (op->type == GGML_TYPE_F32 && op->src[0] &&
+ (op->src[0]->type == GGML_TYPE_Q8_0 || op->src[0]->type == GGML_TYPE_Q4_0 ||
+ op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32) &&
+ op->src[1] && op->src[1]->type == GGML_TYPE_F32 && op->src[2] && op->src[2]->type == GGML_TYPE_I32) {
+ // Check first dimension contiguity requirements (matching CPU backend)
+ bool src0_first_dim_contiguous = (op->src[0]->nb[0] == ggml_type_size(op->src[0]->type));
+ bool src1_first_dim_contiguous = (op->src[1]->nb[0] == ggml_type_size(op->src[1]->type));
+ bool src2_first_dim_contiguous = (op->src[2]->nb[0] == ggml_type_size(op->src[2]->type));
+ bool dst_first_dim_contiguous = (op->nb[0] == sizeof(float));
+
+ // Check destination stride ordering (only for dimensions with ne > 1)
+ bool dst_properly_ordered = true;
+ for (int d = 0; d < 3; d++) {
+ if (op->ne[d] > 1 && op->ne[d + 1] > 1 && op->nb[d] > op->nb[d + 1]) {
+ dst_properly_ordered = false;
+ }
+ }
+
+ // Validate tensor dimension constraints from GGML definition
+ bool dims_valid = (op->src[0]->ne[3] == 1) && // as is 3d (one matrix per expert)
+ (op->src[1]->ne[3] == 1) && // b is 3d
+ (op->src[2]->ne[2] == 1 && op->src[2]->ne[3] == 1) && // ids is 2d
+ (op->src[2]->ne[1] == op->src[1]->ne[2]) && // must have expert list per b row
+ (op->src[0]->ne[0] == op->src[1]->ne[0]) && // K dimension must match
+ (op->src[2]->ne[0] % op->src[1]->ne[1] == 0); // can broadcast
+
+ supported = src0_first_dim_contiguous && src1_first_dim_contiguous && src2_first_dim_contiguous &&
+ dst_first_dim_contiguous && dst_properly_ordered && dims_valid;
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_ROPE:
+ // Support F32 x I32 -> F32 RoPE for the modes implemented by rope_f32.
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_I32 && ggml_is_contiguous(op) && et_ggml_is_row_contiguous(op->src[0])) {
+ const int mode = ggml_get_op_params_i32(op, 2);
+ const int ndims = ggml_get_op_params_i32(op, 1);
+ const bool is_normal = mode == GGML_ROPE_TYPE_NORMAL;
+ const bool is_neox = mode == GGML_ROPE_TYPE_NEOX;
+ const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE;
+ const bool zero_view_offset = op->src[0]->view_src == nullptr || op->src[0]->view_offs == 0;
+ const bool has_sections = ggml_get_op_params_i32(op, 11) > 0 || ggml_get_op_params_i32(op, 12) > 0 ||
+ ggml_get_op_params_i32(op, 13) > 0;
+
+ supported =
+ zero_view_offset && ndims <= 512 &&
+ (is_normal || (is_neox && ndims % 16 == 0) || (is_imrope && ndims % 16 == 0 && has_sections));
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_RMS_NORM:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->ne[0] % 16 == 0 && ggml_is_contiguous(op) && et_ggml_is_row_contiguous(op->src[0]);
+ break;
+ case GGML_OP_NORM:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->ne[0] % 16 == 0 && ggml_is_contiguous(op) && et_ggml_is_row_contiguous(op->src[0]);
+ break;
+ case GGML_OP_L2_NORM:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->ne[0] % 16 == 0 && ggml_is_contiguous(op) && et_ggml_is_row_contiguous(op->src[0]);
+ break;
+ case GGML_OP_GROUP_NORM:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && et_ggml_is_row_contiguous(op->src[0]) &&
+ ggml_get_op_params_i32(op, 0) > 0;
+ break;
+ case GGML_OP_IM2COL:
+ supported = op->src[0] && op->src[1] &&
+ ((op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32) ||
+ (op->type == GGML_TYPE_F16 &&
+ (op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32))) &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[1]) &&
+ op->nb[0] == ggml_type_size(op->type) && op->src[1]->nb[0] == ggml_type_size(op->src[1]->type);
+ break;
+ case GGML_OP_CONV_2D:
+ {
+ // First-cut conv_2d_f32_me kernel constraints. Anything outside
+ // this falls back to CPU (it's a strict subset on purpose).
+ if (!op->src[0] || !op->src[1]) {
+ supported = false;
+ break;
+ }
+ if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 ||
+ op->src[1]->type != GGML_TYPE_F32) {
+ supported = false;
+ break;
+ }
+ if (!ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0]) || !ggml_is_contiguous(op->src[1])) {
+ supported = false;
+ break;
+ }
+
+ const ggml_tensor * flt = op->src[0]; // [Kw, Kh, Cin, Cout]
+ const ggml_tensor * in = op->src[1]; // [W, H, Cin, N]
+ const int32_t s0 = ggml_get_op_params_i32(op, 0);
+ const int32_t s1 = ggml_get_op_params_i32(op, 1);
+ const int32_t p0 = ggml_get_op_params_i32(op, 2);
+ const int32_t p1 = ggml_get_op_params_i32(op, 3);
+ const int32_t d0 = ggml_get_op_params_i32(op, 4);
+ const int32_t d1 = ggml_get_op_params_i32(op, 5);
+
+ const int64_t Kw = flt->ne[0];
+ const int64_t Kh = flt->ne[1];
+ const int64_t Cin = flt->ne[2];
+ const int64_t Cout = flt->ne[3];
+ const int64_t H = in->ne[1];
+ (void) in->ne[0];
+
+ if (s0 < 1 || s1 < 1 || !(d0 == 1 && d1 == 1) || Cin % 16 != 0 || Cout % 16 != 0 || in->ne[3] != 1) {
+ supported = false;
+ break;
+ }
+ const int64_t OW = op->ne[0];
+ const int64_t OH = op->ne[1];
+ if (OW <= 0 || OH <= 0) {
+ supported = false;
+ break;
+ }
+ (void) p0;
+ (void) p1;
+
+ // Mirror the kernel's sizing:
+ // if K_TILES * per_KT_bytes <= budget: 1 buffer, n_chunks=1
+ // else: 2 buffers (double-buffer), shrink chunk_KT until
+ // 2*chunk_KT*per_KT_bytes <= budget.
+ const int64_t Hp = H + 2 * p1;
+ const int64_t OW_pad = (OW + 15) & ~15;
+ const int64_t Wp_a = OW_pad;
+ const bool need_stage = (OW % 16 != 0);
+ const int64_t stage_bytes = need_stage ? (Cout * OH * OW_pad * 4) : 0;
+ const int64_t L2SCP_BUDGET = 1500 * 1024;
+ // Per-hart partial-TenC scratch (mirrors kernel MAX_TILES_PER_HART=2):
+ // 32 minions x 2 tiles x 1024 bytes = 64 KB per shire.
+ const int64_t scratch_bytes = 32 * 2 * 16 * 16 * 4;
+ const int64_t budget = L2SCP_BUDGET - stage_bytes - scratch_bytes;
+ const int64_t per_KT_bytes = Kh * Kw * Cout * 16 * 4 + Kw * 16 * Hp * Wp_a * 4;
+ const int64_t K_TILES = Cin / 16;
+
+ int64_t chunk_KT_calc;
+ int64_t n_chunks_calc;
+ if (K_TILES * per_KT_bytes <= budget) {
+ chunk_KT_calc = K_TILES;
+ n_chunks_calc = 1;
+ } else {
+ chunk_KT_calc = K_TILES;
+ while (chunk_KT_calc > 1 && 2 * chunk_KT_calc * per_KT_bytes > budget) {
+ chunk_KT_calc--;
+ }
+ while (chunk_KT_calc > 1 && K_TILES % chunk_KT_calc != 0) {
+ chunk_KT_calc--;
+ }
+ if (chunk_KT_calc < 1) {
+ supported = false;
+ break;
+ }
+ n_chunks_calc = K_TILES / chunk_KT_calc;
+ }
+
+ if (n_chunks_calc > 1) {
+ const int64_t M_TILES = Cout / 16;
+ const int64_t w_tiles = (OW + 15) / 16;
+ const int64_t total_tiles = OH * w_tiles * M_TILES;
+ // MAX_TILES_PER_HART = 2 (mirrors kernel constant).
+ const int64_t max_workers = (need_stage ? 32 : 1024) * 2;
+ if (total_tiles > max_workers) {
+ supported = false;
+ break;
+ }
+ }
+
+ supported = true;
+ break;
+ }
+ case GGML_OP_SCALE:
+ // F32 contiguous, total elements must be cache line aligned (16 floats)
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) && (ggml_nelements(op) % 16 == 0);
+ break;
+ case GGML_OP_GLU:
+ // Note: we only require row-wise contiguity (ggml_is_contiguous_1) so that
+ // strided views over a packed up_proj tensor (the common split-GLU layout)
+ // are accepted. The kernel walks rows via nb[1] strides, so the inner
+ // dimension just needs to be densely packed.
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_nelements(op) % 16 == 0 && ggml_is_contiguous_1(op) && ggml_is_contiguous_1(op->src[0])) {
+ // Check GLU variant - support SWIGLU, SWIGLU_OAI, GEGLU, GEGLU_ERF, GEGLU_QUICK, REGLU
+ ggml_glu_op glu_type = ggml_get_glu_op(op);
+ const bool supported_variant = glu_type == GGML_GLU_OP_SWIGLU || glu_type == GGML_GLU_OP_SWIGLU_OAI ||
+ glu_type == GGML_GLU_OP_GEGLU || glu_type == GGML_GLU_OP_GEGLU_ERF ||
+ glu_type == GGML_GLU_OP_GEGLU_QUICK || glu_type == GGML_GLU_OP_REGLU;
+
+ if (op->src[1]) {
+ supported = supported_variant && op->src[1]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous_1(op->src[1]) && op->src[0]->ne[0] == op->ne[0] &&
+ op->src[1]->ne[0] == op->ne[0];
+ } else {
+ supported = supported_variant && op->src[0]->ne[0] == 2 * op->ne[0];
+ }
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_SOFT_MAX:
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) && op->src[0]->ne[0] > 1) {
+ // Check optional mask tensor (F32 only)
+ if (op->src[1]) {
+ supported = op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[1]);
+ if (!supported) {
+ break;
+ }
+ }
+ // Check optional sinks tensor (F32 only)
+ if (op->src[2]) {
+ supported = op->src[2]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[2]);
+ } else {
+ supported = true;
+ }
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_SSM_SCAN:
+ supported = op->type == GGML_TYPE_F32 && ggml_is_contiguous(op) && op->src[0] &&
+ op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && op->src[2] && op->src[2]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op->src[2]) && op->src[3] && op->src[3]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op->src[3]) && op->src[4] && op->src[4]->type == GGML_TYPE_F32 &&
+ op->src[5] && op->src[5]->type == GGML_TYPE_F32 && op->src[6] &&
+ op->src[6]->type == GGML_TYPE_I32 && ggml_is_contiguous(op->src[6]) &&
+ op->src[1]->nb[0] == sizeof(float) && op->src[4]->nb[0] == sizeof(float) &&
+ op->src[5]->nb[0] == sizeof(float) &&
+ op->src[1]->nb[1] == (size_t) op->src[1]->ne[0] * sizeof(float) &&
+ op->src[4]->nb[1] == (size_t) op->src[4]->ne[0] * sizeof(float) &&
+ op->src[5]->nb[1] == (size_t) op->src[5]->ne[0] * sizeof(float) &&
+ op->src[0]->ne[0] == op->src[4]->ne[0] && op->src[0]->ne[1] == op->src[1]->ne[0] &&
+ op->src[0]->ne[2] == op->src[1]->ne[1] && op->src[1]->ne[2] == op->src[2]->ne[1] &&
+ op->src[1]->ne[3] == op->src[2]->ne[2] && op->src[4]->ne[2] == op->src[1]->ne[2] &&
+ op->src[4]->ne[3] == op->src[1]->ne[3] && ggml_are_same_shape(op->src[4], op->src[5]) &&
+ op->src[6]->ne[0] == op->src[1]->ne[3] && op->src[3]->ne[1] == op->src[1]->ne[1] &&
+ (op->src[3]->ne[0] == 1 || op->src[3]->ne[0] == op->src[0]->ne[0]) &&
+ (op->src[1]->ne[1] % op->src[4]->ne[1] == 0);
+ break;
+ case GGML_OP_FLASH_ATTN_EXT:
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] &&
+ (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr &&
+ ggml_is_contiguous_rows(op) && ggml_is_contiguous_rows(op->src[0])) {
+ float max_bias = 0.0f;
+ float logit_softcap = 0.0f;
+ memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(max_bias));
+ memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(logit_softcap));
+
+ const ggml_prec prec = ggml_flash_attn_ext_get_prec(op);
+
+ // Mask must be F16 or F32 if present
+ bool mask_ok = (op->src[3] == nullptr) || (op->src[3]->type == GGML_TYPE_F32) ||
+ (op->src[3]->type == GGML_TYPE_F16);
+
+ // GQA: n_head_q must be a multiple of n_head_kv
+ const int64_t nhq = op->src[0]->ne[2];
+ const int64_t nhk = op->src[1]->ne[2];
+
+ // K/V row stride must match element size
+ const size_t k_elem = op->src[1]->type == GGML_TYPE_F16 ? 2 : 4;
+ const size_t v_elem = op->src[2]->type == GGML_TYPE_F16 ? 2 : 4;
+
+ // Only support matrix engine path (F16 K/V, dk%32==0);
+ // mask scalar F32 fallback to get baseline perf readings
+ const bool me_eligible = op->src[1]->type == GGML_TYPE_F16 && op->src[2]->type == GGML_TYPE_F16 &&
+ (op->src[0]->ne[0] % 32) == 0;
+
+ supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_DEFAULT) &&
+ max_bias == 0.0f && logit_softcap == 0.0f && op->src[0]->nb[0] == sizeof(float) &&
+ op->src[1]->nb[0] == k_elem && op->src[2]->nb[0] == v_elem && op->nb[0] == sizeof(float) &&
+ op->src[0]->ne[0] == op->src[1]->ne[0] && // dk matches
+ op->src[2]->ne[0] == op->ne[0] && // dv matches
+ op->src[2]->ne[0] <= 512 && // dv limit
+ op->src[0]->ne[0] <= 512 && // dk limit
+ nhq % nhk == 0 && // GQA ratio is integer
+ op->src[0]->ne[1] == op->ne[2] && op->src[0]->ne[2] == op->ne[1] &&
+ op->src[0]->ne[3] == op->ne[3] && op->src[1]->ne[1] == op->src[2]->ne[1] &&
+ op->src[1]->ne[2] == op->src[2]->ne[2] && op->src[1]->ne[3] == op->src[2]->ne[3] &&
+ op->src[0]->ne[3] == op->src[1]->ne[3];
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_GET_ROWS:
+ // Support F32/F16/Q4_0/Q8_0/Q4_K data with I32 indices -> F32 output
+ if (op->type == GGML_TYPE_F32 && op->src[0] &&
+ (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 ||
+ op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0 ||
+ op->src[0]->type == GGML_TYPE_Q4_K) &&
+ op->src[1] && op->src[1]->type == GGML_TYPE_I32 && ggml_is_contiguous(op) &&
+ ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) {
+ // Validate dimension constraints from ggml implementation
+ supported = (op->src[0]->ne[2] == op->src[1]->ne[1]) && (op->src[1]->ne[3] == 1);
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_CONT:
+ // Support F32->F32 and F16->F16 CONT operations (rearrange non-contiguous to contiguous)
+ if ((op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && op->src[0] &&
+ op->src[0]->type == op->type && ggml_is_contiguous(op)) {
+ // Defensive check: ensure dst and src0 are not aliased (separate buffers)
+ // While GGML design currently guarantees this, check for future robustness
+ if (op->data && op->src[0]->data && op->data == op->src[0]->data) {
+ GGML_LOG_WARN("ET: CONT operation detected aliased tensors (dst == src0), unsupported");
+ supported = false;
+ } else {
+ supported = true;
+ }
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_CPY:
+ // CPY copies src[0] data into dst layout (same as CONT for same-type)
+ // Special path: zero-element tensors (scalars) are accepted as no-ops
+ if (op->src[0]) {
+ const int64_t nelements = op->ne[0] * op->ne[1] * op->ne[2] * op->ne[3];
+ if (nelements == 0) {
+ // Zero-element / scalar no-op case - always supported
+ supported = true;
+ } else if ((op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && op->src[0]->type == op->type &&
+ ggml_is_contiguous(op)) {
+ // Same-type with contiguous dst - reuse CONT kernel
+ if (op->data && op->src[0]->data && op->data == op->src[0]->data) {
+ GGML_LOG_WARN("ET: CPY operation detected aliased tensors, unsupported");
+ supported = false;
+ } else {
+ supported = true;
+ }
+ } else if (op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op)) {
+ // F32 -> F16 conversion copy
+ supported = true;
+ } else {
+ supported = false;
+ }
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_CONCAT:
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op)) {
+ const int32_t dim = ((const int32_t *) op->op_params)[0];
+ if (dim == 0 && op->src[0]->ne[0] % 16 == 0 && op->src[1]->ne[0] % 16 == 0 &&
+ ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) {
+ // Fast dim==0 path: both source row segments are cacheline-aligned
+ // and contiguous, so the kernel can use vector row copies.
+ supported = true;
+ } else if (dim == 0 && ((op->src[0]->nb[0] % sizeof(float) == 0) || op->src[0]->ne[0] == 1) &&
+ ((op->src[1]->nb[0] % sizeof(float) == 0) || op->src[1]->ne[0] == 1)) {
+ // Slow dim==0 path: scalar, stride-aware copies for non-contiguous
+ // or non-aligned source row segments. Destination remains contiguous.
+ supported = true;
+ } else if (op->ne[0] % 16 == 0 && op->src[0]->ne[0] % 16 == 0 && op->src[1]->ne[0] % 16 == 0 &&
+ ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) {
+ // Dim >= 1 path: full aligned row copies from one source or the other.
+ supported = true;
+ }
+ }
+ break;
+ case GGML_OP_SSM_CONV:
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && op->src[0]->nb[0] == sizeof(float) &&
+ op->src[1]->nb[0] == sizeof(float) && op->src[0]->nb[1] == op->src[0]->ne[0] * sizeof(float) &&
+ op->src[1]->nb[1] == op->src[1]->ne[0] * sizeof(float) && ggml_is_contiguous(op) &&
+ op->src[1]->ne[1] == op->src[0]->ne[1] && op->ne[0] == op->src[0]->ne[1] &&
+ op->ne[1] == op->src[0]->ne[0] - op->src[1]->ne[0] + 1 && op->ne[2] == op->src[0]->ne[2];
+ break;
+ case GGML_OP_PAD:
+ // F32 zero-pad only, no dim0 padding, dst contiguous
+ // ne[0] must be CL-aligned (% 16 == 0) or evenly divide a CL (16 % ne[0] == 0)
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && (op->ne[0] % 16 == 0 || 16 % op->ne[0] == 0) &&
+ op->src[0]->nb[0] == sizeof(float)) {
+ const int32_t lp0 = ((const int32_t *) op->op_params)[0];
+ const int32_t rp0 = ((const int32_t *) op->op_params)[1];
+ const bool circular = (bool) ((const int32_t *) op->op_params)[8];
+ if (lp0 == 0 && rp0 == 0 && !circular) {
+ supported = true;
+ } else {
+ supported = false;
+ }
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_REPEAT:
+ // Two acceptable shapes:
+ // 1. No-op REPEAT (src and dst have identical shape): dispatched
+ // to cont_f32, which handles arbitrary contiguous sizes.
+ // 2. Real REPEAT via repeat_f32 kernel: dst ne[0] cacheline-aligned,
+ // src0 ne[0] cacheline-aligned or 1, dst.ne[i] % src0.ne[i] == 0.
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) && ggml_are_same_shape(op->src[0], op)) {
+ supported = true;
+ } else if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ (op->src[0]->ne[0] == 1 || op->src[0]->ne[0] % 16 == 0) && op->ne[0] % 16 == 0 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) && op->ne[0] % op->src[0]->ne[0] == 0 &&
+ op->ne[1] % op->src[0]->ne[1] == 0 && op->ne[2] % op->src[0]->ne[2] == 0 &&
+ op->ne[3] % op->src[0]->ne[3] == 0) {
+ supported = true;
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_FILL:
+ // F32 contiguous, ne[0] cacheline-aligned for SIMD fill
+ supported = op->type == GGML_TYPE_F32 && ggml_is_contiguous(op) && op->ne[0] % 16 == 0;
+ break;
+ case GGML_OP_DIAG:
+ // F32 contiguous dst, src0 is 1D vector [N,1,...], dst is [N,N,...]
+ // ne[0] must be cacheline-aligned for SIMD zeroing
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ op->ne[0] % 16 == 0 && op->ne[0] == op->ne[1] && op->src[0]->ne[0] == op->ne[0] &&
+ op->src[0]->ne[1] == 1 && ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]);
+ break;
+ case GGML_OP_TRI:
+ // F32 contiguous, same shape in/out
+ // Kernel handles arbitrary ne[0] with aligned fast path + scalar fallback
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]);
+ break;
+ case GGML_OP_SOLVE_TRI:
+ // F32 contiguous, A square, shapes compatible
+ // Only lower-triangular left-side non-unit variant
+ // Require k % 16 == 0 for cache-line-safe column parallelism
+ supported = op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && op->src[0]->ne[0] == op->src[0]->ne[1] &&
+ op->src[0]->ne[1] == op->src[1]->ne[1] && op->src[1]->ne[0] % 16 == 0 &&
+ ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
+ break;
+ case GGML_OP_SET:
+ // Minimal useful support: inplace F32 SET of a contiguous src1 view into
+ // a contiguous dst/base tensor using explicit destination view strides.
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) &&
+ ggml_is_contiguous(op->src[1]) && ggml_are_same_shape(op, op->src[0]) && op->src[1]->ne[0] % 16 == 0) {
+ const bool inplace = (bool) ((const int32_t *) op->op_params)[4];
+ const size_t nb1 = ((const int32_t *) op->op_params)[0];
+ const size_t nb2 = ((const int32_t *) op->op_params)[1];
+ const size_t nb3 = ((const int32_t *) op->op_params)[2];
+ const size_t offset = ((const int32_t *) op->op_params)[3];
+ const size_t nb0 = ggml_element_size(op);
+ const size_t im0 = op->src[1]->ne[0] == 0 ? 0 : op->src[1]->ne[0] - 1;
+ const size_t im1 = op->src[1]->ne[1] == 0 ? 0 : op->src[1]->ne[1] - 1;
+ const size_t im2 = op->src[1]->ne[2] == 0 ? 0 : op->src[1]->ne[2] - 1;
+ const size_t im3 = op->src[1]->ne[3] == 0 ? 0 : op->src[1]->ne[3] - 1;
+
+ const bool view_bounds_ok = offset + im0 * nb0 + im1 * nb1 + im2 * nb2 + im3 * nb3 <= ggml_nbytes(op);
+
+ const bool cacheline_aligned =
+ (nb1 % 64 == 0) && (nb2 % 64 == 0) && (nb3 % 64 == 0) && (offset % 64 == 0);
+
+ supported = inplace && view_bounds_ok && cacheline_aligned;
+ }
+ break;
+ case GGML_OP_RWKV_WKV6:
+ // F32 contiguous, head_size must be multiple of 8 for vectorization
+ // 6 sources: k, v, r, tf, td, state
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && op->src[2] && op->src[2]->type == GGML_TYPE_F32 && op->src[3] &&
+ op->src[3]->type == GGML_TYPE_F32 && op->src[4] && op->src[4]->type == GGML_TYPE_F32 && op->src[5] &&
+ op->src[5]->type == GGML_TYPE_F32 && op->src[0]->ne[0] % 8 == 0 && // head_size multiple of 8
+ ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) &&
+ ggml_is_contiguous(op->src[3]) && ggml_is_contiguous(op->src[4]) && ggml_is_contiguous(op->src[5])) {
+ supported = true;
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_RWKV_WKV7:
+ // F32 contiguous, head_size must be multiple of 8 for vectorization
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
+ op->src[1]->type == GGML_TYPE_F32 && op->src[2] && op->src[2]->type == GGML_TYPE_F32 && op->src[3] &&
+ op->src[3]->type == GGML_TYPE_F32 && op->src[4] && op->src[4]->type == GGML_TYPE_F32 && op->src[5] &&
+ op->src[5]->type == GGML_TYPE_F32 && op->src[6] && op->src[6]->type == GGML_TYPE_F32 &&
+ op->src[2]->ne[0] % 8 == 0 && // head_size multiple of 8
+ ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) &&
+ ggml_is_contiguous(op->src[3]) && ggml_is_contiguous(op->src[4]) && ggml_is_contiguous(op->src[5]) &&
+ ggml_is_contiguous(op->src[6])) {
+ supported = true;
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_GATED_DELTA_NET:
+ // F32, S_v must be multiple of 8 for vectorization
+ // q, k, v may be row-contiguous with strided higher dimensions.
+ // g, beta, state stay contiguous.
+ if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && // q
+ op->src[1] && op->src[1]->type == GGML_TYPE_F32 && // k
+ op->src[2] && op->src[2]->type == GGML_TYPE_F32 && // v
+ op->src[3] && op->src[3]->type == GGML_TYPE_F32 && // g
+ op->src[4] && op->src[4]->type == GGML_TYPE_F32 && // beta
+ op->src[5] && op->src[5]->type == GGML_TYPE_F32 && // state
+ op->src[2]->ne[0] % 8 == 0 && // S_v multiple of 8
+ (op->src[3]->ne[0] == 1 || op->src[3]->ne[0] == op->src[2]->ne[0]) && // g is scalar or per-element
+ op->src[4]->ne[0] == 1 && // beta is scalar per position
+ et_ggml_is_row_contiguous(op->src[0]) && et_ggml_is_row_contiguous(op->src[1]) &&
+ et_ggml_is_row_contiguous(op->src[2]) && ggml_is_contiguous(op->src[3]) &&
+ ggml_is_contiguous(op->src[4]) && ggml_is_contiguous(op->src[5])) {
+ supported = true;
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_VIEW:
+ case GGML_OP_PERMUTE:
+ case GGML_OP_TRANSPOSE:
+ case GGML_OP_RESHAPE:
+ // Metadata-only no-ops, accept any type
+ supported = true;
+ break;
+ case GGML_OP_SET_ROWS:
+ // Support F32 data with I64 indices -> F16/F32 output (scatter operation)
+ if (op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] && op->src[1]->type == GGML_TYPE_I64 &&
+ (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && ggml_is_contiguous_rows(op) &&
+ ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous(op->src[1])) {
+ // Validate dimension constraints from ggml implementation
+ supported = (op->ne[0] == op->src[0]->ne[0]) && // same number of columns
+ (op->ne[2] == op->src[0]->ne[2]) && // same batch size
+ (op->ne[3] == op->src[0]->ne[3]) && // same outer dimension
+ (op->src[0]->ne[1] == op->src[1]->ne[0]) && // src rows = index count
+ (op->src[0]->ne[2] % op->src[1]->ne[1] == 0) && // batch constraint
+ (op->src[0]->ne[3] % op->src[1]->ne[2] == 0) && // outer constraint
+ (op->src[1]->ne[3] == 1); // indices tensor constraint
+ } else {
+ supported = false;
+ }
+ break;
+ case GGML_OP_NONE:
+ // Always support NONE operations - they represent leaf nodes (parameters, inputs, constants)
+ // No computation needed, just memory management
+ supported = true;
+ break;
+ default:
+ supported = false;
+ break;
+ }
+ // if(!supported) {
+ // ggml_et_dump_operator_metadata(op);
+ // }
+ return supported;
+}
+
+static bool ggml_backend_et_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {
+ GGML_UNUSED(dev);
+ return buft->iface.get_name == ggml_backend_et_buffer_type_get_name;
+}
+
+static bool ggml_backend_et_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
+ // GET_ROWS (embedding lookup) uses a large weight (tok_embd) that lives on CPU (dev_input).
+ // The scheduler has no mechanism to cache cross-backend weight copies - it re-copies split
+ // inputs every graph_compute call. For GET_ROWS this means copying the entire embedding table
+ // (e.g. 266MB for Llama 3.1 1B) from host to device on every token, just to look up a few rows.
+ // Keep GET_ROWS on CPU and let the scheduler copy only the small result to the device.
+ // The other backends either only offload if the tensor lives on device or is large enough to
+ // justify the copy cost.
+ if (op->op == GGML_OP_GET_ROWS) {
+ return false;
+ }
+ return true;
+
+ GGML_UNUSED(dev);
+}
+
+static const struct ggml_backend_i ggml_backend_et_i = {
+ /* .get_name = */ ggml_backend_et_get_name,
+ /* .free = */ ggml_backend_et_free,
+ /* .set_tensor_async = */ ggml_backend_et_set_tensor_async,
+ /* .get_tensor_async = */ ggml_backend_et_get_tensor_async,
+ /* .set_tensor_2d_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .cpy_tensor_async = */ NULL,
+ /* .synchronize = */ ggml_backend_et_synchronize,
+ /* .graph_plan_create = */ NULL,
+ /* .graph_plan_free = */ NULL,
+ /* .graph_plan_update = */ NULL,
+ /* .graph_plan_compute = */ NULL,
+ /* .graph_compute = */ ggml_backend_et_graph_compute,
+ /* .event_record = */ NULL,
+ /* .event_wait = */ NULL,
+ /* .graph_optimize = */ NULL,
+};
+
+static const char * ggml_backend_et_device_get_name(ggml_backend_dev_t dev) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ return dev_ctx->name.c_str();
+}
+
+static const char * ggml_backend_et_device_get_description(ggml_backend_dev_t dev) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ return dev_ctx->desc.c_str();
+}
+
+static void ggml_backend_et_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ // Currently getFreeMemory is not available on a runtime without server.
+ // For now, report total memory as free.
+ *free = dev_ctx->total_mem;
+ *total = dev_ctx->total_mem;
+}
+
+static enum ggml_backend_dev_type ggml_backend_et_device_get_type(ggml_backend_dev_t dev) {
+ GGML_UNUSED(dev);
+ return GGML_BACKEND_DEVICE_TYPE_GPU;
+}
+
+static void ggml_backend_et_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) {
+ GGML_UNUSED(dev);
+ props->name = ggml_backend_et_device_get_name(dev);
+ props->description = ggml_backend_et_device_get_description(dev);
+ props->type = ggml_backend_et_device_get_type(dev);
+ ggml_backend_et_device_get_memory(dev, &props->memory_free, &props->memory_total);
+ props->device_id = NULL; // No PCI device ID available
+ props->caps = {
+ /* .async = */ true,
+ /* .host_buffer = */ false,
+ /* .buffer_from_host_ptr = */ false,
+ /* .events = */ false,
+ };
+}
+
+static ggml_backend_t ggml_backend_et_device_init_backend(ggml_backend_dev_t dev, const char * params) {
+ GGML_UNUSED(params);
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ return ggml_backend_et_init(dev_ctx->devidx);
+}
+
+static ggml_backend_buffer_type_t ggml_backend_et_device_get_buffer_type(ggml_backend_dev_t dev) {
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ return dev_ctx->buftype;
+}
+
+static ggml_backend_buffer_type_t ggml_backend_et_device_get_host_buffer_type(ggml_backend_dev_t dev) {
+ GGML_UNUSED(dev);
+ return ggml_backend_cpu_buffer_type();
+}
+
+static const struct ggml_backend_device_i ggml_backend_et_device_i = {
+ /* .get_name = */ ggml_backend_et_device_get_name,
+ /* .get_description = */ ggml_backend_et_device_get_description,
+ /* .get_memory = */ ggml_backend_et_device_get_memory,
+ /* .get_type = */ ggml_backend_et_device_get_type,
+ /* .get_props = */ ggml_backend_et_device_get_props,
+ /* .init_backend = */ ggml_backend_et_device_init_backend,
+ /* .get_buffer_type = */ ggml_backend_et_device_get_buffer_type,
+ /* .get_host_buffer_type = */ ggml_backend_et_device_get_host_buffer_type,
+ /* .buffer_from_host_ptr = */ NULL,
+ /* .supports_op = */ ggml_backend_et_device_supports_op,
+ /* .supports_buft = */ ggml_backend_et_device_supports_buft,
+ /* .offload_op = */ ggml_backend_et_device_offload_op,
+ /* .event_new = */ NULL,
+ /* .event_free = */ NULL,
+ /* .event_synchronize = */ NULL,
+};
+
+/*
+ Backend Registry.
+*/
+
+static const char * ggml_backend_et_reg_get_name(ggml_backend_reg_t reg) {
+ GGML_UNUSED(reg);
+ return GGML_ET_NAME;
+}
+
+static size_t ggml_backend_et_reg_get_device_count(ggml_backend_reg_t reg) {
+ ggml_backend_et_reg_ctx * ctx = (ggml_backend_et_reg_ctx *) reg->context;
+ return ctx->devices.size();
+}
+
+static ggml_backend_dev_t ggml_backend_et_reg_get_device(ggml_backend_reg_t reg, size_t devidx) {
+ ggml_backend_et_reg_ctx * ctx = (ggml_backend_et_reg_ctx *) reg->context;
+ if (devidx >= ctx->devices.size()) {
+ return nullptr;
+ }
+ return ctx->devices[devidx];
+}
+
+static void * ggml_backend_et_get_proc_address(ggml_backend_reg_t reg, const char * name) {
+ GGML_UNUSED(reg);
+ GGML_UNUSED(name);
+ return nullptr;
+}
+
+static const struct ggml_backend_reg_i ggml_backend_et_reg_i = {
+ /* .get_name = */ ggml_backend_et_reg_get_name,
+ /* .get_device_count = */ ggml_backend_et_reg_get_device_count,
+ /* .get_device = */ ggml_backend_et_reg_get_device,
+ /* .get_proc_address = */ ggml_backend_et_get_proc_address,
+};
+
+ggml_backend_reg_t ggml_backend_et_reg(void) {
+ static ggml_backend_reg_t _reg = []() -> ggml_backend_reg_t {
+ ggml_backend_et_reg_ctx * ctx = new ggml_backend_et_reg_ctx;
+
+ if (!ggml_et_driver_init()) {
+ return nullptr;
+ }
+
+ ggml_backend_reg_t r = new ggml_backend_reg{
+ /* .api_version = */ GGML_BACKEND_API_VERSION,
+ /* .iface = */ ggml_backend_et_reg_i,
+ /* .context = */ nullptr, // Set later
+ };
+
+ std::vector<rt::DeviceId> rtids = ggml_et_runtime()->getDevices();
+
+ for (int i = 0; i < ggml_et_devicelayer()->getDevicesCount(); i++) {
+ ggml_backend_dev_t dev = new ggml_backend_device{
+ /* .iface = */ ggml_backend_et_device_i,
+ /* .reg = */ r,
+ /* .context = */ nullptr // Set later
+ };
+
+ rt::DeviceId rtid = rtids[i];
+ rt::DeviceProperties prop = ggml_et_runtime()->getDeviceProperties(rtid);
+
+ // Create device context.
+ ggml_backend_et_device_context * dev_ctx = new ggml_backend_et_device_context;
+ dev_ctx->devidx = i;
+ dev_ctx->rtid = rtid;
+ dev_ctx->name = GGML_ET_NAME + std::to_string(i);
+ dev_ctx->desc = "ET device " + std::to_string(i);
+ dev_ctx->total_mem = static_cast<size_t>(prop.memorySize_);
+ {
+ const char * env = getenv("GGML_ET_UBERKERNEL");
+ dev_ctx->uberkernel_enabled = env && env[0] != '\0' && strcmp(env, "0") != 0;
+ }
+ // Add buffer type for device to device context.
+ ggml_backend_et_buffer_type_context * bufty_ctx = new ggml_backend_et_buffer_type_context;
+ bufty_ctx->devidx = i;
+ bufty_ctx->name = GGML_ET_NAME + std::to_string(i);
+ dev_ctx->buftype = new ggml_backend_buffer_type{ /* .iface = */ ggml_backend_et_buffer_type_i,
+ /* .device = */ dev,
+ /* .context = */ bufty_ctx };
+
+ // Create default stream for ordered execution on this device
+ dev_ctx->default_stream = ggml_et_runtime()->createStream(rtid);
+
+ dev_ctx->trace_buffer = ggml_et_runtime()->mallocDevice(rtid, ET_TRACE_BUFFER_SIZE);
+ // Pre-size each slot's host buffers and device-side scratch so the
+ // first few graph_compute calls don't pay a malloc/grow penalty.
+ for (auto & slot : dev_ctx->uberkernel.slots) {
+ slot.insts.reserve(256);
+ slot.params_blob.reserve(1 << 20);
+ slot.device_insts_capacity = 256 * sizeof(ggml_et_uberkernel_inst);
+ slot.device_params_capacity = 1 << 20;
+ slot.device_insts = ggml_et_runtime()->mallocDevice(rtid, slot.device_insts_capacity);
+ slot.device_params = ggml_et_runtime()->mallocDevice(rtid, slot.device_params_capacity);
+ if (slot.device_insts == nullptr) {
+ slot.device_insts_capacity = 0;
+ }
+ if (slot.device_params == nullptr) {
+ slot.device_params_capacity = 0;
+ }
+ }
+
+ dev->context = dev_ctx;
+
+ ctx->devices.push_back(dev);
+ }
+
+ r->context = ctx;
+ return r;
+ }();
+
+ return _reg;
+}
+
+ggml_guid_t ggml_backend_et_guid(void) {
+ static ggml_guid guid = { 0x4b, 0xe0, 0x72, 0x88, 0xc0, 0xf6, 0x29, 0xb4,
+ 0x79, 0x9f, 0x70, 0x68, 0x71, 0x0f, 0x6d, 0xc8 };
+ return &guid;
+}
+
+ggml_backend_t ggml_backend_et_init(size_t devidx) {
+ if (!ggml_et_driver_init()) {
+ return nullptr;
+ }
+
+ if (devidx >= (size_t) ggml_backend_et_get_device_count()) {
+ return nullptr;
+ }
+
+ ggml_backend_et_context * ctx = new ggml_backend_et_context;
+ ctx->devidx = (int) devidx;
+
+ ggml_backend_t backend = new ggml_backend{
+ /* .guid = */ ggml_backend_et_guid(),
+ /* .iface = */ ggml_backend_et_i,
+ /* .device = */ ggml_backend_et_reg_get_device(ggml_backend_et_reg(), devidx),
+ /* .context = */ ctx,
+ };
+
+ return backend;
+}
+
+bool ggml_backend_is_et(ggml_backend_t backend) {
+ return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_et_guid());
+}
+
+int ggml_backend_et_get_device_count(void) {
+ return ggml_backend_et_reg_get_device_count(ggml_backend_et_reg());
+}
+
+void ggml_backend_et_get_device_description(int devidx, char * description, size_t description_size) {
+ if (devidx < 0 || devidx >= ggml_backend_et_get_device_count()) {
+ snprintf(description, description_size, "ET Device %d (invalid)", devidx);
+ return;
+ }
+
+ ggml_backend_dev_t dev = ggml_backend_et_reg_get_device(ggml_backend_et_reg(), devidx);
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ snprintf(description, description_size, "%s", dev_ctx->desc.c_str());
+}
+
+void ggml_backend_et_get_device_memory(int devidx, size_t * free, size_t * total) {
+ if (devidx < 0 || devidx >= ggml_backend_et_get_device_count()) {
+ *free = 0;
+ *total = 0;
+ return;
+ }
+
+ ggml_backend_dev_t dev = ggml_backend_et_reg_get_device(ggml_backend_et_reg(), devidx);
+ ggml_backend_et_device_get_memory(dev, free, total);
+}
+
+ggml_backend_buffer_type_t ggml_backend_et_buffer_type(size_t dev_num) {
+ if (dev_num >= (size_t) ggml_backend_et_get_device_count()) {
+ return nullptr;
+ }
+
+ ggml_backend_dev_t dev = ggml_backend_et_reg_get_device(ggml_backend_et_reg(), dev_num);
+ ggml_backend_et_device_context * dev_ctx = (ggml_backend_et_device_context *) dev->context;
+ return dev_ctx->buftype;
+}
+
+ggml_backend_buffer_type_t ggml_backend_et_host_buffer_type(void) {
+ static ggml_backend_buffer_type host_buffer_type = {
+ /* .iface = */ ggml_backend_et_buffer_type_i,
+ /* .device = */ nullptr,
+ /* .context = */ nullptr,
+ };
+ return &host_buffer_type;
+}
+
+GGML_BACKEND_DL_IMPL(ggml_backend_et_reg)
#include <algorithm>
#ifdef _WIN32
+# define WIN32_LEAN_AND_MEAN
+# ifndef NOMINMAX
+# define NOMINMAX
+# endif
+# include <windows.h>
# include <sal.h>
#else
# include <semaphore.h>
#endif
#pragma clang diagnostic ignored "-Wnested-anon-types"
+#pragma clang diagnostic ignored "-Wlanguage-extension-token"
#pragma clang diagnostic ignored "-Wgnu-anonymous-struct"
+#pragma clang diagnostic ignored "-Wmicrosoft-enum-value"
#include <AEEStdErr.h>
#include <dspqueue.h>
#include "htp-ops.h"
#include "htp/matmul-ops.h"
#include "htp/flash-attn-ops.h"
+#include "htp/unary-ops.h"
#include "htp_iface.h"
#include "htp-drv.h"
case HTP_TRACE_EVT_HVX_FA_K_PREP: return "HVX_K_PREP";
case HTP_TRACE_EVT_HVX_FA_V_PREP: return "HVX_V_PREP";
case HTP_TRACE_EVT_HMX_COMP: return "HMX_COMP";
+ case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH";
+ case HTP_TRACE_EVT_INIT: return "INIT";
+ case HTP_TRACE_EVT_BUFF: return "BUFF";
default: return "UNKNOWN";
}
}
-static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node,
- const htp_prof_desc & pd) {
+static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node, const htp_prof_desc & pd) {
if (!opt_profile) return;
uint32_t op_usec = pd.usecs;
node.op_name().c_str(), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.kparams, op_usec, op_cycles, pd.cycles_start, mhz, pmu_str);
}
+static void ggml_hexagon_dump_batch_prof(const std::string & sess_name, const htp_opbatch_rsp & rsp) {
+ uint64_t batch_cycles = rsp.cycles_stop - rsp.cycles_start;
+ float batch_mhz = rsp.usecs > 0 ? (float) batch_cycles / rsp.usecs : 0.0f;
+
+ char evt_str[256] = "----";
+ if (opt_profile == 3) {
+ snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
+ rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
+ rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
+ rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
+ }
+
+ GGML_LOG_DEBUG("ggml-hex: %s profile-op OPBATCH|----|n-ops %u|%s|----|----|usec %u cycles %llu start %llu mhz %.1f\n",
+ sess_name.c_str(), rsp.n_ops, evt_str, rsp.usecs, (unsigned long long) batch_cycles, (unsigned long long) rsp.cycles_start, batch_mhz);
+}
+
+static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const htp_opbatch_rsp & rsp,
+ const htp_trace_desc * trace_events, uint32_t n_traces) {
+ if (opt_profile == 3 && trace_events) {
+ uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
+ for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
+ uint32_t count = rsp.n_traces[t];
+ valid_cnt[t] = count > n_traces ? n_traces : count;
+ }
+
+ for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
+ for (uint32_t idx = 0; idx < valid_cnt[t]; idx++) {
+ const auto & e = trace_events[t * n_traces + idx];
+ bool is_stop = (e.info & 0x8000) != 0;
+ uint16_t info = e.info & 0x7FFF;
+ GGML_LOG_DEBUG("ggml-hex: %s trace-evt %s: thread %u info %u %s %u\n",
+ sess_name.c_str(), htp_event_name(e.id), t, info, is_stop ? "stop" : "start", e.cycles);
+ }
+ }
+ }
+}
+
// **
static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) {
return type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || ggml_hexagon_is_repack_type(type);
}
-struct htp_mm_kernel_params;
struct ggml_hexagon_session;
+
static void ggml_hexagon_precompute_matmul_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
struct htp_mm_kernel_params * kparams
);
+static void ggml_hexagon_precompute_unary_params(
+ const struct ggml_hexagon_session * sess,
+ uint32_t op,
+ const struct ggml_tensor * src0,
+ const struct ggml_tensor * src1,
+ const struct ggml_tensor * dst,
+ struct htp_unary_kernel_params * kparams
+);
+
static void ggml_hexagon_precompute_fused_qkv_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
}
}
}
+
+ GGML_UNUSED(size);
}
// repack q4_0_tiled tensor into q4_0 data
}
}
}
+
+ GGML_UNUSED(size);
}
// repack q4_1 data into q4_1_tiled tensor
}
}
}
+
+ GGML_UNUSED(size);
}
// repack q4_1_tiled tensor into q4_1 data
}
}
}
+
+ GGML_UNUSED(size);
}
// repack q8_0 data into q8_0_tiled tensor
}
}
}
+
+ GGML_UNUSED(size);
}
// repack q8_0_tiled tensor into q8_0 data
}
}
}
+
+ GGML_UNUSED(size);
}
// repack mxfp4 data into mxfp4_tiled tensor
}
}
}
+
+ GGML_UNUSED(size);
}
// repack mxfp4_tiled tensor into mxfp4 data
}
}
}
+
+ GGML_UNUSED(size);
}
static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer,
static bool ggml_backend_hexagon_buffer_cpy_tensor(ggml_backend_buffer_t buffer,
const struct ggml_tensor * src,
struct ggml_tensor * dst) {
+ // we might optimize this later, for now take the slow path (ie get/set_tensor)
+ return false;
+
GGML_UNUSED(buffer);
GGML_UNUSED(src);
GGML_UNUSED(dst);
- // we might optimize this later, for now take the slow path (ie get/set_tensor)
- return false;
}
static void ggml_backend_hexagon_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
}
}
-static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer_type_t buffer_type) {
+static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
return 128; // HVX alignment
- GGML_UNUSED(buffer_type);
+ GGML_UNUSED(buft);
}
static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * t) {
return ggml_row_size(t->type, ne0) * ne1 * ne2 * ne3;
}
return ggml_nbytes(t);
+
+ GGML_UNUSED(buft);
}
-static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buffer_type) {
- auto * context = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context);
+static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
+ auto * context = static_cast<ggml_backend_hexagon_buffer_type_context *>(buft->context);
return context->sess->max_bufsize;
}
static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
return opt_hostbuf;
+
GGML_UNUSED(buft);
}
static bool ggml_backend_hexagon_repack_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
return false;
+
GGML_UNUSED(buft);
}
std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr to index
std::unordered_multimap<void*, int> d_map; // tensor data to index
+
+
unsigned int n_bufs; // num buffers in the batch
unsigned int n_tens; // num tensors ...
unsigned int n_ops; // num ops ...
n_bufs_max = HTP_OP_MAX_BUFS;
n_ops_max = batch_size;
- n_tens_max = n_ops_max + n_ops_max * HTP_OP_MAX_INPUTS;
+ n_tens_max = std::min<size_t>(n_ops_max + n_ops_max * HTP_OP_MAX_INPUTS, HTP_OP_MAX_TENSORS);
b_vmem_max = max_vmem;
return bi;
}
+
+
bool same_shape(const htp_tensor * h, const ggml_tensor * t) const {
int64_t ne0 = t->ne[0];
int64_t ne1 = t->ne[1];
int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2];
int64_t nb3 = is_repack ? nb2 * t->ne[2] : t->nb[3];
- return (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) &&
+ return (h->type == t->type) &&
+ (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) &&
(h->nb[0] == t->nb[0]) && (h->nb[1] == nb1) && (h->nb[2] == nb2) && (h->nb[3] == nb3);
}
htp_tensor &h = h_tens[ti];
h.bi = add_buffer(sbuf);
+ h.ti = ti;
h.data = t_offset;
h.type = t->type;
h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3];
}
+
+
h.flags = 0;
- if (ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+ if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
h.flags |= HTP_TENSOR_COMPUTE;
}
o.dst[i] = (i < outputs.size() && outputs[i]) ? add_tensor(outputs[i]) : 0xffff;
}
}
+
+ void finalize_ranges() {
+ }
};
struct ggml_hexagon_opqueue {
if (opt_profile && rsp.n_ops > 0) {
auto & ops = op_cache[rsp.id];
- uint64_t batch_usec = ggml_time_us() - start_usec[rsp.id];
- uint32_t htp_usec = 0;
-
GGML_ASSERT(rsp.n_ops <= ops.size());
const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr;
trace_events = (const htp_trace_desc *) (p_ptr + p_size);
}
- uint32_t trace_idx[HTP_MAX_NTHREADS + 1] = {0};
- uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
-
- if (opt_profile == 3) {
- for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
- uint32_t count = rsp.n_traces[t];
- valid_cnt[t] = count > n_traces ? n_traces : count;
- }
- }
+ ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp);
for (uint32_t i = 0; i < rsp.n_ops; i++) {
- htp_usec += pd[i].usecs;
-
ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]);
-
- if (opt_profile == 3) {
- uint32_t op_duration = pd[i].cycles_stop - pd[i].cycles_start;
-
- for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
- while (trace_idx[t] < valid_cnt[t]) {
- const auto & e = trace_events[t * n_traces + trace_idx[t]];
- uint32_t offset = e.cycles - pd[i].cycles_start;
- if (offset >= 0x80000000) {
- trace_idx[t]++;
- continue;
- }
- if (offset > op_duration) {
- break;
- }
- bool is_stop = (e.info & 0x8000) != 0;
- uint16_t info = e.info & 0x7FFF;
- GGML_LOG_DEBUG("ggml-hex: %s trace-op %s: thread %u event %s info %u %s %u\n",
- shm_buf->sess->c_name(), ops[i].op_name().c_str(), t, htp_event_name(e.id), info, is_stop ? "stop" : "start", e.cycles);
- trace_idx[t]++;
- }
- }
- }
}
- char evt_str[256] = "";
- if (opt_profile == 3) {
- snprintf(evt_str, sizeof(evt_str), " evt [%u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u]",
- rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
- rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
- rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
- }
-
- GGML_LOG_DEBUG("ggml-hex: %s profile-batch n-ops %u batch-dur-usec %lld htp-ops-usec %u%s\n",
- shm_buf->sess->c_name(), rsp.n_ops, (long long) batch_usec, htp_usec, evt_str);
+ ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces);
}
}
};
const uint32_t timeo = opt_oppoll ? 0 : DSPQUEUE_TIMEOUT;
int err = dspqueue_read(this->queue, &flags, 1, &n_dbufs, &dbuf, sizeof(rsp), &rsp_size, (uint8_t *) &rsp, timeo);
- if (err == AEE_EEXPIRED) {
+ if (err == AEE_EEXPIRED || err == AEE_EWOULDBLOCK) {
continue;
}
void ggml_hexagon_session::flush_batch() {
if (op_batch->empty()) { return; }
+ op_batch->finalize_ranges();
+
htp_opbatch_req req {};
dspqueue_buffer dbuf{};
GGML_LOG_DEBUG("ggml-hex: %s allocating new session\n", this->name.c_str());
- domain * my_domain = get_domain(this->domain_id);
+ domain * my_domain = htpdrv_get_domain(this->domain_id);
if (my_domain == NULL) {
GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP\n");
throw std::runtime_error("ggml-hex: failed to get CDSP domain (see log for details)");
}
}
- if (opt_profile) {
- htp_iface_pmu_conf pmu_conf{};
- std::copy(opt_pmu_evt.begin(), opt_pmu_evt.end(), pmu_conf.events);
-
- err = htp_iface_profiler(this->handle, opt_profile, &pmu_conf);
- if (err != 0) {
- GGML_LOG_ERROR("ggml-hex: failed to enable profiling: 0x%08x\n", (unsigned) err);
- }
- }
-
// Allocate buffers and state for op batching
this->op_queue = new ggml_hexagon_opqueue(this, opt_opbatch, opt_opqueue);
throw std::runtime_error("ggml-hex: iface start failed (see log for details)");
}
this->valid_iface = true;
+
+ if (opt_profile) {
+ htp_iface_pmu_conf pmu_conf{};
+ std::copy(opt_pmu_evt.begin(), opt_pmu_evt.end(), pmu_conf.events);
+
+ err = htp_iface_profiler(this->handle, opt_profile, &pmu_conf);
+ if (err != 0) {
+ GGML_LOG_ERROR("ggml-hex: failed to enable profiling: 0x%08x\n", (unsigned) err);
+ }
+ }
}
void ggml_hexagon_session::release() noexcept(true) {
}
return true;
+
+ GGML_UNUSED(sinks);
}
static bool ggml_hexagon_precompute_flash_attn_params(
const struct ggml_tensor * sinks = op->src[4];
if (ggml_hexagon_flash_attn_is_hmx_eligible(sess, q, k, v, sinks)) {
size_t Br = 0, Bc = 0;
- int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads);
+ int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0);
if (ret == 0) {
kparams->kernel_type = HTP_FA_KERNEL_HMX;
kparams->Br = Br;
kparams->u.hmx.g_br = hex_align_up(G * Br, 32);
kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0;
- kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0);
+ kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0);
const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256);
kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes
return false;
}
- GGML_UNUSED(sess);
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_matmul_is_hmx_eligible(
}
return true;
+
+ GGML_UNUSED(dst);
}
static bool ggml_hexagon_precompute_hmx_mm_params(
if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) {
// Try grouped path first
const bool use_dma_activation = (src1->nb[1]/sizeof(float) > (size_t)ne00_padded);
- size_t best_mblocks = SIZE_MAX;
- int best_act_threads = 0;
- size_t best_m_chunk = 0;
- size_t best_n_chunk = 0;
- size_t best_vtcm_size = 0;
-
- int act_threads = n_threads;
- while (act_threads >= 1) {
- const size_t f32_scratch_size = use_dma_activation ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * ne00_padded * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0;
- size_t group_overhead = 256 + f32_scratch_size;
- size_t group_size_per_n, group_size_per_m, group_size_per_mn;
- htp_mm_hmx_get_batched_chunk_costs(ne00_padded, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);
-
- size_t m_chunk_candidate = 0;
- size_t n_chunk_candidate = 0;
- size_t vtcm_size_candidate = 0;
-
- if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded,
- (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT,
- &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
- size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, ne00_padded, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads);
- if (exact_size <= vtcm_budget) {
- size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate;
- if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
- best_mblocks = mblocks;
- best_act_threads = act_threads;
- best_m_chunk = m_chunk_candidate;
- best_n_chunk = n_chunk_candidate;
- best_vtcm_size = exact_size;
- }
- }
- }
- if (act_threads == 1) {
- act_threads = 0;
- } else {
- act_threads /= 2;
- }
- }
-
- if (best_act_threads > 0) {
- m_chunk = best_m_chunk;
- n_chunk = best_n_chunk;
- vtcm_size = best_vtcm_size;
- act_threads_selected = best_act_threads;
+ if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, use_dma_activation, n_threads, pipeline, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
use_grouped = true;
}
}
if (!use_grouped) {
// Fallback to simple 2D path (group_size = 1)
- size_t best_mblocks = SIZE_MAX;
- int best_act_threads = 0;
- size_t best_m_chunk = 0;
- size_t best_n_chunk = 0;
- size_t best_vtcm_size = 0;
-
- // For MUL_MAT_ID the kernel runs one 2D matmul per expert, with M equal to the number of rows routed to that expert.
- // A single expert can receive up to all routed rows (dst->ne[1]*dst->ne[2] = n_expert_used*n_tokens), so size the chunk
- // search for that upper bound rather than ne12 (token positions only).
- // We recompute m_chunk per expert against the actual count in the NPU kernel.
- const int m_id_rows = (int) ((size_t) dst->ne[1] * dst->ne[2]);
- const int m_for_chunks = is_matmul_id ? hex_align_up(m_id_rows, 32) : ne11_padded;
- const int m_for_cost = is_matmul_id ? m_id_rows : ne11;
-
- int act_threads = n_threads;
- while (act_threads >= 1) {
- const size_t act_f32_size = is_matmul_id ? 0 : hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * ne00_padded * sizeof(float), HTP_MM_HMX_TILE_SIZE);
- size_t simple_2d_overhead = 256 + act_f32_size;
- size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn;
- htp_mm_hmx_get_2d_chunk_costs(wtype, ne00_padded, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);
-
- size_t m_chunk_candidate = 0;
- size_t n_chunk_candidate = 0;
- size_t vtcm_size_candidate = 0;
-
- if (htp_mm_hmx_compute_chunks(vtcm_budget, simple_2d_overhead, simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn, m_for_chunks, ne01_padded,
- (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) m_for_cost * HTP_MM_HMX_COST_A_CONVERT,
- &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
- size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, ne00_padded, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size);
- if (exact_size <= vtcm_budget) {
- size_t mblocks = ((size_t) m_for_cost + m_chunk_candidate - 1) / m_chunk_candidate;
- if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
- best_mblocks = mblocks;
- best_act_threads = act_threads;
- best_m_chunk = m_chunk_candidate;
- best_n_chunk = n_chunk_candidate;
- best_vtcm_size = exact_size;
- }
- }
- }
- if (act_threads == 1) {
- act_threads = 0;
- } else {
- act_threads /= 2;
- }
- }
-
- if (best_act_threads > 0) {
- m_chunk = best_m_chunk;
- n_chunk = best_n_chunk;
- vtcm_size = best_vtcm_size;
- act_threads_selected = best_act_threads;
- } else {
+ const int m_id_rows = (int) ((size_t) dst->ne[1] * dst->ne[2]);
+ if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
return false;
}
}
kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
kparams->vtcm_size = vtcm_size;
kparams->vtcm_src0_size = 0;
+ kparams->div_n_act_threads = init_fastdiv_values(act_threads_selected);
+ kparams->div_ne00_padded = init_fastdiv_values(ne00_padded);
kparams->vtcm_src1_size = 0;
kparams->vtcm_dst_size = 0;
kparams->kernel_type = HTP_MM_KERNEL_HMX_2D;
}
return true;
+
+ GGML_UNUSED(src0);
}
static void ggml_hexagon_precompute_hvx_mm_params(
int ne12,
int ne13,
bool is_matmul_id,
+ const size_t src2_row_size,
size_t vtcm_budget,
struct htp_mm_kernel_params * kparams
) {
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0->nb[1], 0, d, true, false, false
+ 0, src0->nb[1], 0, src2_row_size, d, true, false, false
);
if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0->nb[1], 0, 2, true, false, false
+ 0, src0->nb[1], 0, src2_row_size, 2, true, false, false
);
}
kparams->n_prefetch = best_n_prefetch;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false, false
);
if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false, false
);
}
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
kparams->n_prefetch = 16;
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
kparams->src1_row_size = src1->nb[1];
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
kparams->src1_row_size = src1->nb[1];
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+ dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
);
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
}
}
-static void ggml_hexagon_precompute_matmul_params(
+static void ggml_hexagon_precompute_matmul_params_impl(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * dst,
+ const size_t src2_row_size,
struct htp_mm_kernel_params * kparams
) {
memset(kparams, 0, sizeof(*kparams));
}
// Fallback to HVX parameter computation
- ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, vtcm_budget, kparams);
+ ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, src2_row_size, vtcm_budget, kparams);
finalize:
kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11);
kparams->div_ne11 = init_fastdiv_values(ne11);
}
+static void ggml_hexagon_precompute_matmul_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * src0,
+ const struct ggml_tensor * src1,
+ const struct ggml_tensor * dst,
+ struct htp_mm_kernel_params * kparams
+) {
+ ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams);
+}
+
+static void ggml_hexagon_precompute_fused_matmul_add_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * src0,
+ const struct ggml_tensor * src1,
+ const struct ggml_tensor * src2,
+ const struct ggml_tensor * dst,
+ struct htp_mm_kernel_params * kparams
+) {
+ ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2->nb[1], kparams);
+}
+
+static void ggml_hexagon_precompute_unary_params(
+ const struct ggml_hexagon_session * sess,
+ uint32_t op,
+ const struct ggml_tensor * src0,
+ const struct ggml_tensor * src1,
+ const struct ggml_tensor * dst,
+ struct htp_unary_kernel_params * kparams
+) {
+ memset(kparams, 0, sizeof(*kparams));
+
+ const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+ const uint32_t n_threads = (std::min)((uint32_t)sess->n_threads, src0_nrows);
+
+ kparams->n_threads = n_threads;
+
+ const size_t src0_data_row_size = src0->ne[0] * sizeof(float);
+ const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+
+ const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
+ const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
+
+ kparams->src0_row_size_aligned = src0_row_size_aligned;
+ kparams->dst_row_size_aligned = dst_row_size_aligned;
+
+ size_t src1_data_row_size = 0;
+ size_t src1_row_size_aligned = 0;
+ bool broadcast_weight = false;
+
+ if (op == HTP_OP_RMS_NORM_MUL) {
+ GGML_ASSERT(src1 != nullptr);
+ src1_data_row_size = src1->ne[0] * sizeof(float);
+ src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
+ broadcast_weight = (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
+ }
+
+ kparams->src1_row_size_aligned = src1_row_size_aligned;
+ kparams->broadcast_weight = broadcast_weight;
+
+ struct htp_unary_vtcm_layout L;
+ uint32_t col_tile = 0;
+ uint32_t vtcm_row_per_thread = 0;
+
+ htp_unary_vtcm_layout_build(&L, op, src0->ne[0], dst->ne[0],
+ op == HTP_OP_RMS_NORM_MUL ? src1->ne[0] : 0,
+ broadcast_weight, n_threads, sess->vtcm_size,
+ &col_tile, &vtcm_row_per_thread);
+
+ kparams->col_tile = col_tile;
+ kparams->vtcm_row_per_thread = vtcm_row_per_thread;
+ kparams->vtcm_size = L.total_bytes;
+
+ kparams->vtcm_src0_size_per_thread = L.src0_bytes;
+ kparams->vtcm_src1_size_per_thread = L.src1_bytes;
+ kparams->vtcm_dst_size_per_thread = L.dst_bytes;
+
+ kparams->vtcm_src0_size = L.src0_bytes * n_threads;
+ kparams->vtcm_src1_size = L.src1_bytes * n_threads;
+ kparams->vtcm_dst_size = L.dst_bytes * n_threads;
+
+ kparams->block = col_tile ? 0 : ((L.src0_bytes / 2) / src0_row_size_aligned);
+
+ const uint32_t tiles_per_row = col_tile > 0 ? (src0->ne[0] + col_tile - 1) / col_tile : 1;
+ kparams->div_ne01 = init_fastdiv_values(src0->ne[1]);
+ kparams->div_ne02 = init_fastdiv_values(src0->ne[2]);
+ kparams->div_ne012 = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
+ kparams->div_tpr = init_fastdiv_values(tiles_per_row);
+}
+
static void ggml_hexagon_precompute_fused_qkv_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0, // Wk
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0_row_size, src1_row_size, d, false, true, false
+ 0, src0_row_size, src1_row_size, 0, d, false, true, false
);
if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
// Test tiled first
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false
+ 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true, false
);
if (try_tiled && L.total_bytes <= sess->vtcm_size) {
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false
+ 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true, false
);
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0_row_size, src1_row_size, d, false, false, true
+ 0, src0_row_size, src1_row_size, 0, d, false, false, true
);
if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
// Test tiled first
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true
+ 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, false, true
);
if (try_tiled && L.total_bytes <= sess->vtcm_size) {
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
- 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true
+ 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, false, true
);
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
return false;
}
+ if (ggml_is_permuted(src0) || ggml_is_permuted(dst)) {
+ return false;
+ }
if (!ggml_are_same_shape(src0, dst)) {
return false;
}
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_add_id(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
if (dst->type != GGML_TYPE_F32) {
return false;
}
+ if (ggml_is_permuted(src0)) {
+ return false;
+ }
if (!ggml_are_same_shape(src0, dst)) {
return false;
}
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
}
-static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess,
- const struct ggml_tensor * op) {
+static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const struct ggml_tensor * src0 = op->src[0];
const struct ggml_tensor * src1 = op->src[1];
const struct ggml_tensor * dst = op;
return false;
}
- if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+ if (!ggml_is_contiguous_1(src0)) {
+ return false;
+ }
+ if (!ggml_is_contiguous(dst)) {
return false;
}
if (!ggml_are_same_shape(src0, src1)) {
return false;
}
- if (!ggml_is_contiguous(src1)) {
+ if (!ggml_is_contiguous_1(src1)) {
return false;
}
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
return false;
}
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
}
return true;
+
+ GGML_UNUSED(sess);
+}
+
+static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+ const struct ggml_tensor * src1 = op->src[1];
+ const struct ggml_tensor * dst = op;
+
+ const bool is_2D = ((const int32_t *) op->op_params)[6] == 1;
+ if (!is_2D) {
+ return false;
+ }
+
+ // For now support F32->F32 and F32->F16 only.
+ if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) {
+ return false;
+ }
+
+ if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) {
+ return false;
+ }
+
+ // For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs.
+ const int32_t p0 = ((const int32_t *) op->op_params)[2];
+ const int32_t p1 = ((const int32_t *) op->op_params)[3];
+ if (p0 != 0 || p1 != 0) {
+ return false;
+ }
+
+ GGML_UNUSED(sess);
+ return true;
}
static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
return false;
}
- GGML_UNUSED(sess);
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_cumsum(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
return false;
}
- GGML_UNUSED(sess);
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_diag(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
return false;
}
- GGML_UNUSED(sess);
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_solve_tri(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
return false;
}
- GGML_UNUSED(sess);
return true;
+
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_tri(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
case GGML_OP_RMS_NORM: return HTP_OP_RMS_NORM;
case GGML_OP_CONCAT: return HTP_OP_CONCAT;
case GGML_OP_SCALE: return HTP_OP_SCALE;
+ case GGML_OP_CLAMP: return HTP_OP_CLAMP;
case GGML_OP_SQR: return HTP_OP_SQR;
case GGML_OP_SQRT: return HTP_OP_SQRT;
case GGML_OP_SOFT_MAX: return HTP_OP_SOFTMAX;
case GGML_OP_SOLVE_TRI: return HTP_OP_SOLVE_TRI;
case GGML_OP_TRI: return HTP_OP_TRI;
case GGML_OP_PAD: return HTP_OP_PAD;
+ case GGML_OP_IM2COL: return HTP_OP_IM2COL;
case GGML_OP_UNARY:
switch (ggml_get_unary_op(t)) {
if (next_node->op == GGML_OP_MUL && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) {
htp_opnode node(n, {}, HTP_OP_RMS_NORM_MUL);
node.add_fused(next_node);
+
+ auto inputs = node.get_inputs();
+ const struct ggml_tensor * src0 = inputs[0];
+ const struct ggml_tensor * src1 = inputs.size() > 1 ? inputs[1] : nullptr;
+ ggml_hexagon_precompute_unary_params(sess,
+ node.opcode, src0, src1, node.dst(),
+ (struct htp_unary_kernel_params *)node.kernel_params
+ );
+
nodes.push_back(std::move(node));
i++; // skip the fused MUL node
return true;
if (n->op == GGML_OP_MUL_MAT && next_node) {
if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) {
if (next_node->src[0] == n || next_node->src[1] == n) {
+ const struct ggml_tensor * src2 = (next_node->src[0] == n) ? next_node->src[1] : next_node->src[0];
struct htp_mm_kernel_params kparams;
- ggml_hexagon_precompute_matmul_params(sess, n->src[0], n->src[1], next_node, &kparams);
- if ((size_t)kparams.vtcm_size <= sess->vtcm_size) {
+ ggml_hexagon_precompute_fused_matmul_add_params(sess, n->src[0], n->src[1], src2, next_node, &kparams);
+ const int src1_nrows = n->src[1]->ne[1] * n->src[1]->ne[2] * n->src[1]->ne[3];
+ const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1);
+ if (can_fuse && (size_t)kparams.vtcm_size <= sess->vtcm_size) {
htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD);
node.add_fused(next_node);
memcpy(node.kernel_params, &kparams, sizeof(kparams));
nodes.push_back(std::move(node));
i += 1;
return true;
- } else {
+ } else if (can_fuse) {
HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n",
kparams.vtcm_size, sess->vtcm_size);
}
node.node,
(struct htp_fa_kernel_params *)node.kernel_params
);
+ } else if (htp_op_is_unary(node.opcode)) {
+ auto inputs = node.get_inputs();
+ const struct ggml_tensor * src0 = inputs[0];
+ const struct ggml_tensor * src1 = inputs.size() > 1 ? inputs[1] : nullptr;
+ ggml_hexagon_precompute_unary_params(sess,
+ node.opcode, src0, src1, node.dst(),
+ (struct htp_unary_kernel_params *)node.kernel_params
+ );
}
computed_nodes.push_back(std::move(node));
}
}
}
}
+
+ GGML_UNUSED(backend);
}
static struct ggml_backend_i hexagon_backend_i = {
}
static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+ GGML_UNUSED(sess);
+
const struct ggml_tensor * src0 = op->src[0];
const struct ggml_tensor * dst = op;
}
return true;
+ GGML_UNUSED(sess);
}
static bool ggml_hexagon_supported_fill(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
return false;
}
- GGML_UNUSED(sess);
return true;
+ GGML_UNUSED(sess);
}
static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
case GGML_OP_L2_NORM:
case GGML_OP_RMS_NORM:
case GGML_OP_SCALE:
+ case GGML_OP_CLAMP:
supp = ggml_hexagon_supported_unary(sess, op);
break;
case GGML_UNARY_OP_SIGMOID:
case GGML_UNARY_OP_SOFTPLUS:
case GGML_UNARY_OP_TANH:
- supp = ggml_hexagon_supported_unary(sess, op);
- break;
case GGML_UNARY_OP_SILU:
case GGML_UNARY_OP_GELU:
case GGML_UNARY_OP_GELU_QUICK:
- supp = ggml_hexagon_supported_activations(sess, op);
+ supp = ggml_hexagon_supported_unary(sess, op);
break;
default:
break;
supp = ggml_hexagon_supported_ssm_conv(sess, op);
break;
+ case GGML_OP_IM2COL:
+ supp = ggml_hexagon_supported_im2col(sess, op);
+ break;
+
case GGML_OP_GATED_DELTA_NET:
supp = ggml_hexagon_supported_gated_delta_net(sess, op);
break;
}
return NULL;
+ GGML_UNUSED(reg);
}
template<typename T> std::vector<T> str_to_vec(const char* str) {
// Init Arch first since it affects other defaults
if (!str_arch) {
- int err = get_hex_arch_ver(CDSP_DOMAIN_ID, &opt_arch);
+ int err = htpdrv_get_arch(CDSP_DOMAIN_ID, &opt_arch);
if (err != 0) {
GGML_LOG_ERROR("ggml-hex: failed to query HTP version (err %d) defaulting to v73\n", err);
opt_arch = 73;
+ } else {
+ if (opt_arch < 73) {
+ GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is under supported range, capping at v73\n", opt_arch);
+ opt_arch = 73;
+ } else if (opt_arch > 81) {
+ GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is over supported range, capping at v81\n", opt_arch);
+ opt_arch = 81;
+ }
}
} else {
if (str_arch[0] == 'v' || str_arch[0] == 'V') {
opt_opstage = str_opstage ? strtoul(str_opstage, NULL, 0) : opt_opstage;
opt_opbatch = str_opbatch ? strtoul(str_opbatch, NULL, 0) : opt_opbatch;
opt_opqueue = str_opqueue ? strtoul(str_opqueue, NULL, 0) : opt_opqueue;
- opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 128);
+ opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 256);
opt_oppoll = str_oppoll ? strtoul(str_oppoll, NULL, 0) : opt_oppoll;
opt_opfusion = str_opfusion ? atoi(str_opfusion) : opt_opfusion;
opt_profile = str_profile ? atoi(str_profile) : 0;
-// sample drv interface
-
-#pragma clang diagnostic ignored "-Wgnu-anonymous-struct"
-#pragma clang diagnostic ignored "-Wmissing-prototypes"
-#pragma clang diagnostic ignored "-Wsign-compare"
-
#include <filesystem>
#include <set>
#include <sstream>
#include <string>
+
#ifdef _WIN32
# define WIN32_LEAN_AND_MEAN
# ifndef NOMINMAX
# include <windows.h>
# include <winevt.h>
#else
-# include <unistd.h>
+# include <dlfcn.h>
+# include <unistd.h>
#endif
+
+#pragma clang diagnostic ignored "-Wgnu-anonymous-struct"
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+#pragma clang diagnostic ignored "-Wsign-compare"
+#pragma clang diagnostic ignored "-Wlanguage-extension-token"
+#pragma clang diagnostic ignored "-Wmicrosoft-enum-value"
+#pragma clang diagnostic ignored "-Wnested-anon-types"
+
#include "ggml-impl.h"
#include "htp-drv.h"
#include "libdl.h"
uint32_t max_message_length,
uint32_t *message_length, uint8_t *message,
uint32_t timeout_us);
-
+typedef AEEResult (*dspqueue_read_noblock_pfn_t)(dspqueue_t queue, uint32_t *flags,
+ uint32_t max_buffers, uint32_t *num_buffers,
+ struct dspqueue_buffer *buffers,
+ uint32_t max_message_length,
+ uint32_t *message_length, uint8_t *message);
typedef int (*fastrpc_mmap_pfn_t)(int domain, int fd, void *addr, int offset, size_t length, enum fastrpc_map_flags flags);
typedef int (*fastrpc_munmap_pfn_t)(int domain, int fd, void *addr, size_t length);
fastrpc_mmap_pfn_t fastrpc_mmap_pfn = nullptr;
fastrpc_munmap_pfn_t fastrpc_munmap_pfn = nullptr;
-dspqueue_create_pfn_t dspqueue_create_pfn = nullptr;
-dspqueue_close_pfn_t dspqueue_close_pfn = nullptr;
-dspqueue_export_pfn_t dspqueue_export_pfn = nullptr;
-dspqueue_write_pfn_t dspqueue_write_pfn = nullptr;
-dspqueue_read_pfn_t dspqueue_read_pfn = nullptr;
+dspqueue_create_pfn_t dspqueue_create_pfn = nullptr;
+dspqueue_close_pfn_t dspqueue_close_pfn = nullptr;
+dspqueue_export_pfn_t dspqueue_export_pfn = nullptr;
+dspqueue_write_pfn_t dspqueue_write_pfn = nullptr;
+dspqueue_read_pfn_t dspqueue_read_pfn = nullptr;
+dspqueue_read_noblock_pfn_t dspqueue_read_noblock_pfn = nullptr;
remote_handle64_open_pfn_t remote_handle64_open_pfn = nullptr;
remote_handle64_invoke_pfn_t remote_handle64_invoke_pfn = nullptr;
uint32_t * message_length,
uint8_t * message,
uint32_t timeout_us) {
+#ifdef _WIN32
+ if (timeout_us == 0) {
+ return dspqueue_read_noblock_pfn(queue, flags, max_buffers, num_buffers, buffers, max_message_length,
+ message_length, message);
+ }
+#endif
return dspqueue_read_pfn(queue, flags, max_buffers, num_buffers, buffers, max_message_length, message_length,
message, timeout_us);
}
dlsym(handle.get(), dspqueue_export_pfn_t, dspqueue_export_pfn, dspqueue_export, false);
dlsym(handle.get(), dspqueue_write_pfn_t, dspqueue_write_pfn, dspqueue_write, false);
dlsym(handle.get(), dspqueue_read_pfn_t, dspqueue_read_pfn, dspqueue_read, false);
+ dlsym(handle.get(), dspqueue_read_noblock_pfn_t, dspqueue_read_noblock_pfn, dspqueue_read_noblock, false);
dlsym(handle.get(), remote_handle64_open_pfn_t, remote_handle64_open_pfn, remote_handle64_open, false);
dlsym(handle.get(), remote_handle64_invoke_pfn_t, remote_handle64_invoke_pfn, remote_handle64_invoke, false);
dlsym(handle.get(), remote_handle_control_pfn_t, remote_handle_control_pfn, remote_handle_control, false);
return AEE_SUCCESS;
}
-domain * get_domain(int domain_id) {
+domain * htpdrv_get_domain(int domain_id) {
int i = 0;
int size = sizeof(supported_domains) / sizeof(domain);
return NULL;
}
-int get_hex_arch_ver(int domain, int * arch) {
+int htpdrv_get_arch(int domain, int * arch) {
if (!remote_handle_control_pfn) {
GGML_LOG_ERROR("ggml-hex: remote_handle_control is not supported on this device\n");
return AEE_EUNSUPPORTEDAPI;
return err;
}
- switch (arch_ver.capability & 0xff) {
- case 0x68:
- *arch = 68;
- return 0;
- case 0x69:
- *arch = 69;
- return 0;
- case 0x73:
- *arch = 73;
- return 0;
- case 0x75:
- *arch = 75;
- return 0;
- case 0x79:
- *arch = 79;
- return 0;
- case 0x81:
- *arch = 81;
- return 0;
- }
- return -1;
+ uint32_t val = arch_ver.capability & 0xff;
+ *arch = (int) ((val >> 4) * 10 + (val & 0x0f));
+ return 0;
}
HTPDRV_API int htpdrv_init(void);
/**
- * get_domain API: get domain struct from domain value.
+ * htpdrv_get_domain API: get domain struct from domain value.
*
* @param[in] domain value of a domain
* @return Returns domain struct of the domain if it is supported or else
* returns NULL.
*
*/
-HTPDRV_API domain * get_domain(int domain_id);
+HTPDRV_API domain * htpdrv_get_domain(int domain_id);
/**
- * get_hex_arch_ver API: query the Hexagon processor architecture version information
+ * htpdrv_get_arch API: query the Hexagon processor architecture version information
*
* @param[in] domain_id value of a domain
* @param[out] Arch version (73, 75, ...)
* non-zero if error, return value points to the error.
*
*/
-HTPDRV_API int get_hex_arch_ver(int domain, int * arch);
+HTPDRV_API int htpdrv_get_arch(int domain, int * arch);
#ifdef __cplusplus
}
#include "htp-ops.h"
#include "htp/matmul-ops.h"
#include "htp/flash-attn-ops.h"
+#include "htp/unary-ops.h"
struct htp_opnode {
ggml_tensor * node = nullptr;
path = "hvx";
}
snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+ } else if (htp_op_is_unary(node.opcode)) {
+ const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
+ snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
} else {
snprintf(str, max_size, "----");
}
add_library(${HTP_LIB} SHARED
main.c
htp_iface_skel.c
- worker-pool.c
- hex-dma.c
+ work-queue.c
+ dma-queue.c
hmx-queue.c
+ htp-tensor.c
+ matmul-ops.c
+ flash-attn-ops.c
gated-delta-net-ops.c
binary-ops.c
unary-ops.c
get-rows-ops.c
cpy-ops.c
repeat-ops.c
- argsort-ops.c
ssm-conv.c
cumsum-ops.c
fill-ops.c
diag-ops.c
solve-tri-ops.c
pad-ops.c
- flash-attn-ops.c
- matmul-ops.c
+ argsort-ops.c
+ im2col-ops.c
)
target_compile_definitions(${HTP_LIB} PRIVATE
#include "htp-ctx.h"
#include "htp-ops.h"
#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
#define htp_act_preamble \
const struct htp_tensor * src0 = actx->octx->src[0]; \
const uint32_t nb3 = dst->nb[3];
struct htp_act_context {
- struct htp_ops_context * octx;
+ struct htp_ops_context * octx;
// Precomputed values
- const uint8_t * data_src0;
- const uint8_t * data_src1;
- uint8_t * data_dst;
-
- size_t src0_row_size;
- size_t src1_row_size;
- size_t dst_row_size;
-
- size_t src0_row_size_aligned;
- size_t src1_row_size_aligned;
- size_t dst_row_size_aligned;
-
- size_t src0_spad_half_size;
- size_t src1_spad_half_size;
- size_t dst_spad_half_size;
-
- uint32_t block;
- uint32_t src0_nrows;
- uint32_t src0_nrows_per_thread;
- int nc;
-};
-
-static void glu_swiglu_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
- struct htp_act_context * actx = (struct htp_act_context *) data;
- htp_act_preamble;
-
- size_t src0_row_size = actx->src0_row_size;
- size_t src1_row_size = actx->src1_row_size;
- size_t dst_row_size = actx->dst_row_size;
-
- const uint32_t src0_nrows = actx->src0_nrows;
- const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;
- const uint32_t src0_start_row = src0_nrows_per_thread * ith;
- const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
-
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
- }
+ const uint8_t * data_src0;
+ const uint8_t * data_src1;
+ uint8_t * data_dst;
- uint64_t t1, t2;
- t1 = HAP_perf_get_qtimer_count();
+ size_t src0_row_size;
+ size_t src1_row_size;
+ size_t dst_row_size;
- const uint8_t * restrict data_src0 = actx->data_src0;
- const uint8_t * restrict data_src1 = actx->data_src1;
- uint8_t * restrict data_dst = actx->data_dst;
+ size_t src0_row_stride;
+ size_t src1_row_stride;
- const int nc = actx->nc;
+ size_t src0_row_size_aligned;
+ size_t src1_row_size_aligned;
+ size_t dst_row_size_aligned;
- const size_t src0_row_size_aligned = actx->src0_row_size_aligned;
- const size_t src1_row_size_aligned = actx->src1_row_size_aligned;
- const size_t dst_row_size_aligned = actx->dst_row_size_aligned;
+ size_t src0_spad_half_size;
+ size_t src1_spad_half_size;
+ size_t dst_spad_half_size;
- uint8_t * restrict src0_spad_data = actx->octx->src0_spad.data + (ith * actx->octx->src0_spad.size_per_thread);
- uint8_t * restrict src1_spad_data = actx->octx->src1_spad.data + (ith * actx->octx->src1_spad.size_per_thread);
- uint8_t * restrict dst_spad_data = actx->octx->dst_spad.data + (ith * actx->octx->dst_spad.size_per_thread);
+ uint32_t block;
+ uint32_t src0_nrows;
+ uint32_t src0_nrows_per_thread;
+ int nc;
- size_t src0_spad_half_size = actx->src0_spad_half_size;
- size_t src1_spad_half_size = actx->src1_spad_half_size;
- size_t dst_spad_half_size = actx->dst_spad_half_size;
+ uint8_t * vtcm_src0;
+ uint8_t * vtcm_src1;
+ uint8_t * vtcm_dst;
- const int BLOCK = actx->block;
- if (BLOCK == 0) {
- FARF(ERROR,
- "swiglu-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",
- actx->octx->src0_spad.size_per_thread, src0_row_size_aligned);
- return;
- }
+ size_t vtcm_src0_size_per_thread;
+ size_t vtcm_src1_size_per_thread;
+ size_t vtcm_dst_size_per_thread;
+};
- dma_queue * dma_queue = actx->octx->ctx->dma[ith];
+struct htp_act_vtcm_layout {
+ size_t total_bytes;
+ size_t off_src0;
+ size_t off_src1;
+ size_t off_dst;
- // See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379
- for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
+ size_t src0_bytes_per_thread;
+ size_t src1_bytes_per_thread;
+ size_t dst_bytes_per_thread;
- // Dummy DMA transation for sequencing (interleaving dst,src,dst,...)
- dma_queue_push_vtcm_to_ddr(dma_queue,
- dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
- dst_row_size, dst_row_size_aligned, 0);
+ uint32_t vtcm_row_per_thread;
+};
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_size)),
- src0_row_size_aligned, src0_row_size, block_size);
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_size)),
- src1_row_size_aligned, src1_row_size, block_size);
- }
+static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L,
+ size_t src0_row_size_aligned,
+ size_t src1_row_size_aligned,
+ size_t dst_row_size_aligned,
+ uint32_t n_threads,
+ size_t vtcm_size) {
+ const size_t spad_size_per_row = src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned;
+ const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * spad_size_per_row));
- for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- float * dst_spad = (float *) dma_queue_pop(dma_queue).src;
- float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;
- float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;
-
- for (uint32_t ib = 0; ib < block_size; ib++) {
- const float * src0_spad_ptr = src0_spad + ib * (src0_row_size_aligned / sizeof(float));
- const float * src1_spad_ptr = src1_spad + ib * (src1_row_size_aligned / sizeof(float));
- float * dst_spad_ptr = dst_spad + ib * (dst_row_size_aligned / sizeof(float));
-
- //swiglu(x) = x1 * sigmoid(x0)
- hvx_sigmoid_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, nc);
- hvx_mul_mul_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, (const uint8_t *) dst_spad_ptr,
- (const uint8_t *) src1_spad_ptr, nc);
- }
-
- dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad), dst_row_size,
- dst_row_size_aligned, block_size);
-
- // prefetch N+2 loop iteration if any
- const uint32_t pref_block = (ir + BLOCK * 2);
- if (pref_block < src0_end_row) {
- const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);
- dma_queue_push_ddr_to_vtcm(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_size)),
- src0_row_size_aligned, src0_row_size, pref_block_size);
- dma_queue_push_ddr_to_vtcm(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_size)),
- src1_row_size_aligned, src1_row_size, pref_block_size);
- }
- }
+ L->vtcm_row_per_thread = vtcm_row_per_thread;
- dma_queue_flush(dma_queue);
+ L->src0_bytes_per_thread = src0_row_size_aligned * vtcm_row_per_thread;
+ L->src1_bytes_per_thread = src1_row_size_aligned * vtcm_row_per_thread;
+ L->dst_bytes_per_thread = dst_row_size_aligned * vtcm_row_per_thread;
- t2 = HAP_perf_get_qtimer_count();
+ L->off_src0 = 0;
+ L->off_src1 = L->off_src0 + L->src0_bytes_per_thread * n_threads;
+ L->off_dst = L->off_src1 + L->src1_bytes_per_thread * n_threads;
- FARF(HIGH, "swiglu-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
- ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3,
- (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+ L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
}
-static void glu_swiglu_oai_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
- struct htp_act_context * actx = (struct htp_act_context *) data;
- htp_act_preamble;
-
- uint64_t t1, t2;
- t1 = HAP_perf_get_qtimer_count();
-
- size_t src0_row_size = actx->src0_row_size;
- size_t src1_row_size = actx->src1_row_size;
- size_t dst_row_size = actx->dst_row_size;
-
- const uint32_t src0_nrows = actx->src0_nrows;
- const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;
-
- const uint32_t src0_start_row = src0_nrows_per_thread * ith;
- const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
-
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
+#define htp_glu_op_preamble \
+ const size_t src0_row_size_aligned = actx->src0_row_size_aligned; \
+ const size_t src1_row_size_aligned = actx->src1_row_size_aligned; \
+ const size_t dst_row_size_aligned = actx->dst_row_size_aligned; \
+ const int nc = actx->nc;
+
+// swiglu(x) = x1 * sigmoid(x0)
+static void swiglu_f32(const float * restrict src0,
+ const float * restrict src1,
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_act_context * actx) {
+ htp_glu_op_preamble;
+
+ for (uint32_t ib = 0; ib < num_rows; ib++) {
+ const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+ const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+ uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+ hvx_sigmoid_f32_aa(dst_ptr, src0_ptr, nc);
+ hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
}
+}
- const uint8_t * restrict data_src0 = actx->data_src0;
- const uint8_t * restrict data_src1 = actx->data_src1;
- uint8_t * restrict data_dst = actx->data_dst;
-
- const int nc = actx->nc;
-
- const size_t src0_row_size_aligned = actx->src0_row_size_aligned;
- const size_t src1_row_size_aligned = actx->src1_row_size_aligned;
- const size_t dst_row_size_aligned = actx->dst_row_size_aligned;
-
- uint8_t * restrict src0_spad_data = actx->octx->src0_spad.data + (ith * actx->octx->src0_spad.size_per_thread);
- uint8_t * restrict src1_spad_data = actx->octx->src1_spad.data + (ith * actx->octx->src1_spad.size_per_thread);
- uint8_t * restrict dst_spad_data = actx->octx->dst_spad.data + (ith * actx->octx->dst_spad.size_per_thread);
-
- size_t src0_spad_half_size = actx->src0_spad_half_size;
- size_t src1_spad_half_size = actx->src1_spad_half_size;
- size_t dst_spad_half_size = actx->dst_spad_half_size;
-
- const int BLOCK = actx->block;
- if (BLOCK == 0) {
- FARF(ERROR,
- "swiglu-oai-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least "
- "%zu\n",
- actx->octx->src0_spad.size_per_thread, src0_row_size_aligned);
- return;
- }
+// out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f)
+static void swiglu_oai_f32(const float * restrict src0,
+ const float * restrict src1,
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_act_context * actx) {
+ htp_glu_op_preamble;
const float alpha = ((const float *) (actx->octx->op_params))[2];
const float limit = ((const float *) (actx->octx->op_params))[3];
- dma_queue * dma_queue = actx->octx->ctx->dma[ith];
-
- // See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379
- for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- // Dummy DMA transation for sequencing (interleaving dst,src,dst,...)
- dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
- dst_row_size, dst_row_size_aligned, 0);
-
- dma_queue_push_ddr_to_vtcm(
- dma_queue,
- dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_size)),
- src0_row_size_aligned, src0_row_size, block_size);
- dma_queue_push_ddr_to_vtcm(
- dma_queue,
- dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_size)),
- src1_row_size_aligned, src1_row_size, block_size);
+ for (uint32_t ib = 0; ib < num_rows; ib++) {
+ const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+ const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+ uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+ // x (src0_ptr) = std::min(src0_p[k], limit);
+ hvx_min_scalar_f32((uint8_t *) src0_ptr, src0_ptr, limit, nc);
+ // y1 (src1_ptr) = std::clamp(src1_p[k], -limit, limit);
+ hvx_clamp_scalar_f32((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
+ // y (src1_ptr) = y1 + 1.f
+ hvx_add_scalar_f32((uint8_t *) src1_ptr, src1_ptr, 1.0, nc);
+ // x1 (dst_ptr) = alpha * x
+ hvx_mul_scalar_f32(dst_ptr, src0_ptr, alpha, nc);
+ // x2 (dst_ptr) = sigmoid(x1) = 1/(1+exp(-x1))
+ hvx_sigmoid_f32_aa(dst_ptr, dst_ptr, nc);
+ // out = x * sigmoid(alpha * x) * (y + 1.f)
+ hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
}
-
- for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- float * dst_spad = (float *) dma_queue_pop(dma_queue).src;
- float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;
- float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;
-
- for (uint32_t ib = 0; ib < block_size; ib++) {
- const float * src0_spad_ptr = src0_spad + ib * (src0_row_size_aligned / sizeof(float));
- const float * src1_spad_ptr = src1_spad + ib * (src1_row_size_aligned / sizeof(float));
- float * dst_spad_ptr = dst_spad + ib * (dst_row_size_aligned / sizeof(float));
-
- // x (src0_spad_data) = std::min(src0_p[k], limit);
- hvx_min_scalar_f32((uint8_t *) src0_spad_ptr, (const uint8_t *) src0_spad_ptr, limit, nc);
- // y1 (src1_spad_data) = std::clamp(src1_p[k], -limit, limit);
- hvx_clamp_scalar_f32((uint8_t *) src1_spad_ptr, (const uint8_t *) src1_spad_ptr, -limit, limit, nc);
- // y (src1_spad_data) = y1 + 1.f
- hvx_add_scalar_f32((uint8_t *) src1_spad_ptr, (const uint8_t *) src1_spad_ptr, 1.0, nc);
- // x1 (dst_spad_data) = alpha * (x)
- hvx_mul_scalar_f32((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, alpha, nc);
- // x2 (dst_spad_data) = sigmoid(x1) = 1/(1+exp(-x1))
- hvx_sigmoid_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) dst_spad_ptr, nc);
- // out = x * sigmoid(alpha * x) * (y + 1.f)
- hvx_mul_mul_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, (const uint8_t *) dst_spad_ptr,
- (const uint8_t *) src1_spad_ptr, nc);
- }
-
- dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad), dst_row_size,
- dst_row_size_aligned, block_size);
-
- // prefetch N+2 loop iteration if any
- const uint32_t pref_block = (ir + BLOCK * 2);
- if (pref_block < src0_end_row) {
- const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);
- dma_queue_push_ddr_to_vtcm(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_size)),
- src0_row_size_aligned, src0_row_size, pref_block_size);
- dma_queue_push_ddr_to_vtcm(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_size)),
- src1_row_size_aligned, src1_row_size, pref_block_size);
- }
- }
-
- dma_queue_flush(dma_queue);
-
- t2 = HAP_perf_get_qtimer_count();
-
- FARF(HIGH, "swiglu-oai-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, src0->ne[0],
- src0->ne[1], src0->ne[2], src0->ne[3], src0_start_row, src0_end_row, src1->ne[0], src1->ne[1], src1->ne[2],
- src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
}
+static const float GELU_COEF_A = 0.044715f;
+static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
-static void unary_gelu_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
- struct htp_act_context * actx = (struct htp_act_context *) data;
- htp_act_preamble;
-
- uint64_t t1, t2;
- t1 = HAP_perf_get_qtimer_count();
-
- const size_t src0_row_size = actx->src0_row_size;
- const size_t dst_row_size = actx->dst_row_size;
- const size_t src0_row_size_aligned = actx->src0_row_size_aligned;
- const size_t dst_row_size_aligned = actx->dst_row_size_aligned;
-
- const uint32_t src0_nrows = actx->src0_nrows;
- const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;
-
- const uint32_t src0_start_row = src0_nrows_per_thread * ith;
- const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
-
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
- }
-
- const uint8_t * data_src0 = actx->data_src0;
- uint8_t * data_dst = actx->data_dst;
-
- // nc/ne0 matches.
- const int ne0_val = actx->nc; // == dst->ne[0]
-
- uint8_t * src0_spad_data = actx->octx->src0_spad.data + (ith * actx->octx->src0_spad.size_per_thread);
- uint8_t * dst_spad_data = actx->octx->dst_spad.data + (ith * actx->octx->dst_spad.size_per_thread);
-
- size_t src0_spad_half_size = actx->src0_spad_half_size;
- size_t dst_spad_half_size = actx->dst_spad_half_size;
-
- // In gelu = x*sigmoid(x*1.702)
- const int BLOCK = actx->block;
-
- if (BLOCK == 0) {
- FARF(ERROR, "gelu-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",
- actx->octx->src0_spad.size_per_thread, src0_row_size_aligned);
- return;
- }
-
- dma_queue * dma_queue = actx->octx->ctx->dma[ith];
-
- // See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379
- for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- // Dummy DMA transation for sequencing (interleaving dst,src,dst,...)
- dma_queue_push_vtcm_to_ddr(dma_queue,
- dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
- dst_row_size, dst_row_size_aligned, 0);
-
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_size)),
- src0_row_size_aligned, src0_row_size, block_size);
- }
+static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+ assert((unsigned long) dst % 128 == 0);
+ assert((unsigned long) src0 % 128 == 0);
+ assert((unsigned long) src1 % 128 == 0);
- for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- float* dst_spad = (float *) dma_queue_pop(dma_queue).src;
- float* src0_spad = (float *) dma_queue_pop(dma_queue).dst;
-
- for (uint32_t ib = 0; ib < block_size; ib++) {
- const float* src0_spad_ptr = src0_spad + ib * (src0_row_size_aligned / sizeof(float));
- float* dst_spad_ptr = dst_spad + ib * (dst_row_size_aligned / sizeof(float));
-
- // gelu = x * sigmoid(1.702 * x) // current implementation
- hvx_mul_scalar_f32((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, (float) 1.702, ne0_val);
- hvx_sigmoid_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) dst_spad_ptr, ne0_val);
- hvx_mul_f32_aaa((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, (const uint8_t *) dst_spad_ptr, ne0_val);
- }
-
- dma_queue_push_vtcm_to_ddr(dma_queue,
- dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),
- dst_row_size, dst_row_size_aligned, block_size);
-
- // prefetch N+2 loop iteration if any
- const uint32_t pref_block = (ir + BLOCK * 2);
- if (pref_block < src0_end_row) {
- const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_size)),
- src0_row_size_aligned, src0_row_size, pref_block_size);
- }
- }
+ HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+ const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+ const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
- dma_queue_flush(dma_queue);
+ const uint32_t epv = 128 / sizeof(float);
+ const uint32_t nvec = n / epv;
+ const uint32_t nloe = n % epv;
- t2 = HAP_perf_get_qtimer_count();
+ const float GELU_COEF_A_TIMES_SQRT = GELU_COEF_A * SQRT_2_OVER_PI;
- FARF(HIGH, "gelu-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", ith, nth, ne00, ne01, ne02,
- ne03, src0_start_row, src0_end_row, ne0, ne1, ne2, ne3, (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
-}
+ const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT);
+ const HVX_Vector v_sqrt_2_pi = hvx_vec_splat_f32(SQRT_2_OVER_PI);
+ const HVX_Vector v_half = hvx_vec_splat_f32(0.5f);
+ const HVX_Vector v_one = hvx_vec_splat_f32(1.0f);
+ const HVX_Vector v_two = hvx_vec_splat_f32(2.0f);
+ // Hoisted fast sigmoid / inverse constants to avoid loop-internal overhead
+ const HVX_Vector v_log2f = Q6_V_vsplat_R(FAST_SIGMOID_LOG2F);
+ const HVX_Vector v_c1 = Q6_V_vsplat_R(FAST_SIGMOID_C1);
+ const HVX_Vector v_c2 = Q6_V_vsplat_R(FAST_SIGMOID_C2);
+ const HVX_Vector v_inv_aprox = Q6_V_vsplat_R(0x7EEEEBB3);
+ const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
+ const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
-static void unary_silu_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
- struct htp_act_context * actx = (struct htp_act_context *) data;
- htp_act_preamble;
+ uint32_t i = 0;
- uint64_t t1, t2;
- t1 = HAP_perf_get_qtimer_count();
+ for (; i < nvec; i++) {
+ HVX_Vector x = vsrc0[i];
+ HVX_Vector g = vsrc1[i];
- const size_t src0_row_size = actx->src0_row_size;
- const size_t dst_row_size = actx->dst_row_size;
- const size_t src0_row_size_aligned = actx->src0_row_size_aligned;
- const size_t dst_row_size_aligned = actx->dst_row_size_aligned;
+ HVX_Vector x2 = hvx_vec_mul_f32_f32(x, x);
+ HVX_Vector coef = hvx_vec_mul_f32_f32(x2, v_coef_a_times_sqrt);
+ coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
+ HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
- const uint32_t src0_nrows = actx->src0_nrows;
- const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;
+ // y2 = 2 * inner
+ HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two);
- const uint32_t src0_start_row = src0_nrows_per_thread * ith;
- const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+ // Sigmoid guard check predicates
+ HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2);
+ HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp);
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
- }
+ // Fast sigmoid approximation
+ HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f);
+ v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half);
- const uint8_t * data_src0 = actx->data_src0;
- uint8_t * data_dst = actx->data_dst;
+ HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
+ HVX_Vector x_sig = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+ HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig);
- const int ne0_val = actx->nc; // == dst->ne[0]
+ HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2);
+ v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f);
- uint8_t * src0_spad_data = actx->octx->src0_spad.data + (ith * actx->octx->src0_spad.size_per_thread);
- uint8_t * dst_spad_data = actx->octx->dst_spad.data + (ith * actx->octx->dst_spad.size_per_thread);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1);
+ v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig);
+ v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig);
- size_t src0_spad_half_size = actx->src0_spad_half_size;
- size_t dst_spad_half_size = actx->dst_spad_half_size;
+ HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+ v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
- const int BLOCK = actx->block;
+ HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+ HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
- if (BLOCK == 0) {
- FARF(ERROR, "silu-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",
- actx->octx->src0_spad.size_per_thread, src0_row_size_aligned);
- return;
- }
+ // Fast division (Newton-Raphson with 2 iterations)
+ HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5);
+ HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
+ i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
+ r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
+ r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
+ HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf);
- dma_queue * dma_queue = actx->octx->ctx->dma[ith];
+ HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv));
- // See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379
- for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
+ // Sigmoid guards
+ sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one);
+ sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero());
- // Dummy DMA transation for sequencing (interleaving dst,src,dst,...)
- dma_queue_push_vtcm_to_ddr(dma_queue,
- dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
- dst_row_size, dst_row_size_aligned, 0);
+ // tanh(inner) = 2 * sigmoid(2 * inner) - 1
+ HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two);
+ tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one);
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_size)),
- src0_row_size_aligned, src0_row_size, block_size);
- }
+ HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one);
+ HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half);
+ HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one);
- for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- float* dst_spad = (float *) dma_queue_pop(dma_queue).src;
- float* src0_spad = (float *) dma_queue_pop(dma_queue).dst;
-
- for (uint32_t ib = 0; ib < block_size; ib++) {
- const float* src0_spad_ptr = src0_spad + ib * (src0_row_size_aligned / sizeof(float));
- float* dst_spad_ptr = dst_spad + ib * (dst_row_size_aligned / sizeof(float));
-
- // silu = x * sigmoid(x)
- hvx_sigmoid_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, ne0_val);
- hvx_mul_f32_aaa((uint8_t *) dst_spad_ptr, (const uint8_t *) src0_spad_ptr, (const uint8_t *) dst_spad_ptr, ne0_val);
- }
-
- dma_queue_push_vtcm_to_ddr(dma_queue,
- dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),
- dst_row_size, dst_row_size_aligned, block_size);
-
- // prefetch N+2 loop iteration if any
- const uint32_t pref_block = (ir + BLOCK * 2);
- if (pref_block < src0_end_row) {
- const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_size)),
- src0_row_size_aligned, src0_row_size, pref_block_size);
- }
+ vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g);
}
- dma_queue_flush(dma_queue);
+ if (nloe) {
+ HVX_Vector x = vsrc0[i];
+ HVX_Vector g = vsrc1[i];
- t2 = HAP_perf_get_qtimer_count();
+ HVX_Vector x2 = hvx_vec_mul_f32_f32(x, x);
+ HVX_Vector coef = hvx_vec_mul_f32_f32(x2, v_coef_a_times_sqrt);
+ coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
+ HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
- FARF(HIGH, "silu-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", ith, nth, ne00, ne01, ne02,
- ne03, src0_start_row, src0_end_row, ne0, ne1, ne2, ne3, (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
-}
+ HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two);
-static const float GELU_COEF_A = 0.044715f;
-static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
+ HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2);
+ HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp);
-static void glu_geglu_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
- struct htp_act_context * actx = (struct htp_act_context *) data;
- htp_act_preamble;
+ HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f);
+ v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half);
- size_t src0_row_size = actx->src0_row_size;
- size_t src1_row_size = actx->src1_row_size;
- size_t dst_row_size = actx->dst_row_size;
+ HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
+ HVX_Vector x_sig = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+ HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig);
- uint64_t t1, t2;
- t1 = HAP_perf_get_qtimer_count();
+ HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2);
+ v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f);
- const uint32_t src0_nrows = actx->src0_nrows;
- const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1);
+ v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig);
+ v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig);
- const uint32_t src0_start_row = src0_nrows_per_thread * ith;
- const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+ HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+ v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
- }
+ HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+ HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
- const uint8_t * restrict data_src0 = actx->data_src0;
- const uint8_t * restrict data_src1 = actx->data_src1;
- uint8_t * restrict data_dst = actx->data_dst;
+ HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5);
+ HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
+ i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
+ r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
+ r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
+ HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf);
- const int nc = actx->nc;
+ HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv));
- const size_t src0_row_size_aligned = actx->src0_row_size_aligned;
- const size_t src1_row_size_aligned = actx->src1_row_size_aligned;
- const size_t dst_row_size_aligned = actx->dst_row_size_aligned;
+ sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one);
+ sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero());
- uint8_t * restrict src0_spad_data = actx->octx->src0_spad.data + (ith * actx->octx->src0_spad.size_per_thread);
- uint8_t * restrict src1_spad_data = actx->octx->src1_spad.data + (ith * actx->octx->src1_spad.size_per_thread);
- uint8_t * restrict dst_spad_data = actx->octx->dst_spad.data + (ith * actx->octx->dst_spad.size_per_thread);
+ HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two);
+ tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one);
- size_t src0_spad_half_size = actx->src0_spad_half_size;
- size_t src1_spad_half_size = actx->src1_spad_half_size;
- size_t dst_spad_half_size = actx->dst_spad_half_size;
+ HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one);
+ HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half);
+ HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one);
- const int BLOCK = actx->block;
- if (BLOCK == 0) {
- FARF(ERROR,
- "geglu-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",
- actx->octx->src0_spad.size_per_thread, src0_row_size_aligned);
- return;
+ HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g);
+ hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
}
+}
- dma_queue * dma_queue = actx->octx->ctx->dma[ith];
-
- // See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379
- for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
+// geglu(x, g) = gelu(x) * g
+static void geglu_f32(const float * restrict src0,
+ const float * restrict src1,
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_act_context * actx) {
+ htp_glu_op_preamble;
- // Dummy DMA transation for sequencing (interleaving dst,src,dst,...)
- dma_queue_push_vtcm_to_ddr(dma_queue,
- dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
- dst_row_size, dst_row_size_aligned, 0);
+ for (uint32_t ib = 0; ib < num_rows; ib++) {
+ const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+ const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+ uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_size)),
- src0_row_size_aligned, src0_row_size, block_size);
- dma_queue_push_ddr_to_vtcm(dma_queue,
- dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_size)),
- src1_row_size_aligned, src1_row_size, block_size);
+ hvx_geglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
}
+}
- for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {
- const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
-
- float * dst_spad = (float *) dma_queue_pop(dma_queue).src;
- float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;
- float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;
-
- for (uint32_t ib = 0; ib < block_size; ib++) {
- const uint8_t * src0_spad_ptr = (const uint8_t *)(src0_spad + ib * (src0_row_size_aligned / sizeof(float)));
- const uint8_t * src1_spad_ptr = (const uint8_t *)(src1_spad + ib * (src1_row_size_aligned / sizeof(float)));
- uint8_t * dst_spad_ptr = (uint8_t *)(dst_spad + ib * (dst_row_size_aligned / sizeof(float)));
-
- // geglu tanh implementation
- // geglu(x, g) = gelu(x) * g
- // gelu(x) = 0.5f*x*(1.0f + tanhf(SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x)))
- hvx_mul_f32_aaa(dst_spad_ptr, src0_spad_ptr, src0_spad_ptr, nc); // res = x*x
- hvx_mul_scalar_f32_aa(dst_spad_ptr, (const uint8_t *)dst_spad_ptr, GELU_COEF_A, nc); // res = res * GELU_COEF_A
- hvx_add_scalar_f32_aa(dst_spad_ptr, (const uint8_t *)dst_spad_ptr, 1.0f, nc); // res = res + 1.0f
- hvx_mul_f32_aaa(dst_spad_ptr, src0_spad_ptr, (const uint8_t *)dst_spad_ptr, nc); // res = res * x
- hvx_mul_scalar_f32_aa(dst_spad_ptr, (const uint8_t*)dst_spad_ptr, SQRT_2_OVER_PI, nc); // res = result * SQRT_2_OVER_PI
- hvx_tanh_f32_aa((uint8_t *) dst_spad_ptr, (const uint8_t *) dst_spad_ptr, nc); // res = tanh(res)
- hvx_add_scalar_f32_aa(dst_spad_ptr, (const uint8_t*)dst_spad_ptr, 1.0f, nc); // res = res + 1.0f
- hvx_mul_f32_aaa(dst_spad_ptr, src0_spad_ptr, (const uint8_t *)dst_spad_ptr, nc); // res = res * x
- hvx_mul_scalar_f32_aa(dst_spad_ptr, (const uint8_t *)dst_spad_ptr, 0.5f, nc); // res = res + 0.5f
- hvx_mul_f32_aaa(dst_spad_ptr, (const uint8_t *)dst_spad_ptr, src1_spad_ptr, nc); // res = res * g
- }
-
- dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad), dst_row_size,
- dst_row_size_aligned, block_size);
-
- // prefetch N+2 loop iteration if any
- const uint32_t pref_block = (ir + BLOCK * 2);
- if (pref_block < src0_end_row) {
- const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);
- dma_queue_push_ddr_to_vtcm(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_size)),
- src0_row_size_aligned, src0_row_size, pref_block_size);
- dma_queue_push_ddr_to_vtcm(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_size)),
- src1_row_size_aligned, src1_row_size, pref_block_size);
- }
+#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR) \
+ static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) { \
+ struct htp_act_context * actx = (struct htp_act_context *) data; \
+ htp_act_preamble; \
+ \
+ struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL; \
+ \
+ size_t src0_row_size = actx->src0_row_size; \
+ size_t src1_row_size = actx->src1_row_size; \
+ size_t dst_row_size = actx->dst_row_size; \
+ \
+ size_t src0_row_stride = actx->src0_row_stride; \
+ size_t src1_row_stride = actx->src1_row_stride; \
+ \
+ const uint32_t src0_nrows = actx->src0_nrows; \
+ const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread; \
+ \
+ const uint32_t src0_start_row = src0_nrows_per_thread * ith; \
+ const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); \
+ \
+ /* no work for this thread */ \
+ if (src0_start_row >= src0_end_row) { \
+ return; \
+ } \
+ \
+ const uint8_t * restrict data_src0 = actx->data_src0; \
+ const uint8_t * restrict data_src1 = actx->data_src1; \
+ uint8_t * restrict data_dst = actx->data_dst; \
+ \
+ const size_t src0_row_size_aligned = actx->src0_row_size_aligned; \
+ const size_t src1_row_size_aligned = actx->src1_row_size_aligned; \
+ const size_t dst_row_size_aligned = actx->dst_row_size_aligned; \
+ \
+ uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread); \
+ uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread); \
+ uint8_t * restrict dst_spad_data = actx->vtcm_dst + (ith * actx->vtcm_dst_size_per_thread); \
+ \
+ size_t src0_spad_half_size = actx->src0_spad_half_size; \
+ size_t src1_spad_half_size = actx->src1_spad_half_size; \
+ size_t dst_spad_half_size = actx->dst_spad_half_size; \
+ \
+ const int BLOCK = actx->block; \
+ if (BLOCK == 0) { \
+ FARF(ERROR, \
+ OP_STR \
+ " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", \
+ actx->vtcm_src0_size_per_thread, src0_row_size_aligned); \
+ return; \
+ } \
+ \
+ dma_queue * dma_queue = actx->octx->ctx->dma[ith]; \
+ \
+ /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */ \
+ for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \
+ const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); \
+ \
+ /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */ \
+ dma_queue_push_vtcm_to_ddr(dma_queue, \
+ dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)), \
+ dst_row_size, dst_row_size_aligned, 0); \
+ \
+ dma_queue_push( \
+ dma_queue, \
+ dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)), \
+ src0_row_size_aligned, src0_row_stride, src0_row_size, block_size); \
+ dma_queue_push( \
+ dma_queue, \
+ dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)), \
+ src1_row_size_aligned, src1_row_stride, src1_row_size, block_size); \
+ } \
+ \
+ for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) { \
+ const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); \
+ \
+ float * dst_spad = (float *) dma_queue_pop(dma_queue).src; \
+ float * src0_spad = (float *) dma_queue_pop(dma_queue).dst; \
+ float * src1_spad = (float *) dma_queue_pop(dma_queue).dst; \
+ \
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); \
+ CORE_EXPR; \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); \
+ \
+ dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad), \
+ dst_row_size, dst_row_size_aligned, block_size); \
+ \
+ /* prefetch N+2 loop iteration if any */ \
+ const uint32_t pref_block = (ir + BLOCK * 2); \
+ if (pref_block < src0_end_row) { \
+ const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block); \
+ dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)), \
+ src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size); \
+ dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)), \
+ src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size); \
+ } \
+ } \
+ \
+ dma_queue_flush(dma_queue); \
+ \
}
- dma_queue_flush(dma_queue);
-
- t2 = HAP_perf_get_qtimer_count();
-
- FARF(HIGH, "geglu-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
- ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3,
- (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
-}
+DEFINE_GLU_PER_THREAD(swiglu, "swiglu-f32", swiglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
+DEFINE_GLU_PER_THREAD(swiglu_oai, "swiglu-oai-f32", swiglu_oai_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
+DEFINE_GLU_PER_THREAD(geglu, "geglu-f32", geglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
static int execute_op_activations_f32(struct htp_ops_context * octx) {
const struct htp_tensor * src0 = octx->src[0];
const struct htp_tensor * src1 = octx->src[1];
const struct htp_tensor * dst = octx->dst;
- if (((src0->ne[0] * SIZEOF_FP32) != src0->nb[1]) || ((dst->ne[0] * SIZEOF_FP32) != dst->nb[1])) {
- FARF(ERROR, "Non-contiguous tensors are not supported at this time \n");
+ if ((dst->ne[0] * SIZEOF_FP32) != dst->nb[1]) {
+ FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
return HTP_STATUS_NO_SUPPORT;
}
const char * op_type = NULL;
switch (octx->op) {
- case HTP_OP_UNARY_SILU:
- act_op_func = (worker_callback_t)unary_silu_f32_per_thread;
- op_type = "silu-f32";
- break;
-
case HTP_OP_GLU_SWIGLU:
act_op_func = (worker_callback_t)glu_swiglu_f32_per_thread;
op_type = "swiglu-f32";
act_op_func = (worker_callback_t)glu_swiglu_oai_f32_per_thread;
op_type = "swiglu-oai-f32";
break;
- case HTP_OP_UNARY_GELU:
- act_op_func = (worker_callback_t)unary_gelu_f32_per_thread;
- op_type = "gelu-f32";
- break;
case HTP_OP_GLU_GEGLU:
act_op_func = (worker_callback_t)glu_geglu_f32_per_thread;
const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
const uint32_t n_threads = MIN(octx->n_threads, src0_nrows);
- size_t src0_row_size = src0->nb[1];
- size_t src1_row_size = src1 ? src1->nb[1] : src0->nb[1];
- size_t dst_row_size = dst->nb[1];
+ // row_size = bytes of useful data per row (what the kernel touches / what DMA copies).
+ // row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
+ const size_t nc_bytes = dst->ne[0] * SIZEOF_FP32;
+ const size_t src0_row_size = nc_bytes;
+ const size_t src1_row_size = nc_bytes;
+ const size_t dst_row_size = nc_bytes;
+ const size_t src0_row_stride = src0->nb[1];
+ const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
- // VTCM scratchpads for all tensors
- // N rows per thread, padded to HVX vector size
- size_t spad_size_per_row = (src0_row_size_aligned + src1_row_size_aligned) + dst_row_size_aligned;
- size_t vtcm_row_per_thread = (octx->ctx->vtcm_size)/ (n_threads* spad_size_per_row);
+ struct htp_act_vtcm_layout L;
+ htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
+ octx->ctx->vtcm_size);
// Make sure the reserved vtcm size is sufficient
- if (vtcm_row_per_thread == 0) {
+ if (L.vtcm_row_per_thread == 0) {
FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
- spad_size_per_row * n_threads);
+ (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
return HTP_STATUS_VTCM_TOO_SMALL;
}
- octx->src0_spad.size_per_thread = src0_row_size_aligned * vtcm_row_per_thread;
- octx->src1_spad.size_per_thread = src1_row_size_aligned * vtcm_row_per_thread;
- octx->dst_spad.size_per_thread = dst_row_size_aligned * vtcm_row_per_thread;
-
- octx->dst_spad.size = n_threads* octx->dst_spad.size_per_thread;
- octx->src0_spad.size = n_threads* octx->src0_spad.size_per_thread;
- octx->src1_spad.size = n_threads* octx->src1_spad.size_per_thread;
-
- octx->src0_spad.data = octx->ctx->vtcm_base;
- octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
- octx->dst_spad.data = octx->src1_spad.data + octx->src1_spad.size;
-
- octx->src0_spad.src = NULL;
- octx->src1_spad.src = NULL;
- octx->dst_spad.src = NULL;
-
if (src1) {
- FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n",
+ FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n",
op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
- src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], octx->src0_spad.size, octx->src1_spad.size,
- octx->dst_spad.size);
+ src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], L.src0_bytes_per_thread * n_threads,
+ L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
} else {
- FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n", op_type,
+ FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n", op_type,
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
- octx->src0_spad.size, octx->src1_spad.size, octx->dst_spad.size);
+ L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
}
if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
actx.src1_row_size_aligned = src1_row_size_aligned;
actx.dst_row_size_aligned = dst_row_size_aligned;
- actx.src0_spad_half_size = octx->src0_spad.size_per_thread / 2;
- actx.src1_spad_half_size = octx->src1_spad.size_per_thread / 2;
- actx.dst_spad_half_size = octx->dst_spad.size_per_thread / 2;
+ actx.src0_row_stride = src0_row_stride;
+ actx.src1_row_stride = src1_row_stride;
+
+ uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+ actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+ actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+ actx.vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+
+ actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
+ actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
+ actx.vtcm_dst_size_per_thread = L.dst_bytes_per_thread;
+
+ actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
+ actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
+ actx.dst_spad_half_size = L.dst_bytes_per_thread / 2;
actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
actx.src0_nrows = src0_nrows;
}
int op_activations(struct htp_ops_context * octx) {
- int err = HTP_STATUS_OK;
-
switch (octx->src[0]->type) {
case HTP_TYPE_F32:
- err = execute_op_activations_f32(octx);
- break;
+ return execute_op_activations_f32(octx);
default:
- err = HTP_STATUS_NO_SUPPORT;
- break;
+ return HTP_STATUS_NO_SUPPORT;
}
-
- return err;
}
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)
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->trace[i]; \
+ 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;
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];
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];
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->trace[i];
+ 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;
// 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) {
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],
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;
}
#include "htp-ctx.h"
#include "htp-ops.h"
#include "htp-ops.h"
+#include "htp-tensor.h"
#ifndef MIN
#define MIN(a, b) ((a) < (b) ? (a) : (b))
#include "ggml-common.h"
#include "htp-ctx.h"
#include "htp-ops.h"
+#include "htp-tensor.h"
#include "hvx-types.h"
#include "hvx-utils.h"
#include "hex-dma.h"
int op_cumsum(struct htp_ops_context * octx) {
const struct htp_tensor * dst = octx->dst;
- int err = HTP_STATUS_OK;
-
switch (dst->type) {
case HTP_TYPE_F32:
- err = op_cumsum_f32(octx);
- break;
+ return op_cumsum_f32(octx);
default:
- err = HTP_STATUS_NO_SUPPORT;
- break;
+ return HTP_STATUS_NO_SUPPORT;
}
-
- return err;
}
--- /dev/null
+#include "dma-queue.h"
+
+#include <stdbool.h>
+#include <stdlib.h>
+#include <string.h>
+
+#pragma clang diagnostic ignored "-Wunused-function"
+
+static inline uint32_t pow2_ceil(uint32_t x) {
+ if (x <= 1) {
+ return 1;
+ }
+ int p = 2;
+ x--;
+ while (x >>= 1) {
+ p <<= 1;
+ }
+ return p;
+}
+
+static inline uintptr_t align_up(uintptr_t addr, size_t align) {
+ return (addr + align - 1) & ~(align - 1);
+}
+
+size_t dma_queue_sizeof(size_t capacity) {
+ capacity = pow2_ceil(capacity);
+
+ size_t size_q = sizeof(dma_queue);
+ size_t offset_r = align_up(size_q, HEX_L2_LINE_SIZE);
+ size_t size_r = sizeof(dma_ring);
+ size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE);
+ size_t size_desc = capacity * sizeof(dma_descriptor_2d);
+ size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
+ size_t size_dptr = capacity * sizeof(dma_ptr);
+
+ return offset_dptr + size_dptr;
+}
+
+size_t dma_queue_alignof(void) {
+ return HEX_L2_LINE_SIZE;
+}
+
+dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace) {
+ capacity = pow2_ceil(capacity);
+
+ size_t size_q = sizeof(dma_queue);
+ size_t offset_r = align_up(size_q, HEX_L2_LINE_SIZE);
+ size_t size_r = sizeof(dma_ring);
+ size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE);
+ size_t size_desc = capacity * sizeof(dma_descriptor_2d);
+ size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
+ size_t size_dptr = capacity * sizeof(dma_ptr);
+
+ size_t total_size = offset_dptr + size_dptr;
+ memset(ptr, 0, total_size);
+
+ dma_queue * q = (dma_queue *) ptr;
+ dma_ring * r = (dma_ring *) ((uintptr_t) ptr + offset_r);
+
+ q->ring = r;
+ q->nocache = 0;
+ q->alias = false;
+
+ r->trace = trace;
+ r->vtcm_base = vtcm_base;
+ r->vtcm_end = vtcm_base + vtcm_size;
+ r->capacity = capacity;
+ r->idx_mask = capacity - 1;
+ r->push_idx = 0;
+ r->pop_idx = 0;
+
+ r->desc = (dma_descriptor_2d *) ((uintptr_t) ptr + offset_desc);
+ r->dptr = (dma_ptr *) ((uintptr_t) ptr + offset_dptr);
+ r->tail = &r->desc[capacity - 1];
+
+ FARF(HIGH, "dma-queue: capacity %u, unified memory size %zu\n", capacity, total_size);
+
+ return q;
+}
+
+void dma_queue_free(dma_queue_t q) {
+ (void) q;
+}
+
+size_t dma_queue_alias_sizeof(void) {
+ return sizeof(dma_queue);
+}
+
+dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache) {
+ dma_queue * q = (dma_queue *) ptr;
+ memset(q, 0, sizeof(dma_queue));
+
+ q->ring = main_q->ring;
+ q->nocache = nocache;
+ q->alias = true;
+
+ return q;
+}
+
+void dma_queue_alias_free(dma_queue_t q) {
+ (void) q;
+}
+
+
--- /dev/null
+#ifndef HTP_DMA_H
+#define HTP_DMA_H
+
+#include <HAP_farf.h>
+#include <hexagon_types.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include "hex-utils.h"
+
+#include "hex-profile.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Define the HW descriptor structs here since the ones in HexSDK are a bit out of date
+typedef struct dma_descriptor_1d_s {
+ void * next;
+ uint32_t size:24;
+ uint32_t desc_size:2;
+ uint32_t dst_comp:1;
+ uint32_t src_comp:1;
+ uint32_t dst_bypass:1;
+ uint32_t src_bypass:1;
+ uint32_t order:1;
+ uint32_t done:1;
+ void * src;
+ void * dst;
+} dma_descriptor_1d;
+
+#if __HVX_ARCH__ < 75
+
+typedef struct dma_descriptor_2d_s {
+ void * next;
+ uint32_t reserved0:24;
+ uint32_t desc_size:2;
+ uint32_t dst_comp:1;
+ uint32_t src_comp:1;
+ uint32_t dst_bypass:1;
+ uint32_t src_bypass:1;
+ uint32_t order:1;
+ uint32_t done:1;
+ void * src;
+ void * dst;
+ uint32_t desc_type:8;
+ uint32_t reserved1:24;
+ uint32_t row_size:16;
+ uint32_t nrows:16;
+ uint32_t src_stride:16;
+ uint32_t dst_stride:16;
+ uint32_t src_offset:16;
+ uint32_t dst_offset:16;
+} dma_descriptor_2d;
+
+#else
+
+typedef struct dma_descriptor_2d_s {
+ void * next;
+ uint32_t dst_stride:24;
+ uint32_t desc_size:2;
+ uint32_t dst_comp:1;
+ uint32_t src_comp:1;
+ uint32_t dst_bypass:1;
+ uint32_t src_bypass:1;
+ uint32_t order:1;
+ uint32_t done:1;
+ void * src;
+ void * dst;
+ uint32_t desc_type:8;
+ uint32_t reserved0:24;
+ uint32_t row_size:24;
+ uint32_t nrows_lo:8;
+ uint32_t nrows_hi:8;
+ uint32_t src_stride:24;
+ uint32_t offset:24;
+ uint32_t reserved1:8;
+} dma_descriptor_2d;
+
+#endif
+
+typedef struct {
+ void *dst;
+ const void *src;
+} dma_ptr;
+
+typedef struct dma_ring_s dma_ring;
+struct dma_ring_s {
+ dma_descriptor_2d * desc; // descriptor pointers
+ dma_descriptor_2d * tail; // tail pointer
+ dma_ptr * dptr; // dst/src pointers
+ uint32_t push_idx;
+ uint32_t pop_idx;
+ uint32_t capacity;
+ uint32_t idx_mask;
+ struct htp_thread_trace * trace;
+ uintptr_t vtcm_base;
+ uintptr_t vtcm_end;
+};
+
+typedef struct dma_queue_s dma_queue;
+typedef dma_queue * dma_queue_t;
+
+struct dma_queue_s {
+ dma_ring * ring; // Points to the descriptor ring state
+ uint8_t nocache; // Queue-specific bypass flag
+ bool alias; // When set, dma_queue_delete will not free the ring
+};
+
+
+
+size_t dma_queue_sizeof(size_t capacity);
+size_t dma_queue_alignof(void);
+dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace);
+void dma_queue_free(dma_queue_t q);
+
+size_t dma_queue_alias_sizeof(void);
+dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache);
+void dma_queue_alias_free(dma_queue_t q);
+
+// TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
+// but those do not seem to always compiler properly.
+static inline void dmstart(void * next) {
+ asm volatile(" release(%0):at" : : "r"(next));
+ asm volatile(" dmstart(%0)" : : "r"(next));
+}
+
+static inline void dmlink(void * cur, void * next) {
+ asm volatile(" release(%0):at" : : "r"(next));
+ asm volatile(" dmlink(%0, %1)" : : "r"(cur), "r"(next));
+}
+
+static inline unsigned int dmpoll(void) {
+ unsigned int ret = 0;
+ asm volatile(" %0 = dmpoll" : "=r"(ret) : : "memory");
+ return ret;
+}
+
+static inline unsigned int dmwait(void) {
+ unsigned int ret = 0;
+ asm volatile(" %0 = dmwait" : "=r"(ret) : : "memory");
+ return ret;
+}
+
+static inline dma_ptr dma_make_ptr(void *dst, const void *src)
+{
+ dma_ptr p = { dst, src };
+ return p;
+}
+
+static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) {
+ return (uintptr_t) ptr >= q->ring->vtcm_base && (uintptr_t) ptr < q->ring->vtcm_end;
+}
+
+static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
+ dma_ring * r = q->ring;
+ if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
+ return false;
+ }
+
+ dma_descriptor_1d * desc = (dma_descriptor_1d *) &r->desc[r->push_idx];
+ desc->src = (void *) dptr.src;
+ desc->dst = (void *) dptr.dst;
+ desc->size = size;
+
+ r->dptr[r->push_idx] = dptr;
+
+ htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
+
+ if (size) {
+ desc->next = NULL;
+ desc->desc_size = 0; // 1D mode
+ desc->src_bypass = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache;
+ desc->dst_bypass = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache;
+ desc->order = 0;
+ desc->done = 0;
+
+ dmlink(r->tail, desc);
+ r->tail = (dma_descriptor_2d *) desc;
+ } else {
+ desc->desc_size = 0;
+ desc->done = 1;
+ }
+
+ r->push_idx = (r->push_idx + 1) & r->idx_mask;
+ return true;
+}
+
+static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+ dma_ring * r = q->ring;
+ if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
+ return false;
+ }
+
+ dma_descriptor_2d * desc = &r->desc[r->push_idx];
+
+ desc->next = NULL;
+ desc->reserved0 = 0;
+ desc->reserved1 = 0;
+ desc->desc_size = 1; // 2d mode
+ desc->src_bypass = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache;
+ desc->dst_bypass = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache;
+ desc->src_comp = 0;
+ desc->dst_comp = 0;
+ desc->order = 0;
+ desc->done = 0;
+ desc->src_stride = src_stride;
+ desc->dst_stride = dst_stride;
+ desc->src = (void *) dptr.src;
+ desc->dst = (void *) dptr.dst;
+ desc->row_size = row_size;
+
+#if __HVX_ARCH__ < 75
+ desc->desc_type = 0; // 2d (16-bit) mode
+ desc->nrows = nrows;
+ desc->src_offset = 0;
+ desc->dst_offset = 0;
+#else
+ desc->desc_type = 9; // 2d (24-bit) mode
+ desc->nrows_lo = (nrows & 0xff);
+ desc->nrows_hi = (nrows >> 8);
+ desc->offset = 0;
+#endif
+
+ r->dptr[r->push_idx] = dptr;
+
+ htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
+
+ if (nrows) {
+ dmlink(r->tail, desc);
+ r->tail = desc;
+ } else {
+ desc->done = 1;
+ }
+
+ r->push_idx = (r->push_idx + 1) & r->idx_mask;
+ return true;
+}
+
+static inline dma_ptr dma_queue_pop(dma_queue * q) {
+ dma_ring * r = q->ring;
+ dma_ptr dptr = { NULL };
+
+ if (r->push_idx == r->pop_idx) {
+ return dptr;
+ }
+
+ dma_descriptor_2d * desc = &r->desc[r->pop_idx];
+
+ // Wait for desc to complete
+ if (!desc->done) {
+ while (!desc->done) {
+ dmpoll();
+ }
+ }
+
+ dptr = r->dptr[r->pop_idx];
+
+ htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
+
+ r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
+ return dptr;
+}
+
+static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
+ dma_ring * r = q->ring;
+ dma_ptr dptr = { NULL };
+
+ if (r->push_idx == r->pop_idx) {
+ return dptr;
+ }
+
+ dptr = r->dptr[r->pop_idx];
+
+ htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
+
+ r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
+ return dptr;
+}
+
+static inline bool dma_queue_empty(dma_queue * q) {
+ return q->ring->push_idx == q->ring->pop_idx;
+}
+
+static inline void dma_queue_flush(dma_queue * q) {
+ while (dma_queue_pop(q).dst != NULL) ;
+}
+
+static inline uint32_t dma_queue_depth(dma_queue * q) {
+ return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask;
+}
+
+static inline uint32_t dma_queue_capacity(dma_queue * q) {
+ return q->ring->capacity;
+}
+
+#if __HVX_ARCH__ < 75
+
+// Overflow-safe DMA push: all 2d descriptor fields (row_size, nrows, src_stride, dst_stride) are 16-bit, max 65535.
+// This version transparently handles values that exceed the 16-bit limit and submits chained DMA transtions.
+
+#define DMA_MAX_FIELD_VAL 65535u
+
+static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+ // Fast path: everything fits in 16 bits
+ if (nrows == 0 || __builtin_expect(
+ row_size <= DMA_MAX_FIELD_VAL &&
+ nrows <= DMA_MAX_FIELD_VAL &&
+ src_stride <= DMA_MAX_FIELD_VAL &&
+ dst_stride <= DMA_MAX_FIELD_VAL, 1)) {
+ return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
+ }
+
+ // Contiguous block
+ // Use 1d DMA mode which supports sizes up to 24-bits (16MB)
+ if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
+ size_t total = row_size * nrows;
+ return dma_queue_push_single_1d(q, dptr, total);
+ }
+
+ // Stride overflow - fall back to row-by-row.
+ {
+ const uint8_t *src = (const uint8_t *) dptr.src;
+ uint8_t *dst = (uint8_t *) dptr.dst;
+ size_t r = 0;
+ while (r + 1 < nrows) {
+ dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
+ if (!dma_queue_push_single_1d(q, p, row_size)) {
+ dma_queue_flush(q);
+ } else {
+ r++;
+ }
+ }
+ dma_queue_flush(q);
+ dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
+ return dma_queue_push_single_1d(q, p, row_size);
+ }
+}
+
+#else // HVX_ARCH >= 75
+
+static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+ // On v75 and up we always use 2d 24-bit mode
+ return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
+}
+
+#endif
+
+static inline bool dma_queue_push_ddr_to_vtcm(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
+ return dma_queue_push(q, dptr, dst_row_size, src_row_size, src_row_size, nrows);
+}
+
+static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
+ return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
+}
+
+#define DMA_CACHE_MAX_SIZE 256U
+
+typedef struct {
+ uint8_t *base;
+ uint32_t line_size;
+ uint32_t capacity;
+ uint32_t src[DMA_CACHE_MAX_SIZE];
+ uint16_t age[DMA_CACHE_MAX_SIZE];
+} dma_cache;
+
+static inline void dma_cache_init(dma_cache *c, uint8_t *base, uint32_t line_size, uint32_t capacity)
+{
+ c->capacity = (capacity > DMA_CACHE_MAX_SIZE) ? DMA_CACHE_MAX_SIZE : capacity;
+ c->base = base;
+ c->line_size = line_size;
+
+ for (unsigned i=0; i < c->capacity; i++) {
+ c->src[i] = 0;
+ c->age[i] = 0;
+ }
+}
+
+static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * src, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows)
+{
+ uint32_t o_idx = 0;
+ uint16_t o_age = 0;
+ uint8_t * dst = 0;
+
+ for (unsigned i=0; i < c->capacity; i++) {
+ if (c->src[i] == (uint32_t) src) {
+ c->age[i] = 0;
+ dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
+ } else {
+ c->age[i]++;
+ if (c->age[i] > o_age) { o_age = c->age[i]; o_idx = i; }
+ }
+ }
+ if (!dst) {
+ c->age[o_idx] = 0;
+ c->src[o_idx] = (uint32_t) src;
+ dst = c->base + o_idx * c->line_size; // normal nrows dma
+ return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
+ }
+
+ return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0);
+}
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#endif /* HTP_DMA_H */
#include "hvx-reduce.h"
#include "hvx-flash-attn.h"
#include "htp-vtcm.h"
-#include "worker-pool.h"
+#include "work-queue.h"
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
uint32_t g_br; // hex_align_up(G * Br, 32) - actual tile row dim
// VTCM buffers (allocated by vtcm_seq_alloc)
+ __fp16 * vtcm_q_dma; // Q DMA fetch buffer
__fp16 * vtcm_q_tiles; // Q tile format [g_br, D]
__fp16 * vtcm_o_tiles[2]; // O ping-pong [g_br, D]
__fp16 * vtcm_k_fp16[2]; // K DMA double-buffer [Bc, D]
__fp16 * vtcm_v_fp16[2]; // V DMA double-buffer [Bc, D]
- __fp16 * vtcm_k_tiles; // K tiles (transposed)
+ __fp16 * vtcm_k_tiles[2]; // K tiles (transposed, double-buffered)
__fp16 * vtcm_v_tiles[2]; // V tiles (column-major, double-buffered)
- __fp16 * vtcm_s_tiles; // S = QK^T [g_br, Bc]
- __fp16 * vtcm_p_tiles; // P = softmax(S) [g_br, Bc]
+ __fp16 * vtcm_s_tiles[2]; // S = QK^T [g_br, Bc] (double-buffered)
+ __fp16 * vtcm_p_tiles[2]; // P = softmax(S) [g_br, Bc]
__fp16 * vtcm_d_tiles; // Diagonal rescale [g_br, g_br]
+ __fp16 * vtcm_d_inv_l; // Diagonal rescale (1/l) [g_br, g_br]
HVX_Vector * vtcm_m_vec; // Row max [g_br]
HVX_Vector * vtcm_l_vec; // Row sum [g_br]
HVX_Vector * vtcm_s_rowmax; // Softmax intermediate [g_br]
if (ir0 >= ir1) return;
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
dma_queue * dma = octx->ctx->dma[ith];
const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3);
const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2);
- // Fetch Q row
- const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
- dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
-
const __fp16 * mp_base = NULL;
if (mask) {
const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]);
}
- // Prefetch first two blocks
- for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
- const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
- const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+ // Precalculate next row variables if there is a next row
+ bool has_next_ir = (ir + 1 < ir1);
+ uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0;
+ const uint8_t * next_q_row_ptr = NULL;
+ const __fp16 * next_mp_base = NULL;
+
+ const uint8_t * next_k_src0 = NULL;
+ const uint8_t * next_v_src0 = NULL;
+ const uint8_t * next_m_src0 = NULL;
+ uint32_t next_block_size0 = 0;
+
+ const uint8_t * next_k_src1 = NULL;
+ const uint8_t * next_v_src1 = NULL;
+ const uint8_t * next_m_src1 = NULL;
+ uint32_t next_block_size1 = 0;
+
+ if (has_next_ir) {
+ const uint32_t next_ir = ir + 1;
+ const uint32_t next_iq3 = fastdiv(next_ir, &factx->src0_div21);
+ const uint32_t next_iq2 = fastdiv(next_ir - next_iq3*neq2*neq1, &factx->src0_div1);
+ const uint32_t next_iq1 = (next_ir - next_iq3*neq2*neq1 - next_iq2 * neq1);
- // K
- const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
- uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
- dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
+ next_ik3 = fastdiv(next_iq3, &factx->broadcast_rk3);
+ next_ik2 = fastdiv(next_iq2, &factx->broadcast_rk2);
- // V
- const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
- uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
- dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
+ next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3);
+ next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2);
+
+ next_q_row_ptr = (const uint8_t *) q->data + (next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3);
- // Mask
if (mask) {
- const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
- // Mask is 1D contiguous for this row
- dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+ const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2);
+ const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3);
+ next_mp_base = (const __fp16 *) ((const uint8_t *) mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]);
+ }
+
+ // Precalculate next K/V block 0 source pointers
+ {
+ const uint32_t ic_start = 0;
+ next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+ next_k_src0 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
+ next_v_src0 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
+ if (mask) {
+ next_m_src0 = (const uint8_t *) (next_mp_base + ic_start);
+ }
+ }
+
+ // Precalculate next K/V block 1 source pointers (if n_blocks > 1)
+ if (factx->n_blocks > 1) {
+ const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE;
+ next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+ next_k_src1 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
+ next_v_src1 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
+ if (mask) {
+ next_m_src1 = (const uint8_t *) (next_mp_base + ic_start);
+ }
+ }
+ }
+
+ if (ir == ir0) {
+ // Fetch Q row
+ const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
+ dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
+
+ // Prefetch first two blocks
+ for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
+ const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
+ const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+
+ // K
+ const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
+ uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
+ dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
+
+ // V
+ const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
+ uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
+ dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
+
+ // Mask
+ if (mask) {
+ const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
+ // Mask is 1D contiguous for this row
+ dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+ }
}
}
const HVX_Vector slope_vec = hvx_vec_splat_f16(slope);
const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
+ const HVX_Vector v_cap = (factx->logit_softcap != 0.0f) ? hvx_vec_splat_f16(factx->logit_softcap) : Q6_V_vzero();
+ const HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
+ const HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
+ const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
+ const uint32_t stride_v2 = factx->size_v_row_padded * 2;
for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) {
const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
// 2. Softcap (in FP16)
if (factx->logit_softcap != 0.0f) {
- const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
scores_f16 = hvx_vec_tanh_f16(scores_f16);
scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap);
}
// 3. Mask (in FP16)
if (mask) {
HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base;
- HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
- HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf);
m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16);
HVX_Vector v_max = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); // splat block max in FP32
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);
+ if (ib + 1 == factx->n_blocks && has_next_ir) {
+ // Queue next row's Q row!
+ dma_queue_push(dma, dma_make_ptr(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
+
+ if (factx->n_blocks % 2 == 0) {
+ // Queue next row's block 0 (into buffer slot 0)
+ uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
+ uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
+
+ // K (block 0 of next row)
+ dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
+
+ // V (block 0 of next row)
+ dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
+
+ // Mask (block 0 of next row)
+ if (mask) {
+ dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
+ }
+ }
+ }
+
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
{
- const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
-
// 4. Online Softmax Update
HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec);
HVX_Vector diff_vec = HVX_OP_SUB_F32(M_vec, M_new_vec);
S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec);
// 5. Accumulate V (F16 * F16 -> F32 accumulator)
- __fp16 __attribute__((aligned(128))) p_arr[VLEN_FP16];
- hvx_vec_store_a(p_arr, 128, P);
+ const uint8_t * v_ptr = v_base;
for (uint32_t j = 0; j < current_block_size; j += 2) {
if (j + 1 == current_block_size) {
- if (p_arr[j] != 0.0f) {
- const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded;
- hvx_mad_f32_f16_aa(VKQ32, v_ptr, (p_arr + j), DV);
- }
+ HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
+ hvx_mad_f32_f16_aa_vec(VKQ32, v_ptr, S0, DV);
break;
}
- if (p_arr[j] == 0.0f && p_arr[j + 1] == 0.0f) {
- continue;
- }
+ HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
+ HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2));
- const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded;
- hvx_mad_f32_f16_aa_rx2(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, (p_arr + j), (p_arr + j + 1), DV);
+ hvx_mad_f32_f16_aa_rx2_vec(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, S0, S1, DV);
+ v_ptr += stride_v2;
}
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
}
}
+ if (has_next_ir) {
+ if (factx->n_blocks % 2 == 0) {
+ // Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
+ if (factx->n_blocks > 1) {
+ uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
+ uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
+
+ // K (block 1 of next row)
+ dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
+
+ // V (block 1 of next row)
+ dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
+
+ // Mask (block 1 of next row)
+ if (mask) {
+ dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
+ }
+ }
+ } else {
+ // Queue next row's block 0 (into buffer slot 0)
+ {
+ uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
+ uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
+
+ // K (block 0 of next row)
+ dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
+
+ // V (block 0 of next row)
+ dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
+
+ // Mask (block 0 of next row)
+ if (mask) {
+ dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
+ }
+ }
+
+ // Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
+ if (factx->n_blocks > 1) {
+ uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
+ uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
+
+ // K (block 1 of next row)
+ dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
+
+ // V (block 1 of next row)
+ dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
+
+ // Mask (block 1 of next row)
+ if (mask) {
+ dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
+ }
+ }
+ }
+ }
+
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir);
// sinks
float M = hvx_vec_get_f32(M_vec);
void * curr_k;
uint32_t kv_start;
uint32_t rows_per_t;
+ size_t buf_idx;
} fa_k_int_args_t;
static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
return;
}
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
- hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK,
+ hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
args->src_stride, start, end);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
}
-static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start, size_t buf_idx) {
+ work_queue_t wp = factx->octx->ctx->work_queue;
uint32_t n = 1;
if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
n = factx->n_threads;
}
uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
- fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t };
+ fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t, buf_idx };
if (n > 1) {
- worker_pool_run_func(wp, fa_k_interleave_thread, &args, n);
+ work_queue_run(wp, fa_k_interleave_thread, &args, n);
} else {
fa_k_interleave_thread(1, 0, &args);
}
__fp16 * v_tiles_dst = (__fp16 *) args->v_tiles_dst;
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
args->src_stride, (uint32_t) args->n_col_tiles, start, end);
void * v_tiles_dst,
size_t n_col_tiles,
uint32_t kv_start) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ work_queue_t wp = factx->octx->ctx->work_queue;
uint32_t n = 1;
if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
n = factx->n_threads;
uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
fa_v_int_args_t args = { factx, kv_rows, src_stride, v_src, v_tiles_dst, n_col_tiles, kv_start, rows_per_t };
if (n > 1) {
- worker_pool_run_func(wp, fa_v_interleave_thread, &args, n);
+ work_queue_run(wp, fa_v_interleave_thread, &args, n);
} else {
fa_v_interleave_thread(1, 0, &args);
}
const size_t start = (size_t) i * rows_per_t;
const size_t end = hex_smin(start + rows_per_t, factx->g_br);
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
// Parallel initialization of per-block state
}
}
- // Initialize vtcm_d_tiles to 0
+ // Initialize vtcm_d_tiles and vtcm_d_inv_l to 0
const size_t d_bytes_per_t = hex_align_up(d_tile_bytes / n, 128);
const size_t d_start = i * d_bytes_per_t;
const size_t d_end = hex_smin(d_start + d_bytes_per_t, d_tile_bytes);
if (d_start < d_tile_bytes) {
hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
+ hvx_splat_u8_a((char *) factx->vtcm_d_inv_l + d_start, 0, d_end - d_start);
}
}
assert(factx->DK == factx->DV);
- const size_t o_tile_bytes = factx->o_tile_bytes;
- const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2));
+ const bool use_q_dma = (factx->vtcm_q_dma != NULL);
__fp16 * q_tiles = factx->vtcm_q_tiles;
if (use_q_dma) {
const size_t g_rows_end = hex_smin(end, n_rows_g);
const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;
- uint8_t * q_flat = (uint8_t *) factx->vtcm_o_tiles[0];
+ uint8_t * q_flat = (uint8_t *) factx->vtcm_q_dma;
if (factx->is_q_fp32) {
switch (d_limit) {
case 2: hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
uint32_t kv_head,
uint32_t ib3,
size_t n_rows_g) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ work_queue_t wp = factx->octx->ctx->work_queue;
uint32_t n = 1;
if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
n = factx->n_threads;
args.q_transposed = q->nb[1] < q->nb[2];
atomic_init(&args.barrier, n);
if (n > 1) {
- worker_pool_run_func(wp, fa_q_load_thread, &args, n);
+ work_queue_run(wp, fa_q_load_thread, &args, n);
} else {
fa_q_load_thread(1, 0, &args);
}
return;
}
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
const struct htp_tensor * dst = args->dst;
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
- for (size_t r = start; r < end; ++r) {
- const size_t q_idx = fastdiv(r, &factx->div_G);
- const size_t h_idx = fastmodulo(r, G, &factx->div_G);
+ size_t q_idx = fastdiv(start, &factx->div_G);
+ size_t h_idx = fastmodulo(start, G, &factx->div_G);
+ for (size_t r = start; r < end; ++r) {
float * out = (float *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
(q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);
*(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
}
}
+
+ h_idx++;
+ if (h_idx == G) {
+ h_idx = 0;
+ q_idx++;
+ }
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
}
return;
}
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
const struct htp_tensor * dst = args->dst;
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
- for (size_t r = start; r < end; ++r) {
- const size_t q_idx = fastdiv(r, &factx->div_G);
- const size_t h_idx = fastmodulo(r, G, &factx->div_G);
+ size_t q_idx = fastdiv(start, &factx->div_G);
+ size_t h_idx = fastmodulo(start, G, &factx->div_G);
+ for (size_t r = start; r < end; ++r) {
__fp16 * out = (__fp16 *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
(q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);
*(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
}
}
+
+ h_idx++;
+ if (h_idx == G) {
+ h_idx = 0;
+ q_idx++;
+ }
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
}
uint32_t kv_head,
uint32_t ib3,
size_t n_rows_g) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ work_queue_t wp = factx->octx->ctx->work_queue;
uint32_t n = 1;
if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
n = factx->n_threads;
fa_o_store_args_t args = { factx, dst, o_tile_src, q_start, kv_head, ib3, n_rows_g, rows_per_t };
worker_callback_t store_fn = factx->is_dst_fp32 ? fa_o_store_thread_f32 : fa_o_store_thread_f16;
if (n > 1) {
- worker_pool_run_func(wp, store_fn, &args, n);
+ work_queue_run(wp, store_fn, &args, n);
} else {
store_fn(1, 0, &args);
}
typedef struct {
struct hmx_fa_context * factx;
+ size_t buf_idx;
size_t kv_rows;
size_t n_rows_g;
size_t n_col_tiles;
return;
}
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, (uint16_t) (args->q_start * G + vec_start * 64));
// Per-thread row scratch: thread i uses bufs at offset i * 2 * stride
uint32_t r0 = r / HMX_FP16_TILE_N_ROWS;
uint32_t r1 = r % HMX_FP16_TILE_N_ROWS;
- const __fp16 * s_ld_base = factx->vtcm_s_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc;
- __fp16 * p_st_base = factx->vtcm_p_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc;
+ const __fp16 * s_ld_base = factx->vtcm_s_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
+ __fp16 * p_st_base = factx->vtcm_p_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
// Decode 2 rows from S tiles into per-thread row buffers
if (has_softcap) {
my_row_buf1[ci] = hvx_vec_mul_f16_f16(t1, v_cap);
}
} else {
- for (size_t c = 0; c < kv_rows; c += 64) {
+ size_t c = 0;
+ for (; c + 64 < kv_rows; c += 128) {
+ size_t ci0 = c / 64;
+ size_t ci1 = ci0 + 1;
+ const __fp16 * in_dtile0 = s_ld_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
+ const __fp16 * in_dtile1 = s_ld_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_s_in0_0 = ((const HVX_Vector *) in_dtile0) + r1 / 2;
+ const HVX_Vector * pv_s_in1_0 = pv_s_in0_0 + 16;
+ const HVX_Vector * pv_s_in0_1 = ((const HVX_Vector *) in_dtile1) + r1 / 2;
+ const HVX_Vector * pv_s_in1_1 = pv_s_in0_1 + 16;
+
+ HVX_VectorPair vp_s_drow0 = Q6_W_vdeal_VVR(*pv_s_in1_0, *pv_s_in0_0, -2);
+ my_row_buf0[ci0] = Q6_V_lo_W(vp_s_drow0);
+ my_row_buf1[ci0] = Q6_V_hi_W(vp_s_drow0);
+
+ HVX_VectorPair vp_s_drow1 = Q6_W_vdeal_VVR(*pv_s_in1_1, *pv_s_in0_1, -2);
+ my_row_buf0[ci1] = Q6_V_lo_W(vp_s_drow1);
+ my_row_buf1[ci1] = Q6_V_hi_W(vp_s_drow1);
+ }
+ for (; c < kv_rows; c += 64) {
size_t ci = c / 64;
const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
HVX_Vector v_s_rowmax0 = v_neg_inf;
HVX_Vector v_s_rowmax1 = v_neg_inf;
- for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 64;
- const size_t ne = hex_smin(kv_rows - c, 64);
- HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
+ if (has_mask) {
+ for (size_t c = 0; c < kv_rows; c += 64) {
+ size_t ci = c / 64;
+ const size_t ne = hex_smin(kv_rows - c, 64);
+ HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
- if (has_mask) {
HVX_Vector v_mask0, v_mask1;
if (mask_broadcast) {
my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_mask0_scaled), v_neg_inf);
my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_mask1_scaled), v_neg_inf);
}
- } else {
+
+ v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
+ v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
+ }
+ } else {
+ size_t c = 0;
+ for (; c + 64 < kv_rows; c += 128) {
+ size_t ci0 = c / 64;
+ size_t ci1 = ci0 + 1;
+ v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci0]);
+ v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci0]);
+ v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci1]);
+ v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci1]);
+ }
+ for (; c < kv_rows; c += 64) {
+ size_t ci = c / 64;
+ const size_t ne = hex_smin(kv_rows - c, 64);
+ HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
if (ne < 64) {
my_row_buf0[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf0[ci], v_neg_inf);
my_row_buf1[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf1[ci], v_neg_inf);
}
+ v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
+ v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
}
-
- v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
- v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
}
v_s_rowmax0 = hvx_vec_reduce_max_f16(v_s_rowmax0);
HVX_Vector v_p_rowsum0 = v_zero;
HVX_Vector v_p_rowsum1 = v_zero;
- for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 64;
+ size_t c = 0;
+ for (; c + 64 < kv_rows; c += 128) {
+ size_t ci0 = c / 64;
+ size_t ci1 = ci0 + 1;
+
+ HVX_Vector v_s_minus_m0_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci0], v_dup_m0);
+ HVX_Vector v_s_minus_m1_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci0], v_dup_m1);
+ HVX_Vector v_s_minus_m0_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci1], v_dup_m0);
+ HVX_Vector v_s_minus_m1_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci1], v_dup_m1);
+
+ HVX_Vector v_p_row0_hf_0 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_0));
+ HVX_Vector v_p_row1_hf_0 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_0));
+ HVX_Vector v_p_row0_hf_1 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_1));
+ HVX_Vector v_p_row1_hf_1 = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_1));
+
+ __fp16 * out_dtile0 = p_st_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
+ __fp16 * out_dtile1 = p_st_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_p_out0_0 = ((HVX_Vector *) out_dtile0) + r1 / 2;
+ HVX_Vector * pv_p_out1_0 = pv_p_out0_0 + 16;
+ HVX_Vector * pv_p_out0_1 = ((HVX_Vector *) out_dtile1) + r1 / 2;
+ HVX_Vector * pv_p_out1_1 = pv_p_out0_1 + 16;
+
+ HVX_VectorPair vp_p_dual0 = Q6_W_vshuff_VVR(v_p_row1_hf_0, v_p_row0_hf_0, -2);
+ *pv_p_out0_0 = Q6_V_lo_W(vp_p_dual0);
+ *pv_p_out1_0 = Q6_V_hi_W(vp_p_dual0);
+
+ HVX_VectorPair vp_p_dual1 = Q6_W_vshuff_VVR(v_p_row1_hf_1, v_p_row0_hf_1, -2);
+ *pv_p_out0_1 = Q6_V_lo_W(vp_p_dual1);
+ *pv_p_out1_1 = Q6_V_hi_W(vp_p_dual1);
+
+ HVX_VectorPair vp_p0_0 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_0);
+ HVX_VectorPair vp_p1_0 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_0);
+ HVX_VectorPair vp_p0_1 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_1);
+ HVX_VectorPair vp_p1_1 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_1);
+
+ v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_0), Q6_V_hi_W(vp_p0_0)));
+ v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_1), Q6_V_hi_W(vp_p0_1)));
+ v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_0), Q6_V_hi_W(vp_p1_0)));
+ v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_1), Q6_V_hi_W(vp_p1_1)));
+ }
+ for (size_t c_rem = c; c_rem < kv_rows; c_rem += 64) {
+ size_t ci = c_rem / 64;
HVX_Vector v_s_minus_m0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci], v_dup_m0);
HVX_Vector v_s_minus_m1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci], v_dup_m1);
v_content = Q6_V_vror_VR(v_content, 64);
}
- __fp16 * out_base = factx->vtcm_d_tiles + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ __fp16 * out_base = factx->vtcm_d_inv_l + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
}
}
fa_softmax_args_t * sargs,
size_t n_row_tiles,
size_t n_row_tiles_g_br) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ work_queue_t wp = factx->octx->ctx->work_queue;
const size_t n_row_vec_cnt = hmx_ceil_div(sargs->n_rows_g, 64);
worker_callback_t softmax_fn = fa_softmax_thread;
if (factx->n_threads > 1 && n_row_vec_cnt >= 2) {
uint32_t n_use = (uint32_t) hex_smin((size_t) factx->n_threads, n_row_vec_cnt);
sargs->thread_div = init_fastdiv_values(n_use);
- worker_pool_run_func(wp, softmax_fn, sargs, n_use);
+ work_queue_run(wp, softmax_fn, sargs, n_use);
} else {
softmax_fn(1, 0, sargs);
}
}
}
+static inline void fa_prefetch_block(dma_queue * dma, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask,
+ uint32_t b, size_t Bc, size_t size_k_row_padded, size_t size_k_row, size_t size_v_row_padded, size_t size_v_row,
+ uint32_t ik2, uint32_t ik3, uint32_t iv2, uint32_t iv3, uint32_t q_start, uint32_t im3, uint32_t kv_head, uint32_t G,
+ size_t m_line_bytes, size_t n_rows_q, size_t nek1, size_t prefetch_buf, struct hmx_fa_context * factx) {
+ const uint32_t prefetch_start = b * Bc;
+ const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
+ const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
+ const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
+
+ if (mask) {
+ if (__builtin_expect(factx->mask_broadcast, true)) {
+ const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
+ dma_cache_push(dma, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+ } else {
+ fa_push_mask_dma_gqa(dma, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
+ }
+ }
+}
+
// ============================================================================
// Core HMX flash attention algorithm (GQA-merged)
// ============================================================================
int hmx_flash_attn_ext(struct htp_ops_context * octx) {
- struct htp_thread_trace * tr_hvx = octx->ctx ? &octx->ctx->trace[0] : NULL;
- struct htp_thread_trace * tr_hmx = octx->ctx ? &octx->ctx->trace[HTP_MAX_NTHREADS] : NULL;
+ struct htp_thread_trace * tr_hvx = &octx->ctx->trace[0];
+ struct htp_thread_trace * tr_hmx = &octx->ctx->trace[HTP_MAX_NTHREADS];
const struct htp_tensor * q = octx->src[0];
const struct htp_tensor * k = octx->src[1];
const struct htp_tensor * v = octx->src[2];
// Build the VTCM layout once (shared with the host estimator) and place every
// scratch buffer at its computed offset.
struct hmx_fa_vtcm_layout L;
- hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline);
+ hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline, factx.is_q_fp32);
if (L.total_bytes > ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL;
uint8_t * const base = ctx->vtcm_base;
+ factx.vtcm_q_dma = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_dma);
factx.vtcm_q_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles);
factx.vtcm_o_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]);
factx.vtcm_o_tiles[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]);
factx.vtcm_k_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]);
factx.vtcm_v_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]);
factx.vtcm_v_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]);
- factx.vtcm_k_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles);
+ factx.vtcm_k_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles[0]);
+ factx.vtcm_k_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_k_tiles[1], pipeline);
factx.vtcm_v_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]);
factx.vtcm_v_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline);
- factx.vtcm_s_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles);
- factx.vtcm_p_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles);
+ factx.vtcm_s_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles[0]);
+ factx.vtcm_s_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_s_tiles[1], pipeline);
+ factx.vtcm_p_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles[0]);
+ factx.vtcm_p_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_p_tiles[1], pipeline);
factx.vtcm_d_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
+ factx.vtcm_d_inv_l = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_inv_l);
factx.vtcm_m_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
factx.vtcm_l_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
factx.vtcm_s_rowmax = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax);
const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
+ const bool q_transposed = q->nb[1] < q->nb[2];
+ const size_t q_src_stride = q_transposed ? q->nb[2] : q->nb[1];
+ const size_t q_row_bytes_untransposed = factx.G * factx.DK * qo_element_size;
+ const size_t q_row_bytes_trans_factor = factx.DK * qo_element_size;
+ const uint32_t kv_rows0 = hex_smin(Bc, nek1);
+
// ======== Reusable job descriptors for pipeline ========
hmx_fa_qk_job_t qk_job;
hmx_fa_o_update_job_t ou_job;
const uint32_t iv2 = kv_head;
const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);
- // 1. Push Q DMA (if Q DMA is used)
- const size_t o_tile_bytes = factx.o_tile_bytes;
- const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2));
- if (use_q_dma) {
- const bool q_transposed = q->nb[1] < q->nb[2];
- const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
- const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
- const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size;
- const size_t src_stride = q_transposed ? q->nb[2] : q->nb[1];
+ // 1. Push Q and KV DMAs for the very first iteration.
+ // Subsequent iterations are enqueued early at the end of the previous iteration.
+ if (ib3 == 0 && q_start == 0 && kv_head == 0) {
+ const uint8_t * q_ptr = (const uint8_t *) q->data;
+ const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
const size_t n_rows = q_transposed ? factx.G : n_rows_q;
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows);
- }
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
- // 2. Prefetch first KV block
- if (factx.n_kv_blocks > 0) {
- const uint32_t kv_rows0 = hex_smin(Bc, nek1);
+ if (factx.n_kv_blocks > 0) {
+ const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
- const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+ const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
- const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+ if (factx.pipeline && mask) {
+ if (__builtin_expect(factx.mask_broadcast, true)) {
+ const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
+ dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
+ } else {
+ fa_push_mask_dma_gqa(dma, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
+ }
+ }
+ }
}
- // 3. Pop Q DMA (blocks until Q is loaded)
- if (use_q_dma) {
- dma_queue_pop(dma);
- }
+ // 2. Pop Q DMA (blocks until Q is loaded)
+ dma_queue_pop(dma);
// ---- Load Q block & Initialize per-block state ----
fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
const size_t k_src_stride = size_k_row_padded / sizeof(__fp16);
const size_t v_src_stride = size_v_row_padded / sizeof(__fp16);
- struct hmx_queue * hmx_q = ctx->hmx_queue;
+ hmx_queue_t hmx_q = ctx->hmx_queue;
if (factx.pipeline) {
- // Pipeline path
+ // Double-buffered job structs because HMX queue runs asynchronously
+ hmx_fa_qk_job_t qk_job[2];
+ hmx_fa_o_update_job_t ou_job[2];
+
+ // Prefetch block 1 early if there are multiple blocks
+ if (factx.n_kv_blocks > 1) {
+ fa_prefetch_block(dma, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+ ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx);
+ }
+
+ // Prep and start QK-dot(0)
+ void * curr_k0 = dma_queue_pop(dma).dst;
+ fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0);
+
+ qk_job[0].q_tiles = factx.vtcm_q_tiles;
+ qk_job[0].k_tiles = factx.vtcm_k_tiles[0];
+ qk_job[0].s_tiles = factx.vtcm_s_tiles[0];
+ qk_job[0].n_row_tiles = n_row_tiles;
+ qk_job[0].n_col_tiles = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
+ qk_job[0].n_dot_tiles = DK / 32;
+ qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
+ qk_job[0].hmx_scales = factx.vtcm_hmx_scales_qk;
+ hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
+
for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
const uint32_t kv_start = kv_blk * Bc;
const uint32_t kv_rows = hex_smin(Bc, nek1 - kv_start);
const size_t n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
- // Push mask DMA
- if (mask) {
- if (__builtin_expect(factx.mask_broadcast, true)) {
- const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
- dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
- } else {
- fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
- }
- }
-
- // Prefetch next KV block early
- if (kv_blk + 1 < factx.n_kv_blocks) {
- const uint32_t prefetch_start = (kv_blk + 1) * Bc;
- const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
- const size_t prefetch_buf = 1 - buf_idx;
- const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
- const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
- }
-
- // ---- Phase 1: K_int ----
- if (kv_blk > 0) {
- ou_job.o_curr = o_tile_curr;
- ou_job.o_prev = o_tile_prev;
- ou_job.p_tiles = factx.vtcm_p_tiles;
- ou_job.v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
- ou_job.d_tiles = factx.vtcm_d_tiles;
- ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
- ou_job.n_row_tiles = n_row_tiles;
- ou_job.n_col_tiles = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
- ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
- ou_job.n_tiles_per_bc = n_tiles_per_bc;
- ou_job.DV = DV;
- hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
- }
-
- // Wait for current K DMA and interleave
- void * curr_k = dma_queue_pop(dma).dst;
- fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
-
- // ---- Phase 2: qk_dot ----
- qk_job.q_tiles = factx.vtcm_q_tiles;
- qk_job.k_tiles = factx.vtcm_k_tiles;
- qk_job.s_tiles = factx.vtcm_s_tiles;
- qk_job.n_row_tiles = n_row_tiles;
- qk_job.n_col_tiles = n_col_tiles;
- qk_job.n_dot_tiles = DK / 32;
- qk_job.n_tiles_per_bc = n_tiles_per_bc;
- qk_job.hmx_scales = factx.vtcm_hmx_scales_qk;
- hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
-
- // Wait for current V DMA and interleave
+ // ---- 1. Pop and run V-prep for current block ----
void * curr_v = dma_queue_pop(dma).dst;
fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);
- if (kv_blk > 0) {
- hmx_queue_pop(hmx_q);
- hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
- }
-
- hmx_queue_pop(hmx_q);
-
- // ---- Phase 3: softmax + build_D ----
+ // ---- 2. Pop and run mask-prep for current block ----
__fp16 * current_mask_vtcm = NULL;
if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
}
}
+ // ---- 3. Pop and run K-prep for next block & push next QK-dot ----
+ if (kv_blk + 1 < factx.n_kv_blocks) {
+ const uint32_t next_start = (kv_blk + 1) * Bc;
+ const uint32_t next_rows = hex_smin(Bc, nek1 - next_start);
+ const size_t next_buf = 1 - buf_idx;
+
+ void * next_k = dma_queue_pop(dma).dst;
+ fa_phase_k_interleave(&factx, next_rows, k_src_stride, next_k, next_start, next_buf);
+
+ qk_job[next_buf].q_tiles = factx.vtcm_q_tiles;
+ qk_job[next_buf].k_tiles = factx.vtcm_k_tiles[next_buf];
+ qk_job[next_buf].s_tiles = factx.vtcm_s_tiles[next_buf];
+ qk_job[next_buf].n_row_tiles = n_row_tiles;
+ qk_job[next_buf].n_col_tiles = hmx_ceil_div(next_rows, HMX_FP16_TILE_N_COLS);
+ qk_job[next_buf].n_dot_tiles = DK / 32;
+ qk_job[next_buf].n_tiles_per_bc = n_tiles_per_bc;
+ qk_job[next_buf].hmx_scales = factx.vtcm_hmx_scales_qk;
+ hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
+ }
+
+ // ---- 4. Wait for current block's QK-dot to finish ----
+ hmx_queue_pop(hmx_q);
+
+ // ---- 5. Phase 2: softmax + build_D ----
fa_softmax_args_t sargs;
memset(&sargs, 0, sizeof(sargs));
sargs.factx = &factx;
+ sargs.buf_idx = buf_idx;
sargs.kv_rows = kv_rows;
sargs.n_rows_g = n_rows_g;
sargs.n_col_tiles = n_col_tiles;
sargs.mask_vtcm = current_mask_vtcm;
sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
sargs.slopes = factx.vtcm_slopes;
+
+ // Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx])
+ if (kv_blk > 0) {
+ const size_t prev_buf = 1 - buf_idx;
+ ou_job[prev_buf].o_curr = o_tile_curr;
+ ou_job[prev_buf].o_prev = o_tile_prev;
+ ou_job[prev_buf].p_tiles = factx.vtcm_p_tiles[prev_buf];
+ ou_job[prev_buf].v_tiles = factx.vtcm_v_tiles[prev_buf];
+ ou_job[prev_buf].d_tiles = factx.vtcm_d_tiles;
+ ou_job[prev_buf].hmx_scales = factx.vtcm_hmx_scales_id;
+ ou_job[prev_buf].n_row_tiles = n_row_tiles;
+ ou_job[prev_buf].n_col_tiles = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
+ ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
+ ou_job[prev_buf].n_tiles_per_bc = n_tiles_per_bc;
+ ou_job[prev_buf].DV = DV;
+ hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
+ }
+
+ // Run Softmax on HVX (blocking call)
fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
+ // Wait for HMX O update for block kv_blk - 1 to finish
+ if (kv_blk > 0) {
+ hmx_queue_pop(hmx_q);
+ hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
+ }
+
+ // Prefetch block kv_blk + 2
+ if (kv_blk + 2 < factx.n_kv_blocks) {
+ fa_prefetch_block(dma, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+ ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx);
+ }
+
buf_idx = 1 - buf_idx;
}
if (factx.n_kv_blocks > 0) {
const uint32_t last_blk = factx.n_kv_blocks - 1;
const size_t last_cols = hmx_ceil_div(hex_smin(Bc, nek1 - last_blk * Bc), HMX_FP16_TILE_N_COLS);
- ou_job.o_curr = o_tile_curr;
- ou_job.o_prev = o_tile_prev;
- ou_job.p_tiles = factx.vtcm_p_tiles;
- ou_job.v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
- ou_job.d_tiles = factx.vtcm_d_tiles;
- ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
- ou_job.n_row_tiles = n_row_tiles;
- ou_job.n_col_tiles = last_cols;
- ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
- ou_job.n_tiles_per_bc = n_tiles_per_bc;
- ou_job.DV = DV;
- hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
+ ou_job[0].o_curr = o_tile_curr;
+ ou_job[0].o_prev = o_tile_prev;
+ ou_job[0].p_tiles = factx.vtcm_p_tiles[1 - buf_idx];
+ ou_job[0].v_tiles = factx.vtcm_v_tiles[1 - buf_idx];
+ ou_job[0].d_tiles = factx.vtcm_d_tiles;
+ ou_job[0].hmx_scales = factx.vtcm_hmx_scales_id;
+ ou_job[0].n_row_tiles = n_row_tiles;
+ ou_job[0].n_col_tiles = last_cols;
+ ou_job[0].n_row_tiles_g_br = n_row_tiles_g_br;
+ ou_job[0].n_tiles_per_bc = n_tiles_per_bc;
+ ou_job[0].DV = DV;
+ hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[0]));
+
+ // Overlapped: run HVX build diag inv L while HMX is busy executing the update
+ htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
+ fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
+ htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
hmx_queue_pop(hmx_q);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
// Wait for current K DMA and interleave
void * curr_k = dma_queue_pop(dma).dst;
- fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
+ fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start, 0);
{
qk_job.q_tiles = factx.vtcm_q_tiles;
- qk_job.k_tiles = factx.vtcm_k_tiles;
- qk_job.s_tiles = factx.vtcm_s_tiles;
+ qk_job.k_tiles = factx.vtcm_k_tiles[0];
+ qk_job.s_tiles = factx.vtcm_s_tiles[0];
qk_job.n_row_tiles = n_row_tiles;
qk_job.n_col_tiles = n_col_tiles;
qk_job.n_dot_tiles = (size_t) (DK / 32);
{
ou_job.o_curr = o_tile_curr;
ou_job.o_prev = o_tile_prev;
- ou_job.p_tiles = factx.vtcm_p_tiles;
+ ou_job.p_tiles = factx.vtcm_p_tiles[0];
ou_job.v_tiles = factx.vtcm_v_tiles[0];
ou_job.d_tiles = factx.vtcm_d_tiles;
ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
ou_job.DV = DV;
hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
+ if (kv_blk + 1 == factx.n_kv_blocks) {
+ // Overlapped: run HVX build diag inv L while HMX is busy executing the update
+ htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
+ fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
+ htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
+ }
hmx_queue_pop(ctx->hmx_queue);
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
}
}
+ // Enqueue DMAs for the next iteration early so they overlap with O-PROC
+ uint32_t next_kv_head = kv_head + 1;
+ uint32_t next_q_start = q_start;
+ uint32_t next_ib3 = ib3;
+ if (next_kv_head >= n_kv_heads) {
+ next_kv_head = 0;
+ next_q_start = q_start + Br;
+ if (next_q_start >= neq1) {
+ next_q_start = 0;
+ next_ib3 = ib3 + 1;
+ }
+ }
+ bool has_next = (next_ib3 < neq3);
+
+ if (has_next) {
+ const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start);
+ const uint8_t * next_q_ptr = (const uint8_t *) q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
+ const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
+ const size_t next_n_rows = q_transposed ? factx.G : next_n_rows_q;
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
+
+ if (factx.n_kv_blocks > 0) {
+ const uint32_t next_ik2 = next_kv_head;
+ const uint32_t next_iv2 = next_kv_head;
+ uint32_t next_ik3 = ik3;
+ uint32_t next_iv3 = iv3;
+ if (next_ib3 != ib3) {
+ next_ik3 = fastdiv(next_ib3, &kparams->broadcast_rk3);
+ next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3);
+ }
+
+ const uint8_t * next_k_src = (const uint8_t *) k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+
+ const uint8_t * next_v_src = (const uint8_t *) v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+
+ if (factx.pipeline && mask) {
+ uint32_t next_im3 = im3;
+ if (next_ib3 != ib3) {
+ next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3);
+ }
+ if (__builtin_expect(factx.mask_broadcast, true)) {
+ const uint8_t * ms_src = (const uint8_t *) mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
+ dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
+ } else {
+ fa_push_mask_dma_gqa(dma, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
+ }
+ }
+ }
+ }
+
// ---- Final normalization ----
{
- htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
- fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
- htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
-
on_job.o_curr = o_tile_curr;
on_job.o_prev = o_tile_prev;
- on_job.d_tiles = factx.vtcm_d_tiles;
+ on_job.d_tiles = factx.vtcm_d_inv_l;
on_job.hmx_scales = factx.vtcm_hmx_scales_id;
on_job.n_row_tiles = n_row_tiles;
on_job.n_row_tiles_g_br = n_row_tiles_g_br;
}
if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
- worker_pool_run_func(octx->ctx->worker_pool, flash_attn_ext_f16_thread, &factx, octx->n_threads);
+ work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, octx->n_threads);
}
return HTP_STATUS_OK;
struct hmx_fa_vtcm_layout {
// Byte offsets from vtcm_base for each region.
size_t off_q_tiles;
+ size_t off_q_dma;
size_t off_o_tiles[2];
size_t off_k_fp16[2];
size_t off_v_fp16[2];
- size_t off_k_tiles;
- size_t off_v_tiles[2]; // [1] allocated only when pipeline, else 0
- size_t off_s_tiles;
- size_t off_p_tiles;
+ size_t off_k_tiles[2];
+ size_t off_v_tiles[2];
+ size_t off_s_tiles[2];
+ size_t off_p_tiles[2];
size_t off_d_tiles;
+ size_t off_d_inv_l;
size_t off_m_vec;
size_t off_l_vec;
size_t off_s_rowmax;
static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
size_t gqa_factor, size_t DK, size_t DV,
- size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+ size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+ const size_t q_dma_size = hex_align_up(g_br * DK * (is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128);
const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256);
size_t off = 0;
- // Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned)
+ // Group A (Part 1 - HMX Tiled buffers)
VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size);
VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size);
VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size);
- VTCM_LAYOUT_ALLOC(off, off_k_tiles, k_tile_size);
- VTCM_LAYOUT_ALLOC(off, off_v_tiles[0], v_tile_size);
- VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline);
- VTCM_LAYOUT_ALLOC(off, off_s_tiles, s_tile_size);
- VTCM_LAYOUT_ALLOC(off, off_p_tiles, s_tile_size);
VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_d_inv_l, d_tile_size);
- // Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned)
+ // Group B & C share start offset (Group B tiles must be 2KB aligned)
+ size_t off_group_b_c = hex_align_up(off, HTP_FA_HMX_TILE_SIZE);
+
+ // Group B: Compute-only buffers
+ size_t off_group_b = off_group_b_c;
+ VTCM_LAYOUT_ALLOC(off_group_b, off_k_tiles[0], k_tile_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_k_tiles[1], k_tile_size, pipeline);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_v_tiles[0], v_tile_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_v_tiles[1], v_tile_size, pipeline);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_s_tiles[0], s_tile_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_s_tiles[1], s_tile_size, pipeline);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_p_tiles[0], s_tile_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_p_tiles[1], s_tile_size, pipeline);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_s_rowmax, col_vec_size);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_p_rowsum, col_vec_size);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_row_bufs, row_vec_size * 2 * n_threads);
+
+ const size_t group_b_size = off_group_b - off_group_b_c;
+
+ // Group C: Q fetch DMA buffer
+ size_t off_group_c = off_group_b_c;
+ VTCM_LAYOUT_ALLOC(off_group_c, off_q_dma, q_dma_size);
+
+ const size_t group_c_size = off_group_c - off_group_b_c;
+
+ off = off_group_b_c + hex_smax(group_b_size, group_c_size);
+
+ // Group A (Part 2 - remaining non-HMX buffers)
VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size);
VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size);
VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size);
VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size);
VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size);
VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size);
- VTCM_LAYOUT_ALLOC(off, off_s_rowmax, col_vec_size);
- VTCM_LAYOUT_ALLOC(off, off_p_rowsum, col_vec_size);
- VTCM_LAYOUT_ALLOC(off, off_row_bufs, row_vec_size * 2 * n_threads);
VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256);
VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size);
}
// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
struct hmx_fa_vtcm_layout L;
- hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline);
+ hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32);
return L.total_bytes;
}
size_t qo_len,
size_t kv_len,
size_t vtcm_budget,
- size_t n_threads) {
+ size_t n_threads,
+ bool is_q_fp32) {
const size_t T = HMX_FP16_TILE_N_ROWS; // 32
const size_t br_unit = hmx_ceil_div(T, gqa_factor);
const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64
const size_t Bc_limit = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
(kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
// Cost coefficients calibrated from profiling
- const size_t c_q_fixed = 1400; // per-Q-block: q_load + epilogue o_update + o_norm + o_store
- const size_t c_iter_fixed = 200; // per-KV-iter: HMX queue push/pop + DMA pop + barriers
+ const size_t c_q_fixed = 800; // per-Q-block: q_load + epilogue o_update + o_norm + o_store
+ const size_t c_iter_base = 200; // per-KV-iter base (HMX dot/update + DMA)
+ const size_t c_softmax = 600; // per 64-row vector chunk on HVX
size_t best_cost = SIZE_MAX, best_mn = 0;
size_t best_Br = 0, best_Bc = 0;
for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
// Try all Bc candidates from Bc_limit down to bc_unit
for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
- size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline);
+ size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32);
if (vtcm_needed <= vtcm_budget) {
// This Bc fits for this Br!
- const size_t q_blocks = (qo_len + Br - 1) / Br;
- const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
- const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
- const size_t mn = Br * Bc;
+ const size_t q_blocks = (qo_len + Br - 1) / Br;
+ const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
+ const size_t actual_threads = (kv_blocks >= 3 && n_threads >= 2) ? n_threads : 1;
+ const size_t n_rows_g = Br * gqa_factor;
+ const size_t n_row_vec_cnt = (n_rows_g + 63) / 64;
+ const size_t n_use = n_row_vec_cnt < actual_threads ? n_row_vec_cnt : actual_threads;
+ const size_t vecs_per_t = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
+
+ const size_t c_iter_actual = c_iter_base + c_softmax * vecs_per_t;
+ const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
+ const size_t mn = Br * Bc;
if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
best_cost = cost;
--- /dev/null
+#ifndef HEX_BITMAP_H
+#define HEX_BITMAP_H
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <string.h>
+
+static inline void bitmap_set(uint32_t * bitmap, uint32_t idx) {
+ bitmap[idx / 32] |= (1U << (idx % 32));
+}
+
+static inline void bitmap_clear(uint32_t * bitmap, uint32_t idx) {
+ bitmap[idx / 32] &= ~(1U << (idx % 32));
+}
+
+static inline bool bitmap_test(const uint32_t * bitmap, uint32_t idx) {
+ return (bitmap[idx / 32] & (1U << (idx % 32))) != 0;
+}
+
+static inline void bitmap_reset(uint32_t * bitmap, size_t size_in_bits) {
+ memset(bitmap, 0, ((size_in_bits + 31) / 32) * sizeof(uint32_t));
+}
+
+#endif // HEX_BITMAP_H
+++ /dev/null
-#include "hex-dma.h"
-
-#include <stdbool.h>
-#include <stdlib.h>
-#include <string.h>
-
-#pragma clang diagnostic ignored "-Wunused-function"
-
-static inline uint32_t pow2_ceil(uint32_t x) {
- if (x <= 1) {
- return 1;
- }
- int p = 2;
- x--;
- while (x >>= 1) {
- p <<= 1;
- }
- return p;
-}
-
-dma_queue * dma_queue_create(size_t capacity) {
- dma_queue * q = (dma_queue *) memalign(32, sizeof(dma_queue));
- if (q == NULL) {
- FARF(ERROR, "%s: failed to allocate DMA queue\n", __FUNCTION__);
- return NULL;
- }
-
- capacity = pow2_ceil(capacity);
-
- memset(q, 0, sizeof(dma_queue));
- q->capacity = capacity;
- q->idx_mask = capacity - 1;
-
- q->desc = (dma_descriptor_2d *) memalign(64, capacity * sizeof(dma_descriptor_2d));
- memset(q->desc, 0, capacity * sizeof(dma_descriptor_2d));
-
- q->dptr = (dma_ptr *) memalign(4, capacity * sizeof(dma_ptr));
- memset(q->dptr, 0, capacity * sizeof(dma_ptr));
-
- q->tail = &q->desc[capacity - 1];
-
- if (!q->desc && !q->dptr) {
- FARF(ERROR, "%s: failed to allocate DMA queue items\n", __FUNCTION__);
- return NULL;
- }
-
- FARF(HIGH, "dma-queue: capacity %u\n", capacity);
-
- return q;
-}
-
-void dma_queue_delete(dma_queue * q) {
- if (!q) {
- return;
- }
- free(q->desc);
- free(q->dptr);
- free(q);
-}
-
-void dma_queue_flush(dma_queue * q) {
- while (dma_queue_pop(q).dst != NULL) ;
-}
-#ifndef HTP_DMA_H
-#define HTP_DMA_H
-
-#include <HAP_farf.h>
-#include <hexagon_types.h>
-#include <stdbool.h>
-#include <stdint.h>
-#include "hex-utils.h"
-
-#include "hex-profile.h"
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-// Define the HW descriptor structs here since the ones in HexSDK are a bit out of date
-typedef struct dma_descriptor_1d_s {
- void * next;
- uint32_t size:24;
- uint32_t desc_size:2;
- uint32_t dst_comp:1;
- uint32_t src_comp:1;
- uint32_t dst_bypass:1;
- uint32_t src_bypass:1;
- uint32_t order:1;
- uint32_t done:1;
- void * src;
- void * dst;
-} dma_descriptor_1d;
-
-#if __HVX_ARCH__ < 75
-
-typedef struct dma_descriptor_2d_s {
- void * next;
- uint32_t reserved0:24;
- uint32_t desc_size:2;
- uint32_t dst_comp:1;
- uint32_t src_comp:1;
- uint32_t dst_bypass:1;
- uint32_t src_bypass:1;
- uint32_t order:1;
- uint32_t done:1;
- void * src;
- void * dst;
- uint32_t desc_type:8;
- uint32_t reserved1:24;
- uint32_t row_size:16;
- uint32_t nrows:16;
- uint32_t src_stride:16;
- uint32_t dst_stride:16;
- uint32_t src_offset:16;
- uint32_t dst_offset:16;
-} dma_descriptor_2d;
-
-#else
-
-typedef struct dma_descriptor_2d_s {
- void * next;
- uint32_t dst_stride:24;
- uint32_t desc_size:2;
- uint32_t dst_comp:1;
- uint32_t src_comp:1;
- uint32_t dst_bypass:1;
- uint32_t src_bypass:1;
- uint32_t order:1;
- uint32_t done:1;
- void * src;
- void * dst;
- uint32_t desc_type:8;
- uint32_t reserved0:24;
- uint32_t row_size:24;
- uint32_t nrows_lo:8;
- uint32_t nrows_hi:8;
- uint32_t src_stride:24;
- uint32_t offset:24;
- uint32_t reserved1:8;
-} dma_descriptor_2d;
-
-#endif
-
-typedef struct {
- void *dst;
- const void *src;
-} dma_ptr;
-
-typedef struct {
- dma_descriptor_2d * desc; // descriptor pointers
- dma_descriptor_2d * tail; // tail pointer
- dma_ptr * dptr; // dst/src pointers
- uint32_t push_idx;
- uint32_t pop_idx;
- uint32_t capacity;
- uint32_t idx_mask;
- struct htp_thread_trace * trace;
-} dma_queue;
-
-dma_queue * dma_queue_create(size_t capacity);
-void dma_queue_delete(dma_queue * q);
-void dma_queue_flush(dma_queue * q);
-
-// TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
-// but those do not seem to always compiler properly.
-static inline void dmstart(void * next) {
- asm volatile(" release(%0):at" : : "r"(next));
- asm volatile(" dmstart(%0)" : : "r"(next));
-}
-
-static inline void dmlink(void * cur, void * next) {
- asm volatile(" release(%0):at" : : "r"(next));
- asm volatile(" dmlink(%0, %1)" : : "r"(cur), "r"(next));
-}
-
-static inline unsigned int dmpoll(void) {
- unsigned int ret = 0;
- asm volatile(" %0 = dmpoll" : "=r"(ret) : : "memory");
- return ret;
-}
-
-static inline unsigned int dmwait(void) {
- unsigned int ret = 0;
- asm volatile(" %0 = dmwait" : "=r"(ret) : : "memory");
- return ret;
-}
-
-static inline dma_ptr dma_make_ptr(void *dst, const void *src)
-{
- dma_ptr p = { dst, src };
- return p;
-}
-
-static const uint32_t dma_src_l2_bypass_on = 1;
-static const uint32_t dma_dst_l2_bypass_on = 1;
-
-static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
- if (((q->push_idx + 1) & q->idx_mask) == q->pop_idx) {
- FARF(HIGH, "dma-push: queue full\n");
- return false;
- }
-
- dma_descriptor_1d * desc = (dma_descriptor_1d *) &q->desc[q->push_idx];
- desc->src = (void *) dptr.src;
- desc->dst = (void *) dptr.dst;
- desc->size = size;
-
- q->dptr[q->push_idx] = dptr;
-
- if (size) {
- desc->next = NULL;
- desc->desc_size = 0; // 1D mode
- desc->src_bypass = dma_src_l2_bypass_on;
- desc->dst_bypass = dma_dst_l2_bypass_on;
- desc->order = 0;
- desc->done = 0;
-
- htp_trace_event_start(q->trace, HTP_TRACE_EVT_DMA, q->push_idx);
- dmlink(q->tail, desc);
- q->tail = (dma_descriptor_2d *) desc;
- } else {
- desc->desc_size = 0;
- desc->done = 1;
- }
-
- q->push_idx = (q->push_idx + 1) & q->idx_mask;
- return true;
-}
-
-static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
- if (((q->push_idx + 1) & q->idx_mask) == q->pop_idx) {
- FARF(HIGH, "dma-push: queue full\n");
- return false;
- }
-
- dma_descriptor_2d * desc = &q->desc[q->push_idx];
-
- desc->next = NULL;
- desc->reserved0 = 0;
- desc->reserved1 = 0;
- desc->desc_size = 1; // 2d mode
- desc->src_bypass = dma_src_l2_bypass_on;
- desc->dst_bypass = dma_dst_l2_bypass_on;
- desc->src_comp = 0;
- desc->dst_comp = 0;
- desc->order = 0;
- desc->done = 0;
- desc->src_stride = src_stride;
- desc->dst_stride = dst_stride;
- desc->src = (void *) dptr.src;
- desc->dst = (void *) dptr.dst;
- desc->row_size = row_size;
-
-#if __HVX_ARCH__ < 75
- desc->desc_type = 0; // 2d (16-bit) mode
- desc->nrows = nrows;
- desc->src_offset = 0;
- desc->dst_offset = 0;
-#else
- desc->desc_type = 9; // 2d (24-bit) mode
- desc->nrows_lo = (nrows & 0xff);
- desc->nrows_hi = (nrows >> 8);
- desc->offset = 0;
-#endif
-
- q->dptr[q->push_idx] = dptr;
-
- if (nrows) {
- htp_trace_event_start(q->trace, HTP_TRACE_EVT_DMA, q->push_idx);
- dmlink(q->tail, desc);
- q->tail = desc;
- } else {
- desc->done = 1;
- }
-
- // FARF(ERROR, "dma-push: i %u row-size %u nrows %d dst %p src %p\n", q->push_idx, row_size, nrows, dptr.dst, dptr.src);
- q->push_idx = (q->push_idx + 1) & q->idx_mask;
- return true;
-}
-
-static inline dma_ptr dma_queue_pop(dma_queue * q) {
- dma_ptr dptr = { NULL };
-
- if (q->push_idx == q->pop_idx) {
- return dptr;
- }
-
- dma_descriptor_2d * desc = &q->desc[q->pop_idx];
-
- // Wait for desc to complete
- if (!desc->done) {
- while (!desc->done) {
- dmpoll();
- }
- }
- htp_trace_event_stop(q->trace, HTP_TRACE_EVT_DMA, q->pop_idx);
-
- dptr = q->dptr[q->pop_idx];
-
- // FARF(ERROR, "dma-pop: i %u dst %p src %p\n", q->pop_idx, dptr.dst, dptr.src);
- q->pop_idx = (q->pop_idx + 1) & q->idx_mask;
- return dptr;
-}
-
-static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
- dma_ptr dptr = { NULL };
-
- if (q->push_idx == q->pop_idx) {
- return dptr;
- }
-
- dptr = q->dptr[q->pop_idx];
-
- // FARF(ERROR, "dma-pop-nowait: i %u dst %p src %p\n", q->pop_idx, dptr.dst, dptr.src);
- q->pop_idx = (q->pop_idx + 1) & q->idx_mask;
- return dptr;
-}
-
-static inline bool dma_queue_empty(dma_queue * q) {
- return q->push_idx == q->pop_idx;
-}
-
-static inline uint32_t dma_queue_depth(dma_queue * q) {
- return (q->push_idx - q->pop_idx) & q->idx_mask;
-}
-
-static inline uint32_t dma_queue_capacity(dma_queue * q) {
- return q->capacity;
-}
-
-#if __HVX_ARCH__ < 75
-
-// Overflow-safe DMA push: all 2d descriptor fields (row_size, nrows, src_stride, dst_stride) are 16-bit, max 65535.
-// This version transparently handles values that exceed the 16-bit limit and submits chained DMA transtions.
-
-#define DMA_MAX_FIELD_VAL 65535u
-
-static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
- // Fast path: everything fits in 16 bits
- if (nrows == 0 || __builtin_expect(
- row_size <= DMA_MAX_FIELD_VAL &&
- nrows <= DMA_MAX_FIELD_VAL &&
- src_stride <= DMA_MAX_FIELD_VAL &&
- dst_stride <= DMA_MAX_FIELD_VAL, 1)) {
- return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
- }
-
- // Contiguous block
- // Use 1d DMA mode which supports sizes up to 24-bits (16MB)
- if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
- size_t total = row_size * nrows;
- return dma_queue_push_single_1d(q, dptr, total);
- }
-
- // Stride overflow — fall back to row-by-row.
- {
- const uint8_t *src = (const uint8_t *) dptr.src;
- uint8_t *dst = (uint8_t *) dptr.dst;
- for (size_t r = 0; r < nrows; ++r) {
- dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
- if (!dma_queue_push_single_1d(q, p, row_size))
- return false;
- if (r + 1 < nrows)
- dma_queue_pop(q);
- }
- return true;
- }
-}
-
-#else // HVX_ARCH >= 75
-
-static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
- // On v75 and up we always use 2d 24-bit mode
- return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
-}
-
-#endif
-
-static inline bool dma_queue_push_ddr_to_vtcm(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
- return dma_queue_push(q, dptr, dst_row_size, src_row_size, src_row_size, nrows);
-}
-
-static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
- return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
-}
-
-#define DMA_CACHE_MAX_SIZE 256U
-
-typedef struct {
- uint8_t *base;
- uint32_t line_size;
- uint32_t capacity;
- uint32_t src[DMA_CACHE_MAX_SIZE];
- uint16_t age[DMA_CACHE_MAX_SIZE];
-} dma_cache;
-
-static inline void dma_cache_init(dma_cache *c, uint8_t *base, uint32_t line_size, uint32_t capacity)
-{
- c->capacity = (capacity > DMA_CACHE_MAX_SIZE) ? DMA_CACHE_MAX_SIZE : capacity;
- c->base = base;
- c->line_size = line_size;
-
- for (unsigned i=0; i < c->capacity; i++) {
- c->src[i] = 0;
- c->age[i] = 0;
- }
-}
-
-static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * src, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows)
-{
- uint32_t o_idx = 0;
- uint16_t o_age = 0;
- uint8_t * dst = 0;
-
- for (unsigned i=0; i < c->capacity; i++) {
- if (c->src[i] == (uint32_t) src) {
- c->age[i] = 0;
- dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
- } else {
- c->age[i]++;
- if (c->age[i] > o_age) { o_age = c->age[i]; o_idx = i; }
- }
- }
- if (!dst) {
- c->age[o_idx] = 0;
- c->src[o_idx] = (uint32_t) src;
- dst = c->base + o_idx * c->line_size; // normal nrows dma
- return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
- }
-
- return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0);
-}
-
-#ifdef __cplusplus
-} // extern "C"
-#endif
-
-#endif /* HTP_DMA_H */
+#pragma once
+#include "dma-queue.h"
};
static inline void htp_trace_event(struct htp_thread_trace * tr, uint16_t id, uint16_t info, uint32_t type) {
- if (tr && tr->events && tr->count < tr->max_events) {
- uint32_t idx = tr->count;
- tr->events[idx].id = id;
- tr->events[idx].info = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0);
- tr->events[idx].cycles = (uint32_t) hex_get_cycles();
+ if (tr->count < tr->max_events) {
+ uint32_t i = tr->count;
+ tr->events[i].id = id;
+ tr->events[i].info = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0);
+ tr->events[i].cycles = (uint32_t) hex_get_cycles();
tr->count++;
}
}
Q6_l2fetch_AP((void *) p, control);
}
-#define HEX_L2_LINE_SIZE 64
-#define HEX_L2_FLUSH_SIZE (128 * 1024)
+static inline void hex_l2fetch_block(const void * addr, size_t size) {
+ if (size == 0) return;
+ const uint32_t width = 16384; // 16KB rows
+ const uint32_t height = (size + width - 1) / width;
+ hex_l2fetch(addr, width, width, height);
+}
+
+#define HEX_L2_LINE_SIZE 128
+#define HEX_L2_BLOCK_SIZE (HEX_L2_LINE_SIZE * 4) // flush granularity (lines per loop iteration)
+#define HEX_L2_FLUSH_WQ_THRESHOLD (4 * 1024)
+#define HEX_L2_FLUSH_ALL_THRESHOLD (4 * 1024 * 1024)
static inline void hex_l2flush(void * addr, size_t size) {
- if (size > HEX_L2_FLUSH_SIZE) {
- qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
- } else {
- const uint32_t s = (uint32_t) addr;
- const uint32_t e = s + size;
- for (uint32_t i = s; i < e; i += HEX_L2_LINE_SIZE * 4) {
- Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0);
- Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1);
- Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2);
- Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3);
- }
+ const uint32_t s = ((uint32_t) addr) & ~(HEX_L2_LINE_SIZE - 1);
+ const uint32_t e = (((uint32_t) addr) + size + HEX_L2_LINE_SIZE - 1) & ~(HEX_L2_LINE_SIZE - 1);
+ for (uint32_t i = s; i < e; i += HEX_L2_BLOCK_SIZE) {
+ Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0);
+ Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1);
+ Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2);
+ Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3);
}
}
// output : fp16 -> f32p
-static void transfer_output_chunk_fp16_to_fp32(
+static void transfer_output_chunk_fp16_to_fp32_col_chunk(
float *restrict dst,
const float *restrict src2,
const __fp16 *restrict vtcm_src,
uint32_t start_row,
uint32_t n_rows,
- uint32_t n_cols,
+ uint32_t c_len,
+ uint32_t total_n_cols,
uint32_t dst_stride,
uint32_t src2_stride,
uint32_t dst_cols
) {
- assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
- const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
+ assert(c_len % HTP_MM_HMX_TILE_N_COLS == 0);
+ assert(total_n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
+ const size_t tile_row_stride = (total_n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
const HVX_Vector one = hvx_vec_splat_f16(1.0);
- const size_t limit_c = hex_smin(n_cols, dst_cols);
+ const size_t limit_c = hex_smin(c_len, dst_cols);
const size_t limit_c_aligned = (limit_c & ~31);
for (size_t r = 0; r < n_rows; r += 2) {
}
}
+static inline void transfer_output_chunk_fp16_to_fp32(
+ float *restrict dst,
+ const float *restrict src2,
+ const __fp16 *restrict vtcm_src,
+ uint32_t start_row,
+ uint32_t n_rows,
+ uint32_t n_cols,
+ uint32_t dst_stride,
+ uint32_t src2_stride,
+ uint32_t dst_cols
+) {
+ transfer_output_chunk_fp16_to_fp32_col_chunk(
+ dst, src2, vtcm_src, start_row, n_rows, n_cols, n_cols, dst_stride, src2_stride, dst_cols
+ );
+}
+
typedef struct {
const __fp16 *vtcm_src;
float *dst;
}
}
+static void transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+ __fp16 *restrict vtcm_dst,
+ const float *restrict row0, // offset by c_first
+ const float *restrict row1, // offset by c_first
+ uint32_t r,
+ uint32_t k_block,
+ uint32_t c_first,
+ uint32_t c_len,
+ uint32_t k_chunk_valid,
+ bool row0_valid,
+ bool row1_valid) {
+
+ uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
+ uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
+
+ uint32_t c = 0;
+ for (; c + 32 <= k_chunk_valid; c += 32) {
+ HVX_Vector v0 = Q6_V_vzero();
+ HVX_Vector v1 = Q6_V_vzero();
+ if (row0_valid) v0 = *(const HVX_Vector *)(row0 + c);
+ if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c);
+
+ HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+ uint32_t c0 = (c_first + c) / HTP_MM_HMX_TILE_N_COLS; // tile column index
+ uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
+
+ HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
+ tile[r1 / 2] = v_out;
+ }
+ if (c < c_len) {
+ HVX_Vector v0 = Q6_V_vzero();
+ HVX_Vector v1 = Q6_V_vzero();
+ if (row0_valid) v0 = *(const HVX_Vector *)(row0 + c);
+ if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c);
+
+ uint32_t rem = (k_chunk_valid > c) ? (k_chunk_valid - c) : 0;
+ HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0);
+ v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero());
+ v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero());
+
+ HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+ uint32_t c0 = (c_first + c) / HTP_MM_HMX_TILE_N_COLS; // tile column index
+ uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
+
+ HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
+ tile[r1 / 2] = v_out;
+ }
+}
+
static void transfer_activation_chunk_fp32_to_fp16_gathered(
__fp16 *restrict vtcm_dst,
const float *restrict src,
uint32_t start_row,
+ uint32_t vtcm_start_row,
uint32_t n_rows,
uint32_t k_block,
const struct mmid_row_mapping *matrix_rows,
for (r = 0; r < n_rows_tiled; r += 2) {
uint32_t r_idx0 = start_row + r + 0;
uint32_t r_idx1 = start_row + r + 1;
- uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index
- uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
+ uint32_t lr = vtcm_start_row + r; // vtcm-local row
+ uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS; // tile row index
+ uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0];
struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1];
}
for (; r < n_rows_padded; r += 2) {
- uint32_t r_idx0 = start_row + r;
- uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index
- uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
+ uint32_t lr = vtcm_start_row + r; // vtcm-local row
+ uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS; // tile row index
+ uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
const bool row0_valid = (start_row + r + 0) < cne1;
const bool row1_valid = (start_row + r + 1) < cne1;
__fp16 *restrict vtcm_dst,
const float *restrict src,
uint32_t start_row,
+ uint32_t vtcm_start_row,
uint32_t n_rows,
uint32_t k_block,
const struct mmid_row_mapping *matrix_rows,
for (r = 0; r < n_rows_tiled; r += 2) {
uint32_t r_idx0 = start_row + r + 0;
uint32_t r_idx1 = start_row + r + 1;
- uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index
- uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
+ uint32_t lr = vtcm_start_row + r; // vtcm-local row
+ uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS; // tile row index
+ uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0];
struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1];
}
for (; r < n_rows_padded; r += 2) {
- uint32_t r_idx0 = start_row + r;
- uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index
- uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
+ uint32_t lr = vtcm_start_row + r; // vtcm-local row
+ uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS; // tile row index
+ uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
const bool row0_valid = (start_row + r + 0) < cne1;
const bool row1_valid = (start_row + r + 1) < cne1;
float *restrict dst,
const __fp16 *restrict vtcm_src,
uint32_t start_row,
+ uint32_t vtcm_start_row,
uint32_t n_rows,
uint32_t n_cols,
const struct mmid_row_mapping *matrix_rows,
for (size_t r = 0; r < n_rows; r += 2) {
uint32_t r_idx0 = start_row + r + 0;
uint32_t r_idx1 = start_row + r + 1;
- const size_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS;
- const size_t r1 = (r_idx0 % HTP_MM_HMX_TILE_N_ROWS) / 2; // index of the row pair within the tile
+ uint32_t lr = vtcm_start_row + r; // vtcm-local row
+ const size_t r0 = (lr / HTP_MM_HMX_TILE_N_ROWS);
+ const size_t r1 = (lr % HTP_MM_HMX_TILE_N_ROWS) / 2; // index of the row pair within the tile
const __fp16 *row_base = vtcm_src + r0 * tile_row_stride;
if (r_idx0 >= cne1) break;
#define QURT_LOWEST_PRIO (254)
-static inline void hmx_lock(struct hmx_queue *q)
+static inline void hmx_lock(hmx_queue_t q)
{
if (!q->hmx_locked) {
HAP_compute_res_hmx_lock(q->hap_rctx);
}
}
-static inline void hmx_unlock(struct hmx_queue *q)
+static inline void hmx_unlock(hmx_queue_t q)
{
if (q->hmx_locked) {
HAP_compute_res_hmx_unlock(q->hap_rctx);
}
}
-static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) {
+static inline void hmx_queue_process(hmx_queue_t q, bool* killed) {
unsigned int ir = atomic_load(&q->idx_read);
while (ir != atomic_load(&q->idx_write)) {
if (!d->done) {
FARF(HIGH, "hmx-queue-process: ir %u func %p data %p", ir, d->func, d->data);
- enum hmx_queue_signal sig = (enum hmx_queue_signal) (unsigned int) d->func;
+ uintptr_t sig = (uintptr_t) d->func;
switch (sig) {
case HMX_QUEUE_NOOP: /* noop */; break;
case HMX_QUEUE_KILL: *killed = true; break;
}
static void hmx_queue_thread(void * arg) {
- struct hmx_queue * q = (struct hmx_queue *) arg;
+ hmx_queue_t q = (hmx_queue_t) arg;
FARF(HIGH, "hmx-queue-thread: started");
FARF(HIGH, "hmx-queue-thread: stopped");
}
-struct hmx_queue * hmx_queue_create(size_t capacity, uint32_t hap_rctx) {
+size_t hmx_queue_sizeof(size_t capacity, uint32_t stack_size) {
capacity = hex_ceil_pow2(capacity);
+ size_t size_q = hex_align_up(sizeof(struct hmx_queue_s), HEX_L2_LINE_SIZE);
+ size_t size_desc = hex_align_up(capacity * sizeof(struct hmx_queue_desc), HEX_L2_LINE_SIZE);
+ size_t size_stack = stack_size;
+ return size_q + size_desc + size_stack;
+}
+
+size_t hmx_queue_alignof(void) {
+ return HEX_L2_LINE_SIZE;
+}
+
+hmx_queue_t hmx_queue_init(void * ptr, size_t capacity, uint32_t stack_size, uint32_t hap_rctx, struct htp_thread_trace * trace) {
+ capacity = hex_ceil_pow2(capacity);
+ size_t size_q = hex_align_up(sizeof(struct hmx_queue_s), HEX_L2_LINE_SIZE);
+ size_t size_desc = hex_align_up(capacity * sizeof(struct hmx_queue_desc), HEX_L2_LINE_SIZE);
+
+ uint8_t * block = (uint8_t *) ptr;
+
+ hmx_queue_t q = (hmx_queue_t) block; block += size_q;
+ memset(q, 0, sizeof(struct hmx_queue_s));
- struct hmx_queue * q = (struct hmx_queue *) memalign(32, sizeof(struct hmx_queue));
- if (q == NULL) {
- FARF(ERROR, "%s: failed to allocate DMA queue\n", __FUNCTION__);
- return NULL;
- }
- memset(q, 0, sizeof(struct hmx_queue));
q->capacity = capacity;
q->idx_mask = capacity - 1;
q->hap_rctx = hap_rctx;
+ q->external_mem = true;
- q->desc = (struct hmx_queue_desc *) memalign(64, capacity * sizeof(struct hmx_queue_desc));
- if (!q->desc) {
- FARF(ERROR, "hmx-queue: failed to allocate HMX queue descriptors\n");
- return NULL;
- }
+ q->desc = (struct hmx_queue_desc *) block; block += size_desc;
memset(q->desc, 0, capacity * sizeof(struct hmx_queue_desc));
- const size_t stack_size = HMX_QUEUE_THREAD_STACK_SIZE;
- q->stack = (unsigned char *) memalign(64, stack_size);
- if (!q->stack) {
- FARF(ERROR, "hmx-queue: thread stack allocation failed (%zu bytes)", stack_size);
- return NULL;
- }
+ q->stack = block;
memset(q->stack, 0, stack_size);
+ q->trace = trace;
+
// Match caller thread priority (same pattern as worker-pool.c).
int prio = qurt_thread_get_priority(qurt_thread_get_id());
if (prio < 1) {
return q;
}
-void hmx_queue_delete(struct hmx_queue * q) {
+void hmx_queue_free(hmx_queue_t q) {
if (!q) {
return;
}
int status;
qurt_thread_join(q->thread, &status);
-
- free(q->desc);
- free(q->stack);
- free(q);
}
extern "C" {
#endif
-#define HMX_QUEUE_THREAD_STACK_SIZE (16 * 1024)
-
#if __HVX_ARCH__ > 79
#define HMX_QUEUE_POLL_COUNT 2000
#else
atomic_uint done;
};
-struct hmx_queue {
+struct hmx_queue_s {
struct hmx_queue_desc * desc;
atomic_uint idx_write; // updated by producer (push)
atomic_uint idx_read; // updated by consumer (process)
uint32_t hap_rctx;
bool hmx_locked;
struct htp_thread_trace * trace;
+ bool external_mem; // memory owned externally
};
-struct hmx_queue * hmx_queue_create(size_t capacity, uint32_t hap_rctx);
-void hmx_queue_delete(struct hmx_queue * q);
+typedef struct hmx_queue_s * hmx_queue_t;
+
+size_t hmx_queue_sizeof(size_t capacity, uint32_t stack_size);
+size_t hmx_queue_alignof(void);
+hmx_queue_t hmx_queue_init(void * ptr, size_t capacity, uint32_t stack_size, uint32_t hap_rctx, struct htp_thread_trace * trace);
+void hmx_queue_free(hmx_queue_t q);
static inline struct hmx_queue_desc hmx_queue_make_desc(hmx_queue_func func, void * data) {
struct hmx_queue_desc d = { func, data };
return d;
}
-static inline bool hmx_queue_push(struct hmx_queue * q, struct hmx_queue_desc d) {
+static inline bool hmx_queue_push(hmx_queue_t q, struct hmx_queue_desc d) {
unsigned int ir = atomic_load(&q->idx_read);
- unsigned int iw = q->idx_write;
+ unsigned int iw = atomic_load(&q->idx_write);
if (((iw + 1) & q->idx_mask) == ir) {
FARF(HIGH, "hmx-queue-push: queue is full\n");
return true;
}
-static inline bool hmx_queue_signal(struct hmx_queue *q, enum hmx_queue_signal sig) {
+static inline bool hmx_queue_signal(hmx_queue_t q, enum hmx_queue_signal sig) {
return hmx_queue_push(q, hmx_queue_make_desc((hmx_queue_func) sig, NULL));
}
-static inline bool hmx_queue_empty(struct hmx_queue * q) {
- return q->idx_pop == q->idx_write;
+static inline bool hmx_queue_empty(hmx_queue_t q) {
+ return q->idx_pop == atomic_load(&q->idx_write);
}
-static inline uint32_t hmx_queue_depth(struct hmx_queue * q) {
- return (q->idx_read - q->idx_read) & q->idx_mask;
+static inline uint32_t hmx_queue_depth(hmx_queue_t q) {
+ return (atomic_load(&q->idx_write) - atomic_load(&q->idx_read)) & q->idx_mask;
}
-static inline uint32_t hmx_queue_capacity(struct hmx_queue * q) {
+static inline uint32_t hmx_queue_capacity(hmx_queue_t q) {
return q->capacity;
}
-static inline struct hmx_queue_desc hmx_queue_pop_one(struct hmx_queue * q) {
+static inline struct hmx_queue_desc hmx_queue_pop_one(hmx_queue_t q) {
unsigned int ip = q->idx_pop;
- unsigned int iw = q->idx_write;
+ unsigned int iw = atomic_load(&q->idx_write);
struct hmx_queue_desc rd = { NULL, NULL };
if (ip == iw) {
return rd;
}
-static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
+static inline struct hmx_queue_desc hmx_queue_pop(hmx_queue_t q) {
while (1) {
struct hmx_queue_desc d = hmx_queue_pop_one(q);
}
}
-static inline void hmx_queue_flush(struct hmx_queue * q) {
+static inline void hmx_queue_flush(hmx_queue_t q) {
while (hmx_queue_pop_one(q).func != NULL) ;
}
-static inline void hmx_queue_wakeup(struct hmx_queue * q) {
+static inline void hmx_queue_wakeup(hmx_queue_t q) {
hmx_queue_signal(q, HMX_QUEUE_WAKEUP);
}
-static inline void hmx_queue_suspend(struct hmx_queue *q) {
+static inline void hmx_queue_suspend(hmx_queue_t q) {
hmx_queue_signal(q, HMX_QUEUE_SUSPEND);
}
#include "hmx-queue.h"
#include "htp-ops.h"
#include "hex-profile.h"
-#include "worker-pool.h"
+#include "work-queue.h"
+#include "hex-fastdiv.h"
#include <assert.h>
#include <dspqueue.h>
#endif
#define HTP_MAX_MMAPS 16
+#define HTP_MAX_DIRTY_RANGES 16
+
// Memory mapping
struct htp_mmap {
uint64_t size;
const struct htp_tensor * dsts[HTP_OP_MAX_OUTPUTS];
};
+ dma_queue ** src_dma[HTP_OP_MAX_INPUTS];
+ dma_queue ** dst_dma[HTP_OP_MAX_OUTPUTS];
+
// TODO convert these to an array
struct htp_spad src0_spad;
struct htp_spad src1_spad;
// Main context for htp DSP backend
struct htp_context {
- dspqueue_t queue;
- dma_queue * dma[HTP_MAX_NTHREADS];
+ dspqueue_t dsp_queue;
+
struct htp_mmap mmap[HTP_MAX_MMAPS];
- worker_pool_context_t worker_pool;
+ dma_queue_t dma[HTP_MAX_NTHREADS];
+ dma_queue_t dma_cached[HTP_MAX_NTHREADS];
+ work_queue_t work_queue;
+ hmx_queue_t hmx_queue;
+
uint32_t n_threads;
+ struct fastdiv_values n_threads_div;
int thread_id;
int thread_prio;
atomic_bool vtcm_needs_release;
uint64_t max_vmem;
+ struct htp_dirty_range {
+ uint32_t start;
+ uint32_t end;
+ uint32_t bi;
+ } dirty_ranges[HTP_MAX_DIRTY_RANGES];
// Persistent DDR scratchpad for MUL_MAT_ID mappings
void * ddr_spad_base;
struct htp_ops_context octx;
- struct hmx_queue * hmx_queue; // Async HMX queue for pipeline overlap
+ qurt_thread_t main_thread;
+ void * main_stack;
+ atomic_bool killed;
+ size_t footprint;
};
int op_matmul(struct htp_ops_context * octx);
int op_diag(struct htp_ops_context * octx);
int op_solve_tri(struct htp_ops_context * octx);
int op_gated_delta_net(struct htp_ops_context * octx);
-int op_tri(struct htp_ops_context * octx);
int op_pad(struct htp_ops_context * octx);
+int op_im2col(struct htp_ops_context * octx);
#endif /* HTP_CTX_H */
HTP_OP_PAD,
HTP_OP_NORM,
HTP_OP_CONCAT,
+ HTP_OP_CLAMP,
+ HTP_OP_IM2COL,
HTP_OP_INVALID
};
#define HTP_OP_MAX_KERN_PARAMS 32
#define HTP_OP_MAX_BUFS 16
-#define HTP_OP_MAX_REQS 256
-#define HTP_OP_MAX_TENSORS (HTP_OP_MAX_REQS * HTP_OP_MAX_INPUTS + HTP_OP_MAX_REQS)
+#define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)
#define HTP_OP_MAX_VMEM_DEFAULT (3355443200u)
enum htp_tensor_flags {
HTP_TENSOR_COMPUTE = (1U << 0), // Tensor buffer temporal compute data (not weights)
- HTP_TENSOR_FLUSHED = (1U << 1) // Tensor buffer has been flushed (set by the NPU)
+ HTP_TENSOR_DIRTY = (1U << 1) // Tensor buffer is dirty and needs to be flushed
};
// Tensor descriptor
struct htp_tensor {
uint32_t data; // Buffer offset in the messages, and data pointer on the NPU
+ uint32_t reserved; // Reserved for alignment padding (must be multiple of 8)
uint32_t size; // Data size in bytes
uint32_t flags; // Buffer / tensor flags
- uint16_t type; // Data type
+ uint32_t type; // Data type
uint16_t bi; // Buffer index
+ uint16_t ti; // Tensor index
uint32_t ne[HTP_OP_MAX_DIMS]; // Number of elements
uint32_t nb[HTP_OP_MAX_DIMS]; // Stride in bytes (see ggml.h ggml_tensor)
};
enum htp_trace_event_id {
HTP_TRACE_EVT_DMA = 0,
+ HTP_TRACE_EVT_L2FLUSH = 1,
+ HTP_TRACE_EVT_INIT = 2,
+ HTP_TRACE_EVT_BUFF = 3,
HTP_TRACE_EVT_HVX_COMP = 20,
HTP_TRACE_EVT_HVX_A_QUANT = 21,
uint32_t n_tensors; // Number of tensors
uint32_t n_ops; // Number of op profile descriptors
uint32_t n_traces[HTP_MAX_NTHREADS + 1];
- uint8_t pad[8]; // align to 8 bytes
+ uint32_t usecs; // Number of usec
+ uint32_t pad; // align to 8 bytes
+ uint64_t cycles_start; // Start cycle counter
+ uint64_t cycles_stop; // Stop cycle counter
// struct htp_prof_desc profs[]; -- dspqueue buf 0
};
--- /dev/null
+#include "htp-tensor.h"
+
+#include <qurt.h>
+#include <qurt_memory.h>
+#include <HAP_farf.h>
+
+#include "hex-common.h"
+#include "hex-utils.h"
+#include "hex-fastdiv.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "work-queue.h"
+
+struct l2flush_range {
+ uint32_t start; // line-aligned start address
+ uint32_t end; // line-aligned end address
+ uint32_t block_first; // global block index of this range's first block
+ uint32_t n_blocks; // number of HEX_L2_BLOCK_SIZE chunks (last may be partial)
+};
+
+struct l2flush_multi_task {
+ struct htp_thread_trace * trace;
+ struct l2flush_range ranges[HTP_OP_MAX_INPUTS];
+ uint32_t n_ranges;
+ uint32_t total_blocks;
+ uint32_t blocks_per_thread;
+};
+
+static void flush_all_dcache(struct htp_context * ctx) {
+ struct htp_thread_trace * tr = &ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+ qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+ hex_l2fetch_block(ctx, ctx->footprint);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+ memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
+}
+
+static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
+ struct l2flush_multi_task * task = (struct l2flush_multi_task *) data;
+ (void) n;
+
+ const uint32_t gb_first = i * task->blocks_per_thread;
+ uint32_t gb_last = gb_first + task->blocks_per_thread;
+ if (gb_last > task->total_blocks) {
+ gb_last = task->total_blocks;
+ }
+ if (gb_first >= gb_last) {
+ return;
+ }
+
+ struct htp_thread_trace * tr = &task->trace[i];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
+
+ for (uint32_t r = 0; r < task->n_ranges; r++) {
+ const struct l2flush_range * rg = &task->ranges[r];
+ const uint32_t rb_first = rg->block_first;
+ const uint32_t rb_last = rg->block_first + rg->n_blocks;
+
+ const uint32_t lo = gb_first > rb_first ? gb_first : rb_first;
+ const uint32_t hi = gb_last < rb_last ? gb_last : rb_last;
+ if (lo >= hi) {
+ continue;
+ }
+
+ const uint32_t s = rg->start + (lo - rb_first) * HEX_L2_BLOCK_SIZE;
+ uint32_t e = rg->start + (hi - rb_first) * HEX_L2_BLOCK_SIZE;
+ if (e > rg->end) {
+ e = rg->end;
+ }
+ hex_l2flush((void *) (uintptr_t) s, e - s);
+ }
+
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
+}
+
+void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+ const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
+ uint32_t n_pending = 0;
+
+ for (uint32_t i = 0; i < n; i++) {
+ const struct htp_tensor * t = tensors[i];
+ if (!t) continue;
+
+ uint32_t t_start = t->data;
+ uint32_t t_end = t_start + t->size;
+
+ bool merged = false;
+ for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
+ struct htp_dirty_range * r = &ctx->dirty_ranges[j];
+ if (!r->start) continue;
+
+ if (r->start <= t_end && t_start <= r->end) {
+ uint32_t new_start = (t_start < r->start) ? t_start : r->start;
+ uint32_t new_end = (t_end > r->end) ? t_end : r->end;
+ r->start = new_start;
+ r->end = new_end;
+ merged = true;
+ }
+ }
+
+ if (!merged) {
+ pending[n_pending++] = t;
+ }
+ }
+
+ if (n_pending == 0) {
+ return;
+ }
+
+ uint32_t empty_indices[HTP_MAX_DIRTY_RANGES];
+ uint32_t active_indices[HTP_MAX_DIRTY_RANGES];
+ uint32_t n_active = 0;
+ uint32_t n_empty = 0;
+ for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
+ if (ctx->dirty_ranges[j].start) {
+ active_indices[n_active++] = j;
+ } else {
+ empty_indices[n_empty++] = j;
+ }
+ }
+
+ if (n_pending <= n_empty) {
+ for (uint32_t i = 0; i < n_pending; i++) {
+ uint32_t idx = empty_indices[i];
+ struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+ r->start = pending[i]->data;
+ r->end = pending[i]->data + pending[i]->size;
+ r->bi = pending[i]->bi;
+ }
+ return;
+ }
+
+ uint32_t n_evict = n_pending - n_empty;
+ uint32_t total_evict_size = 0;
+ for (uint32_t i = 0; i < n_evict; i++) {
+ uint32_t idx = active_indices[i];
+ struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+ total_evict_size += r->end - r->start;
+ }
+
+ if (total_evict_size > HEX_L2_FLUSH_ALL_THRESHOLD) {
+ flush_all_dcache(ctx);
+ for (uint32_t i = 0; i < n_pending; i++) {
+ struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+ r->start = pending[i]->data;
+ r->end = pending[i]->data + pending[i]->size;
+ r->bi = pending[i]->bi;
+ }
+ return;
+ }
+
+ if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_OP_MAX_INPUTS) {
+ struct l2flush_multi_task task;
+ task.trace = ctx->trace;
+ task.n_ranges = n_evict;
+
+ uint32_t block_acc = 0;
+ for (uint32_t i = 0; i < n_evict; i++) {
+ uint32_t idx = active_indices[i];
+ struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+
+ struct l2flush_range * rg = &task.ranges[i];
+ rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
+ rg->end = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
+ rg->block_first = block_acc;
+ rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+ block_acc += rg->n_blocks;
+ }
+
+ task.total_blocks = block_acc;
+ task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
+
+ work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
+ } else {
+ struct htp_thread_trace * tr = &ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+ for (uint32_t i = 0; i < n_evict; i++) {
+ uint32_t idx = active_indices[i];
+ struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+ uint32_t size = r->end - r->start;
+ hex_l2flush((void *) (uintptr_t) r->start, size);
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+ }
+
+ for (uint32_t i = 0; i < n_evict; i++) {
+ uint32_t idx = active_indices[i];
+ struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+ r->start = pending[i]->data;
+ r->end = pending[i]->data + pending[i]->size;
+ r->bi = pending[i]->bi;
+ }
+
+ for (uint32_t i = 0; i < n_empty; i++) {
+ uint32_t idx = empty_indices[i];
+ struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+ r->start = pending[n_evict + i]->data;
+ r->end = pending[n_evict + i]->data + pending[n_evict + i]->size;
+ r->bi = pending[n_evict + i]->bi;
+ }
+}
+
+static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
+ uint32_t t_start = t->data;
+ uint32_t t_end = t_start + t->size;
+
+ for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+ struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+ if (!r->start) continue;
+
+ if (r->start < t_end && t_start < r->end) {
+ if (t_start <= r->start && r->end <= t_end) {
+ r->start = 0;
+ } else if (t_start <= r->start) {
+ r->start = t_end;
+ } else if (r->end <= t_end) {
+ r->end = t_start;
+ }
+ }
+ }
+}
+
+static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) {
+ uint32_t t_start = t->data;
+ uint32_t t_end = t_start + t->size;
+
+ for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+ struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+ if (!r->start) continue;
+
+ if (r->start < t_end && t_start < r->end) {
+ return true;
+ }
+ }
+ return false;
+}
+
+void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+ const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
+ uint32_t n_dirty = 0;
+ uint64_t total_dirty = 0;
+
+ for (uint32_t i = 0; i < n; i++) {
+ const struct htp_tensor * t = tensors[i];
+ if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) {
+ dirty_tensors[n_dirty++] = t;
+ total_dirty += t->size;
+ }
+ }
+
+ if (total_dirty == 0) {
+ return;
+ }
+
+ if (total_dirty > HEX_L2_FLUSH_ALL_THRESHOLD) {
+ flush_all_dcache(ctx);
+ return;
+ }
+
+ if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
+ struct l2flush_multi_task task;
+ task.trace = ctx->trace;
+ task.n_ranges = 0;
+
+ uint32_t block_acc = 0;
+ for (uint32_t i = 0; i < n_dirty; i++) {
+ const struct htp_tensor * t = dirty_tensors[i];
+ make_tensor_clean(ctx, t);
+
+ struct l2flush_range * rg = &task.ranges[task.n_ranges++];
+ rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
+ rg->end = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
+ rg->block_first = block_acc;
+ rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+ block_acc += rg->n_blocks;
+ }
+
+ task.total_blocks = block_acc;
+ task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
+
+ work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
+ return;
+ }
+
+ struct htp_thread_trace * tr = &ctx->trace[0];
+ for (uint32_t i = 0; i < n_dirty; i++) {
+ const struct htp_tensor * t = dirty_tensors[i];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
+ hex_l2flush((void *) (uintptr_t) t->data, t->size);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
+ make_tensor_clean(ctx, t);
+ }
+}
--- /dev/null
+#ifndef HTP_TENSOR_H
+#define HTP_TENSOR_H
+
+#include <stdint.h>
+#include "htp-ops.h"
+#include "hex-bitmap.h"
+
+static inline void * htp_tensor_data(const struct htp_tensor * t) {
+ return (void *) (uintptr_t) t->data;
+}
+
+static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
+ return (uint32_t *) &t->flags;
+}
+
+struct htp_context;
+void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+
+#endif // HTP_TENSOR_H
}
}
}
+static inline void hvx_mad_f32_f16_aa_vec(float * restrict y, const void * restrict x, HVX_Vector S0, uint32_t n) {
+ const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
+
+ HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
+ HVX_Vector * restrict vy = (HVX_Vector *) y;
+
+ uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; ++i) {
+ vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
+ }
+
+ if (nloe) {
+ HVX_VectorPair xy_p = vy_p[i];
+ xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
+
+ HVX_Vector xy = Q6_V_lo_W(xy_p);
+ i = 2 * i; // index for vy
+
+ if (nloe >= VLEN_FP32) {
+ vy[i] = xy;
+ nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
+ }
+
+ if (nloe) {
+ hvx_vec_store_a(&vy[i], nloe * 4, xy);
+ }
+ }
+}
+
+static inline void hvx_mad_f32_f16_aa_rx2_vec(float * restrict y, const void * restrict x0, const void * restrict x1,
+ HVX_Vector S0, HVX_Vector S1, uint32_t n) {
+ const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
+ const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
+
+ HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
+ HVX_Vector * restrict vy = (HVX_Vector *) y;
+
+ uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; ++i) {
+ vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
+ vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
+ }
+
+ if (nloe) {
+ HVX_VectorPair xy_p = vy_p[i];
+ xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
+ xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
+
+ HVX_Vector xy = Q6_V_lo_W(xy_p);
+ i = 2 * i; // index for vy
+
+ if (nloe >= VLEN_FP32) {
+ vy[i] = xy;
+ nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
+ }
+
+ if (nloe) {
+ hvx_vec_store_a(&vy[i], nloe * 4, xy);
+ }
+ }
+}
static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) {
assert((size_t) dst % 128 == 0);
--- /dev/null
+#ifndef HVX_NORM_H
+#define HVX_NORM_H
+
+#include <stdint.h>
+#include "hvx-base.h"
+#include "hvx-reduce.h"
+#include "hvx-inverse.h"
+#include "hvx-sqrt.h"
+#include "hvx-repl.h"
+
+static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
+ uint8_t * restrict dst,
+ const int num_elems,
+ float epsilon) {
+
+ const HVX_Vector * restrict v_src = (HVX_Vector *) src;
+ HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
+
+ const int nvec = num_elems / VLEN_FP32; // number of full vectors
+ const int nloe = num_elems % VLEN_FP32; // leftover elements
+
+ // Compute sum of squares for full vectors
+ HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
+ HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+ }
+
+ // Handle tail elements using vectorized ops with masking
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+ }
+
+ // Reduce HVX sum
+ sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+
+ HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
+ HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
+ HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
+ HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
+
+ // Scale full vectors
+ HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+ v_dst[i] = Q6_Vsf_equals_Vqf32(v2);
+ }
+
+ // Handle tail elements using vectorized ops with masking
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+ HVX_Vector result = Q6_Vsf_equals_Vqf32(v2);
+
+ // Store with masking to avoid overwriting memory beyond the tensor
+ hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
+ }
+}
+
+static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
+ const uint8_t * restrict weight,
+ uint8_t * restrict dst,
+ const int num_elems,
+ float epsilon) {
+ const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
+ const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
+ HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
+
+ const int nvec = num_elems / VLEN_FP32; // number of full vectors
+ const int nloe = num_elems % VLEN_FP32; // leftover elements
+
+ // Compute sum of squares for full vectors
+ HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
+ HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+ }
+
+ // Handle tail elements using vectorized ops with masking
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+ }
+
+ // Reduce HVX sum
+ sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+
+ HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
+ HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
+ HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
+ HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
+
+ // Scale and multiply
+ HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+ HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
+ HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[i]);
+ v_dst[i] = Q6_Vsf_equals_Vqf32(result);
+ }
+
+ // Handle tail elements using vectorized ops with masking
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+ HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
+ HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[nvec]);
+ HVX_Vector res_v = Q6_Vsf_equals_Vqf32(result);
+
+ // Store with masking to avoid overwriting memory beyond the tensor
+ hvx_vec_store_a(&v_dst[nvec], nloe * 4, res_v);
+ }
+}
+
+static inline void hvx_fast_norm_f32(const uint8_t * restrict src,
+ uint8_t * restrict dst,
+ const int num_elems,
+ float epsilon) {
+
+ const HVX_Vector * restrict v_src = (HVX_Vector *) src;
+ HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
+
+ const int nvec = num_elems / VLEN_FP32; // number of full vectors
+ const int nloe = num_elems % VLEN_FP32; // leftover elements
+
+ // Compute sum of squares and sum of values for full vectors
+ HVX_Vector sum_sq_v = Q6_V_vsplat_R(0x00000000);
+ HVX_Vector sum_x_v = Q6_V_vsplat_R(0x00000000);
+ HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
+ sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
+ }
+
+ // Handle tail elements using vectorized ops with masking
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
+ sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
+ }
+
+ // Reduce HVX sums
+ sum_sq_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_sq_v));
+ sum_x_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_x_v));
+
+ HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
+ HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
+ HVX_Vector mean_sq_v = Q6_Vqf32_vmpy_VsfVsf(sum_sq_v, denom_v);
+ HVX_Vector mean_x_v = Q6_Vqf32_vmpy_VsfVsf(sum_x_v, denom_v);
+ HVX_Vector mean_x_sq_v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(mean_x_v), Q6_Vsf_equals_Vqf32(mean_x_v));
+ HVX_Vector var_v = Q6_Vqf32_vsub_Vqf32Vqf32(mean_sq_v, mean_x_sq_v);
+ HVX_Vector var_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(var_v, epsilon_v);
+
+ // scale = rsqrt(variance + epsilon), mean_x broadcast for subtraction
+ HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(var_epsilon_v));
+ HVX_Vector mean_x_b = hvx_vec_repl_f32(Q6_Vsf_equals_Vqf32(mean_x_v));
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
+ HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
+ v_dst[i] = Q6_Vsf_equals_Vqf32(v3);
+ }
+
+ // Handle tail elements using vectorized ops with masking
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
+ HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
+ HVX_Vector result = Q6_Vsf_equals_Vqf32(v3);
+
+ // Store with masking to avoid overwriting memory beyond the tensor
+ hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
+ }
+}
+
+static inline void hvx_fast_l2_norm_f32(const uint8_t * restrict src,
+ uint8_t * restrict dst,
+ const int num_elems,
+ float epsilon) {
+
+ const HVX_Vector * restrict v_src = (HVX_Vector *) src;
+ HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
+
+ HVX_Vector sum_v = hvx_vec_splat_f32(0.0f);
+
+ const int nvec = num_elems / VLEN_FP32;
+ const int nloe = num_elems % VLEN_FP32;
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
+ }
+
+ // Include tail elements in the sum-of-squares using a predicate mask
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
+ }
+
+ // Compute scale = 1/fmax(sqrt(sum), epsilon) entirely in HVX registers.
+ // hvx_vec_rsqrt_f32 + hvx_vec_inverse_f32 avoids scalar extraction.
+ HVX_Vector sum_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+ HVX_Vector rsqrt_v = hvx_vec_rsqrt_f32(sum_sf); // 1/sqrt(sum)
+ HVX_Vector sqrt_v = hvx_vec_inverse_f32(rsqrt_v); // sqrt(sum)
+ HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+ HVX_Vector denom_v = Q6_Vsf_vmax_VsfVsf(sqrt_v, epsilon_v); // fmax(sqrt(sum), epsilon)
+ HVX_Vector scale_v = hvx_vec_inverse_f32(denom_v); // 1/fmax(sqrt(sum), epsilon)
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ v_dst[i] = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
+ }
+
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector result = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
+ hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
+ }
+}
+
+#endif // HVX_NORM_H
}
}
+// Signed 32-bit Integer Max variants
+
+static inline HVX_Vector hvx_vec_reduce_max_n_i32(HVX_Vector in, unsigned int n) {
+ unsigned int total = n * 4; // total vec nbytes
+ unsigned int width = 4; // int32 nbytes
+
+ HVX_Vector max_val = in, max_t;
+ while (width < total) {
+ max_t = Q6_V_vror_VR(max_val, width); // rotate right
+ max_val = Q6_Vw_vmax_VwVw(max_t, max_val); // elementwise signed max
+ width = width << 1;
+ }
+ return max_val;
+}
+
+static inline HVX_Vector hvx_vec_reduce_max_i32(HVX_Vector in) {
+ return hvx_vec_reduce_max_n_i32(in, 32);
+}
+
+static inline int32_t hvx_reduce_max_i32_a(const uint8_t * restrict src, const int num_elems) {
+ HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
+ HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
+ assert((uintptr_t) src % 128 == 0);
+ hvx_reduce_loop_body(HVX_Vector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
+}
+
+static inline int32_t hvx_reduce_max_i32_u(const uint8_t * restrict src, const int num_elems) {
+ HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
+ HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
+ hvx_reduce_loop_body(HVX_UVector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
+}
+
+static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int num_elems) {
+ if (hex_is_aligned((void *) src, 128)) {
+ return hvx_reduce_max_i32_a(src, num_elems);
+ } else {
+ return hvx_reduce_max_i32_u(src, num_elems);
+ }
+}
+
#undef hvx_reduce_loop_body
#undef HVX_REDUCE_MAX_OP
#undef HVX_REDUCE_SUM_OP
#include "hvx-base.h"
#include "hvx-pow.h"
#include "hvx-log.h"
+#include "hvx-norm.h"
#endif /* HVX_UTILS_H */
--- /dev/null
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <HAP_perf.h>
+#include <hexagon_protos.h>
+#include <hexagon_types.h>
+#include <string.h>
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "hvx-utils.h"
+#include "hex-dma.h"
+#include "hex-profile.h"
+#include "htp-vtcm.h"
+
+struct htp_im2col_context {
+ struct htp_ops_context * octx;
+ uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel)
+
+ uint32_t pe_rows_per_thread; // N*OH rows per worker
+ uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256
+ uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256
+
+ // Patch-embed DMA path VTCM ping-pong.
+ uint8_t * pe_vtcm_src; // base of the 2x src buffers region
+ uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region
+ uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes
+ uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes
+};
+
+// Per-op VTCM layout for the patch-embed DMA path
+struct htp_im2col_vtcm_layout {
+ size_t off_src;
+ size_t off_dst;
+ size_t src_bytes_per_thread;
+ size_t dst_bytes_per_thread;
+ size_t total_bytes;
+};
+
+static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L,
+ size_t src_row_bytes,
+ size_t dst_row_bytes,
+ uint32_t n_threads) {
+ L->src_bytes_per_thread = 2 * src_row_bytes;
+ L->dst_bytes_per_thread = 2 * dst_row_bytes;
+
+ L->off_src = 0;
+ L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads;
+ L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
+}
+
+#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
+ static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
+ struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
+ struct htp_ops_context * octx = ictx->octx; \
+ struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
+ const struct htp_tensor * restrict src1 = octx->src[1]; \
+ const struct htp_tensor * restrict dst = octx->dst; \
+ const int32_t s0 = octx->op_params[0]; \
+ const int32_t s1 = octx->op_params[1]; \
+ const int32_t p0 = octx->op_params[2]; \
+ const int32_t p1 = octx->op_params[3]; \
+ const int32_t d0 = octx->op_params[4]; \
+ const int32_t d1 = octx->op_params[5]; \
+ const uint32_t N = src1->ne[3]; \
+ const uint32_t IC = src1->ne[2]; \
+ const uint32_t IH = src1->ne[1]; \
+ const uint32_t IW = src1->ne[0]; \
+ const uint32_t KH = octx->src[0]->ne[1]; \
+ const uint32_t KW = octx->src[0]->ne[0]; \
+ const uint32_t OH = dst->ne[2]; \
+ const uint32_t OW = dst->ne[1]; \
+ const uint32_t patch_stride = IC * KH * KW; \
+ const float * restrict src_data = (const float *) src1->data; \
+ DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
+ const uint32_t npatches = N * OH * OW; \
+ const uint32_t patch_start = ictx->npatches_per_thread * ith; \
+ const uint32_t patch_end = MIN(patch_start + ictx->npatches_per_thread, npatches); \
+ if (patch_start >= patch_end) { \
+ return; \
+ } \
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
+ for (uint32_t p = patch_start; p < patch_end; p++) { \
+ const uint32_t iow = p % OW; \
+ const uint32_t ioh = (p / OW) % OH; \
+ const uint32_t in = p / (OW * OH); \
+ DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \
+ for (uint32_t iic = 0; iic < IC; iic++) { \
+ const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \
+ for (uint32_t ikh = 0; ikh < KH; ikh++) { \
+ const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \
+ DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
+ if (iih < 0 || iih >= (int32_t) IH) { \
+ SPLAT_FN(out_run, 0.0f, KW); \
+ continue; \
+ } \
+ const int32_t iiw0 = (int32_t) iow * s0 - p0; \
+ const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \
+ if (d0 == 1) { \
+ /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \
+ const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \
+ int32_t hi = (int32_t) IW - iiw0; \
+ if (hi > (int32_t) KW) { \
+ hi = (int32_t) KW; \
+ } \
+ if (hi <= lo) { \
+ SPLAT_FN(out_run, 0.0f, KW); \
+ } else { \
+ if (lo > 0) { \
+ SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \
+ } \
+ COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \
+ (uint32_t) (hi - lo)); \
+ if (hi < (int32_t) KW) { \
+ SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \
+ } \
+ } \
+ continue; \
+ } \
+ for (uint32_t ikw = 0; ikw < KW; ikw++) { \
+ const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \
+ out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \
+ (DST_CTYPE) 0.0f : \
+ (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \
+ } \
+ } \
+ } \
+ } \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
+ }
+
+IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
+IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")
+
+#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
+ static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
+ struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
+ struct htp_ops_context * octx = ictx->octx; \
+ struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
+ const struct htp_tensor * restrict src1 = octx->src[1]; \
+ const struct htp_tensor * restrict dst = octx->dst; \
+ const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \
+ const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \
+ const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \
+ const uint32_t patch_stride = IC * KH * KW; \
+ const float * restrict src_data = (const float *) src1->data; \
+ DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
+ dma_queue * dmaq = octx->ctx->dma[ith]; \
+ uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \
+ uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \
+ float * srcb = (float *) src_base; \
+ DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \
+ const uint32_t nrows = N * OH; \
+ const uint32_t per_thread = ictx->pe_rows_per_thread; \
+ const uint32_t row_start = per_thread * ith; \
+ const uint32_t row_end = MIN(row_start + per_thread, nrows); \
+ if (row_start >= row_end) \
+ return; \
+ for (uint32_t r = row_start; r < row_end; r++) { \
+ const uint32_t in = r / OH; \
+ const uint32_t ioh = r % OH; \
+ for (uint32_t ikh = 0; ikh < KH; ikh++) { \
+ int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
+ int ok = (iih >= 0 && iih < (int32_t) IH); \
+ for (uint32_t iic = 0; iic < IC; iic++) { \
+ float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \
+ const float * _vsrc = \
+ ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \
+ dma_queue_push_ddr_to_vtcm( \
+ dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
+ IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \
+ } \
+ } \
+ for (uint32_t i = 0; i < IC * KH; i++) \
+ dma_queue_pop(dmaq); \
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \
+ for (uint32_t iow = 0; iow < OW; iow++) { \
+ DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \
+ for (uint32_t ikh = 0; ikh < KH; ikh++) { \
+ int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
+ for (uint32_t iic = 0; iic < IC; iic++) { \
+ DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
+ if (iih < 0 || iih >= (int32_t) IH) { \
+ SPLAT_FN(out_run, 0.0f, KW); \
+ continue; \
+ } \
+ const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \
+ COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \
+ } \
+ } \
+ } \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \
+ DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \
+ dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \
+ OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \
+ dma_queue_flush(dmaq); \
+ } \
+ }
+
+IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
+IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32")
+
+static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
+ const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];
+ const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];
+ const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];
+ const int is_2D = octx->op_params[6] == 1;
+ if (!is_2D) {
+ return false;
+ }
+ if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
+ return false;
+ }
+ const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
+ if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
+ return false; // non-overlapping
+ }
+ if (p0 != 0 || p1 != 0) {
+ return false; // no padding
+ }
+ if (d0 != 1 || d1 != 1) {
+ return false; // no dilation
+ }
+ return true;
+}
+
+// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path.
+// Returns false if it doesn't fit the VTCM budget (caller falls back).
+static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx,
+ struct htp_im2col_context * ictx,
+ uint32_t n_threads) {
+ const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
+ const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
+ const uint32_t OW = octx->dst->ne[1];
+ const uint32_t patch_stride = IC * KH * KW;
+
+ ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256);
+ const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
+ ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256);
+
+ // 2 src + 2 dst buffers per thread (ping-pong), src region first then dst.
+ struct htp_im2col_vtcm_layout L;
+ htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads);
+ if (L.total_bytes > octx->ctx->vtcm_size) {
+ return false;
+ }
+
+ uint8_t * const base = octx->ctx->vtcm_base;
+ ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
+ ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+ ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
+ ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
+ return true;
+}
+
+int op_im2col(struct htp_ops_context * octx) {
+ const struct htp_tensor * src1 = octx->src[1];
+ const struct htp_tensor * dst = octx->dst;
+
+ if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) {
+ FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported");
+ return HTP_STATUS_NO_SUPPORT;
+ }
+
+ const uint32_t N = src1->ne[3];
+ const uint32_t OH = dst->ne[2];
+ const uint32_t OW = dst->ne[1];
+ const uint32_t npatches = N * OH * OW;
+ const uint32_t n_threads = MIN(octx->n_threads, npatches);
+
+ if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
+ return HTP_STATUS_OK;
+ }
+
+ struct htp_im2col_context ictx = { 0 };
+ ictx.octx = octx;
+ ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;
+
+ // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
+ // everything else (padding/dilation/stride edges) -> pure-DDR kernel.
+ if (im2col_use_patchembed_dma(octx)) {
+ const uint32_t nrows = N * OH;
+ const uint32_t pth = MIN(octx->n_threads, nrows);
+ if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
+ ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
+ if (dst->type == HTP_TYPE_F16) {
+ work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
+ } else {
+ work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
+ }
+ return HTP_STATUS_OK;
+ }
+ // else: doesn't fit -> fall through to the pure-DDR kernel below.
+ }
+
+ if (dst->type == HTP_TYPE_F16) {
+ work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
+ } else {
+ work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads);
+ }
+ return HTP_STATUS_OK;
+}
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
+#include "hex-bitmap.h"
#include "htp-ctx.h"
#include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"
#include "htp_iface.h"
-#include "worker-pool.h"
-
-AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) {
- struct htp_context * ctx;
- int err = 0;
+#include "work-queue.h"
+#include "hex-profile.h"
- ctx = calloc(1, sizeof(*ctx));
- if (ctx == NULL) {
- return AEE_ENOMEMORY;
- }
-
- // Use the context structure as the handle
- *handle = (remote_handle64) ctx;
-
- // Enable FARF logs
- HAP_setFARFRuntimeLoggingParams(0xffff, NULL, 0);
-
- // Set client class
- {
- HAP_power_request_t request;
- memset(&request, 0, sizeof(HAP_power_request_t));
- request.type = HAP_power_set_apptype;
- request.apptype = HAP_POWER_COMPUTE_CLIENT_CLASS;
-
- if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
- return err;
- }
- }
+#define HMX_QUEUE_CAPACITY 16
+#define HMX_QUEUE_STACK_SIZE 16384
+#define WORK_QUEUE_CAPACITY 16
+#define WORK_QUEUE_STACK_SIZE 16384
+#define MAIN_THREAD_STACK_SIZE 32768
- {
- HAP_power_request_t request;
- memset(&request, 0, sizeof(request));
-
- request.type = HAP_power_set_DCVS_v3;
- request.dcvs_v3.set_dcvs_enable = TRUE;
- request.dcvs_v3.dcvs_enable = FALSE;
- request.dcvs_v3.set_bus_params = TRUE;
- request.dcvs_v3.bus_params.min_corner = HAP_DCVS_VCORNER_MAX;
- request.dcvs_v3.bus_params.max_corner = HAP_DCVS_VCORNER_MAX;
- request.dcvs_v3.bus_params.target_corner = HAP_DCVS_VCORNER_MAX;
- request.dcvs_v3.set_core_params = TRUE;
- request.dcvs_v3.core_params.min_corner = HAP_DCVS_VCORNER_MAX;
- request.dcvs_v3.core_params.max_corner = HAP_DCVS_VCORNER_MAX;
- request.dcvs_v3.core_params.target_corner = HAP_DCVS_VCORNER_MAX;
- request.dcvs_v3.set_sleep_disable = TRUE;
- request.dcvs_v3.sleep_disable = TRUE;
-
-#if (__HEXAGON_ARCH__ >= 79)
- HAP_set_dcvs_v3_protected_bus_corners(&request, 1);
-#endif
- if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
- return err;
- }
+_Static_assert(WORK_QUEUE_MAX_N_THREADS >= HTP_MAX_NTHREADS,
+ "work-queue thread cap must be >= HTP_MAX_NTHREADS");
- memset(&request, 0, sizeof(request));
- request.type = HAP_power_set_HVX;
- request.hvx.power_up = TRUE;
- if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
- return err;
- }
- }
+struct htp_handle {
+ struct htp_context * ctx;
+};
-#if __HVX_ARCH__ >= 75
- {
- // Power on HMX and set HMX clock
- HAP_power_request_t request;
- memset(&request, 0, sizeof(HAP_power_request_t));
- request.type = HAP_power_set_HMX_v2;
- request.hmx_v2.set_power = TRUE;
- request.hmx_v2.power_up = TRUE;
- request.hmx_v2.set_clock = TRUE;
- request.hmx_v2.target_corner = HAP_DCVS_EXP_VCORNER_MAX;
- request.hmx_v2.min_corner = HAP_DCVS_EXP_VCORNER_MAX;
- request.hmx_v2.max_corner = HAP_DCVS_EXP_VCORNER_MAX;
- request.hmx_v2.perf_mode = HAP_CLK_PERF_HIGH;
- FARF(ALWAYS, "Setting HMX clock\n");
- err = HAP_power_set((void *) ctx, &request);
- if (err != AEE_SUCCESS) {
- FARF(ERROR, "ggml-hex: error setting HMX clock.");
- return err;
- }
- }
-#else
- {
- // Power on HMX
- HAP_power_request_t request;
- memset(&request, 0, sizeof(HAP_power_request_t));
- request.type = HAP_power_set_HMX;
- request.hmx.power_up = TRUE;
- FARF(ALWAYS, "Powering HMX on\n");
- err = HAP_power_set((void *) ctx, &request);
- if (err != AEE_SUCCESS) {
- FARF(ERROR, "ggml-hex: error powering on HMX.");
- return err;
- }
+AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) {
+ (void) uri;
+ struct htp_handle * h = calloc(1, sizeof(*h));
+ if (h == NULL) {
+ return AEE_ENOMEMORY;
}
-#endif
+ *handle = (remote_handle64) h;
return AEE_SUCCESS;
}
AEEResult htp_iface_etm(remote_handle64 handle, uint32_t enable) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h) {
+ return AEE_EBADPARM;
+ }
+
int err = enable ? HAP_user_etm_enable() : HAP_user_etm_disable();
if (err) {
if (err == AEE_EVERSIONNOTSUPPORT) {
}
AEEResult htp_iface_profiler(remote_handle64 handle, uint32_t mode, const htp_iface_pmu_conf* pmu_conf) {
- struct htp_context * ctx = (struct htp_context *) handle;
- if (!ctx) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h || !h->ctx) {
return AEE_EBADPARM;
}
+ struct htp_context * ctx = h->ctx;
if (mode == HTP_PROF_PMU) {
const uint32_t* events = pmu_conf->events;
}
AEEResult htp_iface_close(remote_handle64 handle) {
- struct htp_context * ctx = (struct htp_context *) handle;
-
- if (!ctx) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h) {
return AEE_EBADPARM;
}
- if (ctx->queue) {
- FARF(ERROR, "Closing handle with queue still open");
- return AEE_EITEMBUSY;
- }
+ struct htp_context * ctx = h->ctx;
+ if (ctx) {
+ if (ctx->dsp_queue) {
+ FARF(ERROR, "Closing handle with queue still open");
+ return AEE_EITEMBUSY;
+ }
- // release the mmaps (if any)
- for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
- if (ctx->mmap[i].size) {
+ // release the mmaps (if any)
+ for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
+ if (ctx->mmap[i].size) {
#if __HVX_ARCH__ > 73
- HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size);
+ HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size);
#else
- HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
+ HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
#endif
- ctx->mmap[i].size = 0;
- ctx->mmap[i].base = NULL;
- ctx->mmap[i].fd = -1;
+ ctx->mmap[i].size = 0;
+ ctx->mmap[i].base = NULL;
+ ctx->mmap[i].fd = -1;
+ }
}
- }
- if (ctx->profiler) {
- qurt_pmu_enable(1);
- }
+ if (ctx->profiler) {
+ qurt_pmu_enable(1);
+ }
- if (ctx->etm) {
- HAP_user_etm_disable();
+ if (ctx->etm) {
+ HAP_user_etm_disable();
+ }
+
+ // Free the unified block (ctx is the base address of the block)
+ free(ctx);
+ h->ctx = NULL;
}
- free(ctx);
+ free(h);
return AEE_SUCCESS;
}
AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
- struct htp_context * ctx = (struct htp_context *) handle;
- if (!ctx) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h || !h->ctx) {
return AEE_EBADPARM;
}
+ struct htp_context * ctx = h->ctx;
// See if we already have this mapping
for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
}
AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) {
- struct htp_context * ctx = (struct htp_context *) handle;
- if (!ctx) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h || !h->ctx) {
return AEE_EBADPARM;
}
+ struct htp_context * ctx = h->ctx;
for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
struct htp_mmap *m = &ctx->mmap[i];
}
}
+static void htp_main_thread(void * context);
static void htp_packet_callback(dspqueue_t queue, int error, void * context);
static void htp_error_callback(dspqueue_t queue, int error, void * context);
AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp_queue_id, uint32_t n_hvx, uint32_t n_hmx, uint64_t max_vmem) {
- struct htp_context * ctx = (struct htp_context *) handle;
-
- if (!ctx) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h) {
return AEE_EBADPARM;
}
- if (ctx->queue) {
+ if (h->ctx) {
FARF(ERROR, "Queue already open");
return AEE_EITEMBUSY;
}
- // Import queue created on the CPU
- int err = dspqueue_import(dsp_queue_id, // Queue ID from dspqueue_export
- htp_packet_callback, // Packet callback
- htp_error_callback, // Error callback; no errors expected on the DSP
- (void *) ctx, // Callback context
- &ctx->queue);
+ // Cache the original FastRPC thread priority, then calculate compute priority
+ int fastrpc_tid = qurt_thread_get_id();
+ int fastrpc_prio = qurt_thread_get_priority(fastrpc_tid);
+ int main_prio = fastrpc_prio - 10;
+ if (main_prio < 1) main_prio = 1;
+
+ dspqueue_t dsp_queue = NULL;
+ bool use_callbacks = false;
+
+ // Import queue with NULL callbacks to avoid starting dspueue internal threads
+ int err = dspqueue_import(dsp_queue_id, NULL, NULL, (void *) h, &dsp_queue);
+ if (err == AEE_EBADPARM) {
+ // Fallback for devices that don't support NULL callbacks
+ FARF(HIGH, "dspqueue import with NULL callbacks failed, trying with callbacks");
+ use_callbacks = true;
+ err = dspqueue_import(dsp_queue_id, htp_packet_callback, htp_error_callback, (void *) h, &dsp_queue);
+ }
+
if (err) {
FARF(ERROR, "Queue import failed with 0x%08x", (unsigned) err);
return err;
}
+ qurt_sysenv_max_hthreads_t hw_threads;
+ qurt_sysenv_get_max_hw_threads(&hw_threads);
+ uint32_t hw_nhvx = (qurt_hvx_get_units() >> 8) & 0xFF;
+
+ if (n_hvx == 0) {
+ n_hvx = hw_nhvx;
+ }
+ if (n_hvx > hw_threads.max_hthreads) {
+ n_hvx = hw_threads.max_hthreads;
+ }
+ if (n_hvx > HTP_MAX_NTHREADS) {
+ n_hvx = HTP_MAX_NTHREADS;
+ }
+
+ // layout segments of our contiguous block
+
+ // 1. htp_context : sits at the base (block is 4K-aligned via memalign below)
+ size_t offset = sizeof(struct htp_context);
+
+ // 2. main_stack
+ size_t offset_main_stack = 0;
+ size_t size_main_stack = 0;
+ if (!use_callbacks) {
+ offset_main_stack = hex_align_up(offset, 4096);
+ size_main_stack = MAIN_THREAD_STACK_SIZE;
+ offset = offset_main_stack + size_main_stack;
+ }
+
+ // 3. work_queue
+ size_t wq_align = work_queue_alignof();
+ size_t offset_wq = hex_align_up(offset, wq_align);
+ size_t size_wq = work_queue_sizeof(n_hvx, WORK_QUEUE_CAPACITY, WORK_QUEUE_STACK_SIZE);
+ offset = offset_wq + size_wq;
+
+ // 4. dma_queue
+ size_t dma_align = dma_queue_alignof();
+ size_t offset_dma = hex_align_up(offset, dma_align);
+ size_t size_dma = 0;
+ for (uint32_t i = 0; i < n_hvx; i++) {
+ size_dma = hex_align_up(size_dma, dma_queue_alignof());
+ size_dma += dma_queue_sizeof(256);
+ size_dma = hex_align_up(size_dma, dma_queue_alignof());
+ size_dma += dma_queue_alias_sizeof();
+ }
+ offset = offset_dma + size_dma;
+
+ // 5. hmx_queue
+ size_t offset_hmx = 0;
+ size_t size_hmx = 0;
+ if (n_hmx) {
+ size_t hmx_align = hmx_queue_alignof();
+ offset_hmx = hex_align_up(offset, hmx_align);
+ size_hmx = hmx_queue_sizeof(HMX_QUEUE_CAPACITY, HMX_QUEUE_STACK_SIZE);
+ offset = offset_hmx + size_hmx;
+ }
+
+ size_t footprint = hex_align_up(offset, 128);
+
+ void * block = memalign(4096, footprint);
+ if (!block) {
+ FARF(ERROR, "Unable to allocate unified block of size %zu\n", footprint);
+ dspqueue_close(dsp_queue);
+ return AEE_ENOMEMORY;
+ }
+ memset(block, 0, footprint);
+
+ h->ctx = (struct htp_context *) block;
+ struct htp_context * ctx = h->ctx;
+ ctx->footprint = footprint;
+
+ ctx->thread_id = fastrpc_tid;
+ ctx->thread_prio = main_prio;
ctx->max_vmem = max_vmem;
- ctx->thread_id = qurt_thread_get_id();
- ctx->thread_prio = qurt_thread_get_priority(ctx->thread_id);
+ ctx->dsp_queue = dsp_queue;
- // allocate VTCM
err = vtcm_alloc(ctx);
if (err != AEE_SUCCESS) {
FARF(ERROR, "Unable to allocate VTCM");
+ htp_iface_stop(handle);
return AEE_ENOMEMORY;
}
- ctx->hmx_enabled = n_hmx;
- ctx->hmx_queue = NULL;
- if (n_hmx) {
- ctx->hmx_queue = hmx_queue_create(16, ctx->vtcm_rctx);
- if (ctx->hmx_queue) {
- ctx->hmx_queue->trace = &ctx->trace[HTP_MAX_NTHREADS];
- } else {
- FARF(ERROR, "hmx-queue-create failed");
- ctx->hmx_enabled = false;
+ HAP_setFARFRuntimeLoggingParams(0xffff, NULL, 0);
+
+ // Set client class
+ {
+ HAP_power_request_t request;
+ memset(&request, 0, sizeof(HAP_power_request_t));
+ request.type = HAP_power_set_apptype;
+ request.apptype = HAP_POWER_COMPUTE_CLIENT_CLASS;
+
+ if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
+ htp_iface_stop(handle);
+ return err;
}
}
- FARF(HIGH, "HMX %s (n_hmx=%d)", ctx->hmx_enabled ? "enabled" : "disabled", n_hmx);
- qurt_sysenv_max_hthreads_t hw_threads;
- qurt_sysenv_get_max_hw_threads(&hw_threads);
- uint32_t hw_nhvx = (qurt_hvx_get_units() >> 8) & 0xFF;
+ // DCVS setup
+ {
+ HAP_power_request_t request;
+ memset(&request, 0, sizeof(request));
- if (n_hvx == 0) {
- n_hvx = hw_nhvx;
+ request.type = HAP_power_set_DCVS_v3;
+ request.dcvs_v3.set_dcvs_enable = TRUE;
+ request.dcvs_v3.dcvs_enable = FALSE;
+ request.dcvs_v3.set_bus_params = TRUE;
+ request.dcvs_v3.bus_params.min_corner = HAP_DCVS_VCORNER_MAX;
+ request.dcvs_v3.bus_params.max_corner = HAP_DCVS_VCORNER_MAX;
+ request.dcvs_v3.bus_params.target_corner = HAP_DCVS_VCORNER_MAX;
+ request.dcvs_v3.set_core_params = TRUE;
+ request.dcvs_v3.core_params.min_corner = HAP_DCVS_VCORNER_MAX;
+ request.dcvs_v3.core_params.max_corner = HAP_DCVS_VCORNER_MAX;
+ request.dcvs_v3.core_params.target_corner = HAP_DCVS_VCORNER_MAX;
+ request.dcvs_v3.set_sleep_disable = TRUE;
+ request.dcvs_v3.sleep_disable = TRUE;
+
+#if (__HEXAGON_ARCH__ >= 79)
+ HAP_set_dcvs_v3_protected_bus_corners(&request, 1);
+#endif
+ if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
+ htp_iface_stop(handle);
+ return err;
+ }
+
+ memset(&request, 0, sizeof(request));
+ request.type = HAP_power_set_HVX;
+ request.hvx.power_up = TRUE;
+ if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
+ htp_iface_stop(handle);
+ return err;
+ }
}
- if (n_hvx > hw_threads.max_hthreads) {
- n_hvx = hw_threads.max_hthreads;
+
+#if __HVX_ARCH__ >= 75
+ {
+ // Power on HMX and set HMX clock
+ HAP_power_request_t request;
+ memset(&request, 0, sizeof(HAP_power_request_t));
+ request.type = HAP_power_set_HMX_v2;
+ request.hmx_v2.set_power = TRUE;
+ request.hmx_v2.power_up = TRUE;
+ request.hmx_v2.set_clock = TRUE;
+ request.hmx_v2.target_corner = HAP_DCVS_EXP_VCORNER_MAX;
+ request.hmx_v2.min_corner = HAP_DCVS_EXP_VCORNER_MAX;
+ request.hmx_v2.max_corner = HAP_DCVS_EXP_VCORNER_MAX;
+ request.hmx_v2.perf_mode = HAP_CLK_PERF_HIGH;
+ FARF(ALWAYS, "Setting HMX clock\n");
+ err = HAP_power_set((void *) ctx, &request);
+ if (err != AEE_SUCCESS) {
+ FARF(ERROR, "ggml-hex: error setting HMX clock.");
+ htp_iface_stop(handle);
+ return err;
+ }
}
- if (n_hvx > HTP_MAX_NTHREADS) {
- n_hvx = HTP_MAX_NTHREADS;
+#else
+ {
+ // Power on HMX
+ HAP_power_request_t request;
+ memset(&request, 0, sizeof(HAP_power_request_t));
+ request.type = HAP_power_set_HMX;
+ request.hmx.power_up = TRUE;
+ FARF(ALWAYS, "Powering HMX on\n");
+ err = HAP_power_set((void *) ctx, &request);
+ if (err != AEE_SUCCESS) {
+ FARF(ERROR, "ggml-hex: error powering on HMX.");
+ htp_iface_stop(handle);
+ return err;
+ }
}
+#endif
+
+ ctx->hmx_enabled = n_hmx;
+ ctx->hmx_queue = NULL;
+ if (n_hmx) {
+ void * hmx_ptr = (void *) ((uintptr_t) block + offset_hmx);
+ ctx->hmx_queue = hmx_queue_init(hmx_ptr, HMX_QUEUE_CAPACITY, HMX_QUEUE_STACK_SIZE, ctx->vtcm_rctx, &ctx->trace[HTP_MAX_NTHREADS]);
+ }
+ FARF(HIGH, "HMX %s (n_hmx=%d)", ctx->hmx_enabled ? "enabled" : "disabled", n_hmx);
ctx->n_threads = n_hvx;
+ ctx->n_threads_div = init_fastdiv_values(ctx->n_threads);
+
+ // Initialize DMA queues
+ uint8_t * dma_ptr_curr = (uint8_t *) ((uintptr_t) block + offset_dma);
+ size_t size_dma_q = dma_queue_sizeof(256);
+ size_t size_dma_alias = dma_queue_alias_sizeof();
+
for (int i = 0; i < ctx->n_threads; i++) {
- ctx->dma[i] = dma_queue_create(256); // queue depth
- if (ctx->dma[i]) {
- ctx->dma[i]->trace = &ctx->trace[i];
- }
+ dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
+ ctx->dma_cached[i] = dma_queue_init(dma_ptr_curr, 256, (uintptr_t) ctx->vtcm_base, ctx->vtcm_size, &ctx->trace[i]);
+ dma_ptr_curr += size_dma_q;
+
+ dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
+ ctx->dma[i] = dma_queue_alias_init(dma_ptr_curr, ctx->dma_cached[i], 1);
+ dma_ptr_curr += size_dma_alias;
}
ctx->ddr_spad_size = 512 * 1024; // 512 KB
ctx->ddr_spad_base = memalign(128, ctx->ddr_spad_size);
- // init worker pool
- err = worker_pool_init(&ctx->worker_pool, n_hvx);
- if (err != AEE_SUCCESS) {
- FARF(ERROR, "Unable to create worker pool");
- if (ctx->ddr_spad_base) {
- free(ctx->ddr_spad_base);
- ctx->ddr_spad_base = NULL;
- ctx->ddr_spad_size = 0;
+ void * wq_ptr = (void *) ((uintptr_t) block + offset_wq);
+ ctx->work_queue = work_queue_init(wq_ptr, n_hvx, WORK_QUEUE_CAPACITY, WORK_QUEUE_STACK_SIZE);
+
+ ctx->main_stack = NULL;
+ ctx->main_thread = 0;
+ atomic_store(&ctx->killed, false);
+
+ if (!use_callbacks) {
+ // Start main compute thread
+ ctx->main_stack = (void *) ((uintptr_t) block + offset_main_stack);
+
+ qurt_thread_attr_t attr;
+ qurt_thread_attr_init(&attr);
+ qurt_thread_attr_set_stack_addr(&attr, ctx->main_stack);
+ qurt_thread_attr_set_stack_size(&attr, size_main_stack);
+ qurt_thread_attr_set_priority(&attr, main_prio);
+ qurt_thread_attr_set_name(&attr, "htp-main");
+
+ int err_thread = qurt_thread_create(&ctx->main_thread, &attr, htp_main_thread, ctx);
+ if (err_thread) {
+ FARF(ERROR, "Unable to create htp main thread: %d", err_thread);
+ htp_iface_stop(handle);
+ return AEE_ENOMEMORY;
}
- return err;
}
FARF(HIGH, "session %u started: n-hvx %u vtcm-size %zu vtcm-rctx %u n-threads %u thread-id %d thread-prio %d \n",
}
AEEResult htp_iface_stop(remote_handle64 handle) {
- struct htp_context * ctx = (struct htp_context *) handle;
- if (!ctx) {
+ struct htp_handle * h = (struct htp_handle *) handle;
+ if (!h || !h->ctx) {
return AEE_EBADPARM;
}
+ struct htp_context * ctx = h->ctx;
- if (!ctx->queue) {
- FARF(ERROR, "Queue not open");
- return AEE_EBADSTATE;
+ if (ctx->main_thread) {
+ atomic_store(&ctx->killed, true);
+ int status;
+ (void) qurt_thread_join(ctx->main_thread, &status);
+ ctx->main_thread = 0;
}
- // Close queue. dspqueue_close() will also wait for callbacks to finish.
- int err = dspqueue_close(ctx->queue);
- ctx->queue = NULL;
+ int err = dspqueue_close(ctx->dsp_queue); ctx->dsp_queue = NULL;
if (err != 0) {
FARF(ERROR, "Queue close failed with 0x%08x", (unsigned) err);
return err;
}
- if (ctx->worker_pool) {
- // Release worker pool
- worker_pool_release(&ctx->worker_pool);
- }
+ work_queue_free(ctx->work_queue);
for (int i = 0; i < ctx->n_threads; i++) {
- dma_queue_delete(ctx->dma[i]);
+ dma_queue_alias_free(ctx->dma[i]);
+ dma_queue_free(ctx->dma_cached[i]);
}
if (ctx->hmx_queue) {
- hmx_queue_delete(ctx->hmx_queue);
+ hmx_queue_free(ctx->hmx_queue);
ctx->hmx_queue = NULL;
}
ctx->hmx_enabled = false;
ctx->ddr_spad_size = 0;
}
+ free(ctx);
+ h->ctx = NULL;
+
return AEE_SUCCESS;
}
case HTP_OP_RMS_NORM:
case HTP_OP_RMS_NORM_MUL:
case HTP_OP_SCALE:
+ case HTP_OP_CLAMP:
case HTP_OP_SQR:
case HTP_OP_SQRT:
case HTP_OP_UNARY_SOFTPLUS:
case HTP_OP_UNARY_SIGMOID:
+ case HTP_OP_UNARY_SILU:
+ case HTP_OP_UNARY_GELU:
case HTP_OP_UNARY_NEG:
case HTP_OP_UNARY_EXP:
case HTP_OP_UNARY_TANH:
case HTP_OP_L2_NORM:
return op_unary(octx);
- case HTP_OP_UNARY_SILU:
- case HTP_OP_UNARY_GELU:
case HTP_OP_GLU_SWIGLU:
case HTP_OP_GLU_SWIGLU_OAI:
case HTP_OP_GLU_GEGLU:
case HTP_OP_PAD:
return op_pad(octx);
+ case HTP_OP_IM2COL:
+ return op_im2col(octx);
+
case HTP_OP_CONCAT:
return op_concat(octx);
return op_gated_delta_net(octx);
case HTP_OP_TRI:
- return op_tri(octx);
+ return op_unary(octx);
case HTP_OP_INVALID:
break;
-
- // No default to catch missing cases
}
FARF(ERROR, "Unknown Op %u", octx->op);
}
}
-static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t idx, struct htp_tensor *t) {
+static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) {
uint32_t offset = t->data;
uint32_t size = t->size;
uint32_t bi = t->bi;
- t->data = bufs[bi].base + offset; // update data to the actual pointer
+ t->data = (uint32_t) (bufs[bi].base + offset); // update data to the actual pointer
FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t n_tens) {
for (uint32_t i=0; i < n_tens; i++) {
- prep_tensor(ctx, bufs, i, tens + i);
+ prep_tensor(ctx, bufs, tens, i, tens + i);
}
}
// Prep input tensors
for (uint32_t i=0; i<HTP_OP_MAX_INPUTS; i++) {
- struct htp_tensor *src = op->src[i] == 0xffff ? NULL : tens + op->src[i];
-
- octx->src[i] = src;
- if (!src) continue;
-
- if (!(src->flags & HTP_TENSOR_FLUSHED) && (src->flags & HTP_TENSOR_COMPUTE)) {
- // flush compute buffers on input
- hex_l2flush((void *) src->data, src->size);
+ uint16_t src_idx = op->src[i];
+ if (src_idx == 0xffff) {
+ octx->src[i] = NULL;
+ octx->src_dma[i] = NULL;
+ continue;
}
+ struct htp_tensor *src = tens + src_idx;
+ octx->src[i] = src;
+ octx->src_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
+
FARF(HIGH, "prep-src #%u: data %p size %u : %u:%u:%u:%u", op->src[i], (void*) src->data, src->size,
src->ne[0], src->ne[1], src->ne[3], src->ne[3]);
}
+ htp_tensor_flush_all(octx->ctx, octx->src, HTP_OP_MAX_INPUTS);
+
// Prep output tensors
for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) {
uint16_t dst_idx = op->dst[i];
if (dst_idx == 0xffff) {
- octx->dsts[i] = NULL;
+ octx->dsts[i] = NULL;
+ octx->dst_dma[i] = NULL;
continue;
}
struct htp_tensor *dst = tens + dst_idx;
- octx->dsts[i] = dst;
+ octx->dsts[i] = dst;
+ octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
int status = execute_op(octx);
+ htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+
octx->src0_spad.src = NULL;
octx->src1_spad.src = NULL;
octx->src2_spad.src = NULL;
octx->src3_spad.src = NULL;
octx->dst_spad.src = NULL;
- // flush buffers on output
- for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) {
- if (octx->dsts[i]) {
- struct htp_tensor *dst = (struct htp_tensor *)octx->dsts[i];
- hex_l2flush((void *) dst->data, dst->size);
- dst->flags |= HTP_TENSOR_FLUSHED;
+ return status;
+}
+
+static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_req * req, const struct dspqueue_buffer * dbuf) {
+ dspqueue_t queue = ctx->dsp_queue;
+ int err;
+
+ const uint32_t n_bufs = req->n_bufs;
+ const uint32_t n_tens = req->n_tensors;
+ const uint32_t n_ops = req->n_ops;
+
+ const uint32_t b_size = sizeof(struct htp_buf_desc) * n_bufs;
+ const uint32_t t_size = sizeof(struct htp_tensor) * n_tens;
+ const uint32_t o_size = sizeof(struct htp_op_desc) * n_ops;
+ const uint32_t p_size = sizeof(struct htp_prof_desc) * n_ops;
+ const uint32_t tr_size = (HTP_MAX_NTHREADS + 1) * req->n_traces * sizeof(struct htp_trace_desc);
- FARF(HIGH, "post-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, op->dst[i], (void*) dst->data, dst->size,
- dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
+ if (dbuf->size < b_size + t_size + o_size + p_size + tr_size) {
+ FARF(ERROR, "invalid opbatch memory block size %u (req %u)", dbuf->size, b_size + t_size + o_size + p_size + tr_size);
+ return;
+ }
+
+ FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
+ n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);
+
+ // Setup descriptor pointers
+ uint8_t * m_ptr = dbuf->ptr;
+ struct htp_buf_desc* bufs = (struct htp_buf_desc*) m_ptr; m_ptr += b_size;
+ struct htp_tensor* tens = (struct htp_tensor*) m_ptr; m_ptr += t_size;
+ struct htp_op_desc* ops = (struct htp_op_desc*) m_ptr; m_ptr += o_size;
+ struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;
+
+ struct profile_data batch_prof;
+ profile_start(HTP_PROF_BASIC, &batch_prof);
+
+ memset(ctx->trace, 0, sizeof(ctx->trace));
+ if (ctx->profiler == HTP_PROF_TRACE) {
+ struct htp_trace_desc * trace_events = (struct htp_trace_desc *) (m_ptr + p_size);
+ for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
+ ctx->trace[t].events = &trace_events[t * req->n_traces];
+ ctx->trace[t].max_events = req->n_traces;
}
}
- return status;
+ // Clean cache at the start of the batch
+ htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+ qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+ hex_l2fetch_block(ctx, ctx->footprint);
+ memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
+ htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+
+ htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
+ prep_op_bufs(ctx, bufs, n_bufs);
+ htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
+
+ prep_tensors(ctx, bufs, tens, n_tens);
+
+ struct htp_ops_context *octx = &ctx->octx;
+ memset(octx, 0, sizeof(*octx));
+ octx->n_threads = ctx->n_threads;
+ octx->ctx = ctx;
+
+ work_queue_wakeup(ctx->work_queue);
+ if (ctx->hmx_queue) {
+ hmx_queue_wakeup(ctx->hmx_queue);
+ }
+
+ int op_status = HTP_STATUS_OK;
+ for (uint32_t i = 0; i < n_ops && op_status == HTP_STATUS_OK; i++) {
+ struct profile_data prof;
+
+ profile_start(ctx->profiler, &prof);
+
+ op_status = proc_op_req(octx, tens, i, &ops[i]);
+
+ profile_stop(ctx->profiler, &prof);
+
+ if (ctx->profiler) {
+ pds[i].opcode = ops[i].opcode;
+ pds[i].usecs = prof.usecs;
+ pds[i].cycles_start = prof.cycles_start;
+ pds[i].cycles_stop = prof.cycles_stop;
+ for (int j = 0; j < HEX_NUM_PMU_COUNTERS; j++) {
+ pds[i].pmu[j] = prof.pmu_counters[j];
+ }
+ }
+ }
+
+ if (ctx->hmx_queue) {
+ hmx_queue_suspend(ctx->hmx_queue);
+ hmx_queue_flush(ctx->hmx_queue);
+ }
+ work_queue_suspend(ctx->work_queue);
+
+ // Flush remaining dirty tensors at the end of the batch
+ htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+ qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+ htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+
+ profile_stop(HTP_PROF_BASIC, &batch_prof);
+
+ struct htp_opbatch_rsp rsp;
+ memset(&rsp, 0, sizeof(rsp));
+ rsp.id = req->id;
+ rsp.status = op_status;
+ rsp.n_bufs = n_bufs;
+ rsp.n_tensors = n_tens;
+ rsp.n_ops = n_ops;
+ rsp.usecs = batch_prof.usecs;
+ rsp.cycles_start = batch_prof.cycles_start;
+ rsp.cycles_stop = batch_prof.cycles_stop;
+
+ if (ctx->profiler == HTP_PROF_TRACE) {
+ for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
+ rsp.n_traces[t] = ctx->trace[t].count;
+ }
+ }
+
+ struct dspqueue_buffer write_dbuf = *dbuf;
+ write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
+
+ err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
+ if (err != 0) {
+ FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
+ }
}
+#define DSPQUEUE_READ_TIMEOUT_USEC 5000
#define DSPQUEUE_POLL_TIMEOUT_USEC 100
#define DSPQUEUE_POLL_COUNT 100
-static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
- struct htp_context * ctx = (struct htp_context *) context;
-
+static void process_ops(struct htp_context * ctx) {
+ dspqueue_t queue = ctx->dsp_queue;
int err;
uint32_t poll_count = DSPQUEUE_POLL_COUNT;
vtcm_acquire(ctx);
- while (!ctx->vtcm_needs_release) {
+ while (!ctx->vtcm_needs_release && !atomic_load(&ctx->killed)) {
struct htp_opbatch_req req;
uint32_t r_size = sizeof(req);
// Reset poll count for valid requests
poll_count = DSPQUEUE_POLL_COUNT;
- const uint32_t n_bufs = req.n_bufs;
- const uint32_t n_tens = req.n_tensors;
- const uint32_t n_ops = req.n_ops;
-
- const uint32_t b_size = sizeof(struct htp_buf_desc) * n_bufs;
- const uint32_t t_size = sizeof(struct htp_tensor) * n_tens;
- const uint32_t o_size = sizeof(struct htp_op_desc) * n_ops;
- const uint32_t p_size = sizeof(struct htp_prof_desc) * n_ops;
- const uint32_t tr_size = (HTP_MAX_NTHREADS + 1) * req.n_traces * sizeof(struct htp_trace_desc);
-
- if (dbuf.size < b_size + t_size + o_size + p_size + tr_size) {
- FARF(ERROR, "invalid opbatch memory block size %u (req %u)", dbuf.size, b_size + t_size + o_size + p_size + tr_size);
- break;
- }
-
- FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req.id,
- n_bufs, n_tens, n_ops, req.n_traces, dbuf.size, b_size, t_size, o_size);
-
- // Setup descriptor pointers
- uint8_t * m_ptr = dbuf.ptr;
- struct htp_buf_desc* bufs = (struct htp_buf_desc*) m_ptr; m_ptr += b_size;
- struct htp_tensor* tens = (struct htp_tensor*) m_ptr; m_ptr += t_size;
- struct htp_op_desc* ops = (struct htp_op_desc*) m_ptr; m_ptr += o_size;
- struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;
-
- prep_op_bufs(ctx, bufs, n_bufs);
- prep_tensors(ctx, bufs, tens, n_tens);
-
- struct htp_ops_context *octx = &ctx->octx;
- memset(octx, 0, sizeof(*octx));
- octx->n_threads = ctx->n_threads;
- octx->ctx = ctx;
-
- if (ctx->profiler == HTP_PROF_TRACE) {
- memset(ctx->trace, 0, sizeof(ctx->trace));
- struct htp_trace_desc * trace_events = (struct htp_trace_desc *) (m_ptr + p_size);
- for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
- ctx->trace[t].events = &trace_events[t * req.n_traces];
- ctx->trace[t].max_events = req.n_traces;
- }
- } else {
- for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
- ctx->trace[t].events = NULL;
- ctx->trace[t].max_events = 0;
- }
- }
-
- int op_status = HTP_STATUS_OK;
- uint32_t op_wakeup = n_ops / 2; // half-way throgh the batch
-
- hmx_queue_wakeup(ctx->hmx_queue);
-
- for (uint32_t i=0; i < n_ops; i++) {
- struct profile_data prof;
-
- if (i == op_wakeup) {
- dspqueue_write_early_wakeup_noblock(queue, 0, 0);
- }
-
- profile_start(ctx->profiler, &prof);
+ process_opbatch(ctx, &req, &dbuf);
+ }
- op_status = proc_op_req(octx, tens, i, &ops[i]);
+ vtcm_release(ctx);
+}
- profile_stop(ctx->profiler, &prof);
+static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
+ (void) queue;
+ (void) error;
+ struct htp_handle * h = (struct htp_handle *) context;
+ if (h && h->ctx) {
+ process_ops(h->ctx);
+ }
+}
- if (op_status != HTP_STATUS_OK) {
- break;
- }
+static void htp_main_thread(void * context) {
+ struct htp_context * ctx = (struct htp_context *) context;
- if (ctx->profiler) {
- pds[i].opcode = ops[i].opcode;
- pds[i].usecs = prof.usecs;
- pds[i].cycles_start = prof.cycles_start;
- pds[i].cycles_stop = prof.cycles_stop;
- for (int j = 0; j < HEX_NUM_PMU_COUNTERS; j++) {
- pds[i].pmu[j] = prof.pmu_counters[j];
- }
- }
- }
+ FARF(HIGH, "htp-main-thread: started");
- hmx_queue_suspend(ctx->hmx_queue);
+ while (!atomic_load(&ctx->killed)) {
+ uint32_t flags = 0;
+ uint32_t num_buffers = 0;
+ uint32_t message_length = 0;
- struct htp_opbatch_rsp rsp;
- rsp.id = req.id;
- rsp.status = op_status;
- rsp.n_bufs = n_bufs;
- rsp.n_tensors = n_tens;
- rsp.n_ops = n_ops;
- memset(rsp.pad, 0, sizeof(rsp.pad));
- if (ctx->profiler == HTP_PROF_TRACE) {
- for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
- rsp.n_traces[t] = ctx->trace[t].count;
- }
+ int err = dspqueue_peek(ctx->dsp_queue, &flags, &num_buffers, &message_length, 50000);
+ if (err == 0) {
+ process_ops(ctx);
+ } else if (err == AEE_EWOULDBLOCK || err == AEE_EEXPIRED) {
+ continue;
} else {
- memset(rsp.n_traces, 0, sizeof(rsp.n_traces));
- }
-
- dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
-
- err = dspqueue_write(queue, 0, 1, &dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
- if (err != 0) {
- FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
+ FARF(ERROR, "dspqueue_peek failed: 0x%08x", (unsigned) err);
break;
}
}
- vtcm_release(ctx);
+ FARF(HIGH, "htp-main-thread: stopped");
}
#include "hex-dma.h"
#include "hvx-utils.h"
#include "hvx-dump.h"
+#include "hvx-arith.h"
+#include "hvx-reduce.h"
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
// Precomputed values
uint32_t src0_nrows_per_thread;
+ uint32_t src0_row_size_padded;
+ uint32_t src1_nrows;
struct fastdiv_values mm_div_ne12_ne1;
struct fastdiv_values mm_div_ne1;
// Per thread quant tasks
// Precomputed block-parallel quantization values
worker_callback_t quant_task_func;
- uint32_t quant_ib_first[MAX_NUM_WORKERS];
- uint32_t quant_ib_last[MAX_NUM_WORKERS];
- uint32_t quant_r[MAX_NUM_WORKERS];
- uint32_t quant_c[MAX_NUM_WORKERS];
+ uint32_t quant_ib_first[WORK_QUEUE_MAX_N_THREADS];
+ uint32_t quant_ib_last[WORK_QUEUE_MAX_N_THREADS];
+ uint32_t quant_r[WORK_QUEUE_MAX_N_THREADS];
+ uint32_t quant_c[WORK_QUEUE_MAX_N_THREADS];
uint32_t n_quant_tasks;
uint32_t n_quant_rows_per_thread;
atomic_uint quant_barrier;
// Fields for scattered mapping & HMX support in MUL_MAT_ID
const uint32_t * matrix_row_counts;
const struct mmid_row_mapping * matrix_rows;
+ uint32_t mapping_stride;
// Dynamic VTCM pointers allocated sequentially
uint8_t * vtcm_src0;
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
};
-
-
#define htp_matmul_tensors_preamble \
const struct htp_tensor * restrict src0 = octx->src[0]; \
const struct htp_tensor * restrict src1 = octx->src[1]; \
return;
}
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start);
const uint32_t blck_0 = 64;
const uint32_t src0_start_row = src0_nrows_per_thread * ith; \
const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); \
\
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
\
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \
const uint32_t n_prefetch = kparams->n_prefetch; \
const uint32_t src0_start_row = src0_nrows_per_thread * ith; \
const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); \
\
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
\
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \
const uint32_t n_prefetch = kparams->n_prefetch; \
\
uint32_t push_ct = ct_start; \
if (src0_start_row < src0_end_row) { \
+ if (src2) { \
+ float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; \
+ const float * src2_ptr = (const float *) src2->data + src0_start_row; \
+ int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
+ if (slice_size > 0) { \
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), \
+ slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \
+ dma_queue_pop_nowait(dma_queue); \
+ } \
+ } \
for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \
src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \
\
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\
if (push_ct < ct_end) { \
dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \
\
int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
if (copy_cnt > 0) { \
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \
if (src2) { \
- float * dst_ptr = &dst_col[src0_start_row]; \
- const float * src2_ptr = (const float *) src2->data + src0_start_row; \
- float * tmp_ptr = tmp; \
- int remaining = copy_cnt; \
- while (remaining > 0) { \
- int n = MIN(remaining, 32); \
- HVX_Vector v_out = hvx_vmemu(tmp_ptr); \
- HVX_Vector v_z = hvx_vmemu(src2_ptr); \
- hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z)); \
- dst_ptr += n; \
- src2_ptr += n; \
- tmp_ptr += n; \
- remaining -= n; \
- } \
+ hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], \
+ (const uint8_t *) tmp, \
+ (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), \
+ copy_cnt); \
} else { \
hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \
} \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \
} \
}
uint8_t * restrict vtcm_src3_ptr = mmctx->vtcm_src3 + mmctx->vtcm_src3_size_per_thread * ith; \
uint8_t * restrict src1_data = mmctx->vtcm_src1; \
\
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
\
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \
const uint32_t n_prefetch = kparams->n_prefetch; \
uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; \
uint8_t * restrict src1_data = mmctx->vtcm_src1; \
\
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
\
const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \
const uint8_t * restrict src2_row = (const uint8_t *) src2->data; \
return; \
} \
\
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \
\
uint8_t * restrict dst = mmctx->vtcm_src1; \
static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, void * data) {
struct htp_mm_context * mmctx = data;
struct htp_ops_context * octx = mmctx->octx;
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]);
const struct htp_tensor * src = octx->src[1];
static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, void * data) {
struct htp_mm_context * mmctx = data;
struct htp_ops_context * octx = mmctx->octx;
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]);
const struct htp_tensor * src = octx->src[1];
const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U);
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
const size_t dst_row_size = nb1;
const size_t src0_row_size = nb01;
const uint32_t src0_start_row = src0_nrows_per_thread * ith;
const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
const size_t dst_row_size = nb1;
const size_t src0_row_size = nb01;
// Prefill vtcm with 2x src0 rows
if (src0_start_row < src0_end_row) {
+ if (src2) {
+ float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row;
+ const float * src2_ptr = (const float *) src2->data + src0_start_row;
+ int slice_size = (int)src0_end_row - (int)src0_start_row;
+ if (slice_size > 0) {
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr),
+ slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1);
+ dma_queue_pop_nowait(dma_queue);
+ }
+ }
for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
const uint32_t is0 = (ir0 - src0_start_row);
if (is0 >= n_prefetch) {
}
int copy_cnt = src0_end_row - src0_start_row;
- if (src2) {
- float * dst_ptr = &dst_col[src0_start_row];
- const float * src2_ptr = (const float *) src2->data + src0_start_row;
- float * tmp_ptr = tmp;
- int remaining = copy_cnt;
- while (remaining > 0) {
- int n = MIN(remaining, 32);
- HVX_Vector v_out = hvx_vmemu(tmp_ptr);
- HVX_Vector v_z = hvx_vmemu(src2_ptr);
- hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
- dst_ptr += n;
- src2_ptr += n;
- tmp_ptr += n;
- remaining -= n;
+ if (copy_cnt > 0) {
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
+ if (src2) {
+ hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],
+ (const uint8_t *) tmp,
+ (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row),
+ copy_cnt);
+ } else {
+ hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
}
- } else {
- hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
}
}
-#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ids->ne[0] * ids->ne[1] + (i1)]
+#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)]
static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {
htp_matmul_preamble;
return;
}
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
const uint32_t n_prefetch = kparams->n_prefetch;
return;
}
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
const uint32_t n_prefetch = kparams->n_prefetch;
static int hvx_mm_matmul(struct htp_ops_context * octx) {
htp_matmul_tensors_preamble;
+ struct htp_thread_trace * tr = &octx->ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
struct htp_mm_context mmctx_struct = {0};
struct htp_mm_context * mmctx = &mmctx_struct;
mmctx->octx = octx;
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
- dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false);
+ dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false, false);
if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM ||
kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+ mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL;
mmctx->vtcm_src0_stride = src0_row_size_padded;
mmctx->vtcm_src1_stride = src1_row_size;
- if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
- return HTP_STATUS_OK;
-
if (need_quant) {
mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
mmctx->quant_task_func = quant_task_func;
mmctx->n_quant_tasks = 0;
}
- const uint32_t n_matmul_jobs = octx->n_threads;
- worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
+ worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads);
return HTP_STATUS_OK;
}
#define DEQUANTIZE_WORKER_LOOP_IMPL(SUFFIX) \
static void dequantize_tiled_worker_loop_##SUFFIX(unsigned int n, unsigned int i, void *data) { \
tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; \
- struct htp_thread_trace * tr = state->traces ? &state->traces[i] : NULL; \
+ struct htp_thread_trace * tr = &state->traces[i]; \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); \
for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { \
int start = task_id * state->n_tiles_per_task; \
static void convert_f16_worker_loop(unsigned int n, unsigned int i, void *data) {
tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data;
- struct htp_thread_trace * tr = state->traces ? &state->traces[i] : NULL;
+ struct htp_thread_trace * tr = &state->traces[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i);
for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) {
int start = task_id * state->n_tiles_per_task;
static void quantize_f32_worker_loop(unsigned int n, unsigned int i, void *data) {
tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data;
- struct htp_thread_trace * tr = state->traces ? &state->traces[i] : NULL;
+ struct htp_thread_trace * tr = &state->traces[i];
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, i);
for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) {
static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
output_transfer_task_state_t *st = (output_transfer_task_state_t *) data;
- struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+ struct htp_thread_trace * tr = &st->traces[i];
int start_chunk_idx = i * st->n_chunks_per_task;
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx);
uint32_t dma_step_rows_shift;
} activation_transfer_task_state_t;
+typedef struct {
+ __fp16 *dst;
+ const float *src;
+ uint32_t n_rows;
+ uint32_t k_block;
+ uint32_t k_stride;
+ uint32_t k_valid;
+ uint32_t n_col_chunks;
+ struct fastdiv_values n_threads_div;
+ float *vtcm_f32_act;
+ size_t vtcm_f32_act_bytes;
+ struct htp_thread_trace *traces;
+ struct htp_context *ctx;
+ uint32_t dma_step_rows;
+ uint32_t dma_step_rows_shift;
+} activation_transfer_col_chunk_state_t;
+
+static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk(
+ dma_queue *dma_q,
+ __fp16 *restrict vtcm_dst,
+ const float *restrict src,
+ uint32_t n_rows,
+ uint32_t k_block,
+ uint32_t k_stride,
+ uint32_t k_chunk_valid,
+ uint32_t c_first,
+ uint32_t c_len,
+ float *thread_f32_act,
+ struct htp_thread_trace *tr,
+ uint32_t dma_step_rows,
+ uint32_t dma_step_rows_shift) {
+
+ const uint32_t R = dma_step_rows;
+ const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+
+ const uint32_t n_steps = n_rows_padded >> dma_step_rows_shift;
+
+ // Push step 0
+ if (n_steps > 0 && n_rows > 0) {
+ uint32_t nrows_to_fetch = hex_smin(n_rows, R);
+ dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src + c_first),
+ c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch);
+ }
+ // Push step 1
+ if (n_steps > 1) {
+ uint32_t next_r = R * 1;
+ if (next_r < n_rows) {
+ uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
+ const float *next_src = src + next_r * k_stride + c_first;
+ float *next_buf = thread_f32_act + 1 * R * c_len;
+ dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src),
+ c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch);
+ }
+ }
+ for (uint32_t s = 0; s < n_steps; ++s) {
+ uint32_t r = s << dma_step_rows_shift;
+ float *curr_buf = thread_f32_act;
+
+ if (r < n_rows) {
+ curr_buf = (float *) dma_queue_pop(dma_q).dst;
+ }
+
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r);
+ for (uint32_t p = 0; p < (R >> 1); ++p) {
+ uint32_t row_idx = r + (p << 1);
+ float *pair_buf = curr_buf + (p << 1) * c_len;
+ bool r0_valid = ((row_idx + 0) < n_rows);
+ bool r1_valid = ((row_idx + 1) < n_rows);
+
+ transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+ vtcm_dst, pair_buf, pair_buf + c_len, row_idx, k_block, c_first, c_len, k_chunk_valid, r0_valid, r1_valid
+ );
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r);
+
+ // Push step s + 2
+ uint32_t next_s = s + 2;
+ uint32_t next_r = next_s << dma_step_rows_shift;
+ if (next_r < n_rows) {
+ uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
+ const float *next_src = src + next_r * k_stride + c_first;
+ dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src),
+ c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch);
+ }
+ }
+}
+
+static void transfer_activation_chunk_fp32_to_fp16_col_chunk(
+ __fp16 *restrict vtcm_dst,
+ const float *restrict src,
+ uint32_t n_rows,
+ uint32_t k_block,
+ uint32_t k_stride,
+ uint32_t c_first,
+ uint32_t c_len,
+ uint32_t k_chunk_valid) {
+ const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+ const uint32_t n_rows_tiled = (n_rows / HTP_MM_HMX_TILE_N_ROWS) * HTP_MM_HMX_TILE_N_ROWS;
+
+ uint32_t r = 0;
+
+ #pragma unroll(2)
+ for (r = 0; r < n_rows_tiled; r += 2) {
+ const float *ptr_in0 = src + (r + 0) * k_stride + c_first;
+ const float *ptr_in1 = src + (r + 1) * k_stride + c_first;
+
+ transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+ vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, true, true
+ );
+ }
+
+ for (; r < n_rows_padded; r += 2) {
+ const bool row0_valid = r < n_rows;
+ const bool row1_valid = (r + 1) < n_rows;
+
+ const float *ptr_in0 = row0_valid ? (src + (r + 0) * k_stride + c_first) : NULL;
+ const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride + c_first) : NULL;
+
+ transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+ vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, row0_valid, row1_valid
+ );
+ }
+}
+
+static void transfer_activation_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
+ activation_transfer_col_chunk_state_t *st = (activation_transfer_col_chunk_state_t *) data;
+ struct htp_thread_trace * tr = &st->traces[i];
+
+ uint32_t n_blocks = st->k_block / 32;
+ uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div);
+ uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div);
+ uint32_t c_first = b_first * 32;
+ uint32_t c_last = b_last * 32;
+ uint32_t c_len = c_last - c_first;
+
+ if (c_len == 0) {
+ return;
+ }
+
+ uint32_t k_chunk_valid = 0;
+ if (st->k_valid > c_first) {
+ k_chunk_valid = hex_smin(st->k_valid, c_last) - c_first;
+ }
+
+ __fp16 *dst = st->dst;
+ const float *src = st->src;
+
+ if (st->vtcm_f32_act) {
+ size_t thread_scratch_bytes = hex_align_down(fastdiv(st->vtcm_f32_act_bytes, &st->n_threads_div), 128);
+ float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * thread_scratch_bytes);
+
+ transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk(
+ st->ctx->dma[i], dst, src, st->n_rows, st->k_block, st->k_stride, k_chunk_valid,
+ c_first, c_len, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift
+ );
+ } else {
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first);
+ transfer_activation_chunk_fp32_to_fp16_col_chunk(
+ dst, src, st->n_rows, st->k_block, st->k_stride, c_first, c_len, k_chunk_valid
+ );
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first);
+ }
+}
+
static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
dma_queue *dma_q,
__fp16 *restrict vtcm_dst,
static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
activation_transfer_task_state_t *st = (activation_transfer_task_state_t *) data;
- struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+ struct htp_thread_trace * tr = &st->traces[i];
for (unsigned int task_id = i; task_id < (unsigned int)st->n_tasks; task_id += n) {
int chunk_idx = task_id * st->n_chunks_per_task;
static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) {
activation_transfer_gathered_task_state_t *st = data;
- struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+ struct htp_thread_trace * tr = &st->traces[i];
int chunk_idx = i;
int chunk_size = st->n_chunks_per_task;
- int start_row = st->start_row + chunk_idx * chunk_size;
+ int vtcm_start_row = chunk_idx * chunk_size;
+ int start_row = st->start_row + vtcm_start_row;
int n_rows = hex_smin(st->cne1 - start_row, chunk_size);
if (n_rows > 0) {
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
transfer_activation_chunk_fp32_to_fp16_gathered(
- st->dst, st->src, start_row, n_rows, st->k_block,
+ st->dst, st->src, start_row, vtcm_start_row, n_rows, st->k_block,
st->matrix_rows, st->cur_a, st->mapping_stride,
st->ne11, &st->ne11_div, st->nb11, st->nb12, st->cne1, st->k_valid);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
static void transfer_activation_chunk_gathered_worker_flat_fn(unsigned int n, unsigned int i, void *data) {
activation_transfer_gathered_task_state_t *st = data;
- struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+ struct htp_thread_trace * tr = &st->traces[i];
int chunk_idx = i;
int chunk_size = st->n_chunks_per_task;
- int start_row = st->start_row + chunk_idx * chunk_size;
+ int vtcm_start_row = chunk_idx * chunk_size;
+ int start_row = st->start_row + vtcm_start_row;
int n_rows = hex_smin(st->cne1 - start_row, chunk_size);
if (n_rows > 0) {
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
transfer_activation_chunk_fp32_to_fp16_gathered_flat(
- st->dst, st->src, start_row, n_rows, st->k_block,
+ st->dst, st->src, start_row, vtcm_start_row, n_rows, st->k_block,
st->matrix_rows, st->cur_a, st->mapping_stride,
st->nb12, st->cne1, st->k_valid);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
static void transfer_output_chunk_scattered_worker_fn(unsigned int n, unsigned int i, void *data) {
output_transfer_scattered_task_state_t *st = data;
- struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+ struct htp_thread_trace * tr = &st->traces[i];
int chunk_idx = i;
int chunk_size = st->n_chunks_per_task;
- int start_row = st->start_row + chunk_idx * chunk_size;
+ int vtcm_start_row = chunk_idx * chunk_size;
+ int start_row = st->start_row + vtcm_start_row;
int n_rows = hex_smin(st->cne1 - start_row, chunk_size);
if (n_rows > 0) {
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, chunk_idx);
transfer_output_chunk_fp16_to_fp32_scattered(
- st->dst, st->vtcm_src, start_row, n_rows, st->n_cols,
+ st->dst, st->vtcm_src, start_row, vtcm_start_row, n_rows, st->n_cols,
st->matrix_rows, st->cur_a, st->mapping_stride,
st->dst_nb1, st->dst_nb2, st->cne1);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, chunk_idx);
}
}
+typedef struct {
+ float *dst;
+ const float *src2;
+ const __fp16 *vtcm_src;
+ uint32_t n_rows;
+ uint32_t n_cols;
+ uint32_t dst_stride;
+ uint32_t src2_stride;
+ uint32_t dst_cols;
+ struct fastdiv_values n_threads_div;
+ struct htp_thread_trace *traces;
+ struct htp_context *ctx;
+} output_transfer_col_chunk_state_t;
+
+static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
+ (void) n;
+ output_transfer_col_chunk_state_t *st = (output_transfer_col_chunk_state_t *) data;
+ struct htp_thread_trace * tr = &st->traces[i];
+
+ uint32_t n_blocks = st->n_cols / 32;
+ uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div);
+ uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div);
+ uint32_t c_first = b_first * 32;
+ uint32_t c_last = b_last * 32;
+ uint32_t c_len = c_last - c_first;
+
+ if (c_len == 0) return;
+
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
+
+ float *dst = st->dst + c_first;
+ const float *src2 = st->src2 ? (st->src2 + c_first) : NULL;
+ const __fp16 *vtcm_src = st->vtcm_src + b_first * HTP_MM_HMX_TILE_N_ELMS;
+
+ int chunk_dst_cols = (int)st->dst_cols - (int)c_first;
+ if (chunk_dst_cols > 0) {
+ transfer_output_chunk_fp16_to_fp32_col_chunk(
+ dst, src2, vtcm_src, 0, st->n_rows, c_len, st->n_cols,
+ st->dst_stride, st->src2_stride, (uint32_t)chunk_dst_cols
+ );
+ }
+
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
+}
+
static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const float *src2, const __fp16 *vtcm_src,
int n_rows, int n_cols, int dst_stride, uint32_t src2_stride, int dst_cols, int n_threads) {
assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
if (n_rows <= 0) return;
+ uint32_t n_blocks = (uint32_t)n_cols / 32;
+ if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) {
+ struct fastdiv_values n_threads_div = init_fastdiv_values(n_threads);
+ output_transfer_col_chunk_state_t col_state;
+ col_state.dst = dst;
+ col_state.src2 = src2;
+ col_state.vtcm_src = vtcm_src;
+ col_state.n_rows = (uint32_t)n_rows;
+ col_state.n_cols = (uint32_t)n_cols;
+ col_state.dst_stride = (uint32_t)dst_stride;
+ col_state.src2_stride = src2_stride;
+ col_state.dst_cols = (uint32_t)dst_cols;
+ col_state.n_threads_div = n_threads_div;
+ col_state.traces = ctx->trace;
+ col_state.ctx = ctx;
+
+ worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_col_chunk_worker_fn, &col_state, n_threads);
+ return;
+ }
+
size_t n_tot_chunks = n_rows;
size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : hmx_ceil_div(n_rows, n_threads);
n_chunks_per_task = hex_align_up(n_chunks_per_task, 2);
}
}
-static void transfer_activation_chunk_threaded(
- struct htp_context *ctx,
- __fp16 *dst,
- const float *src,
- int n_rows,
- int k_block,
- int k_stride,
- int n_threads,
- int k_valid,
- float *vtcm_f32_act,
- size_t vtcm_f32_act_bytes) {
+struct activation_transfer_params {
+ struct htp_context * ctx;
+ __fp16 * dst;
+ const float * src;
+ int n_rows;
+ int k_block;
+ int k_stride;
+ int n_threads;
+ const struct fastdiv_values * act_threads_div;
+ const struct fastdiv_values * k_div;
+ int k_valid;
+ float * vtcm_f32_act;
+ size_t vtcm_f32_act_bytes;
+};
+
+static void transfer_activation_chunk_threaded(const struct activation_transfer_params * params) {
+ struct htp_context * ctx = params->ctx;
+ __fp16 * dst = params->dst;
+ const float * src = params->src;
+ int n_rows = params->n_rows;
+ int k_block = params->k_block;
+ int k_stride = params->k_stride;
+ int n_threads = params->n_threads;
+ const struct fastdiv_values * act_threads_div = params->act_threads_div;
+ const struct fastdiv_values * k_div = params->k_div;
+ int k_valid = params->k_valid;
+ float * vtcm_f32_act = params->vtcm_f32_act;
+ size_t vtcm_f32_act_bytes = params->vtcm_f32_act_bytes;
+
if (n_rows <= 0) {
return;
}
+ const size_t n_tasks = (n_rows + 31) >> 5;
+ if (n_threads > 1 && k_block > 32 && n_tasks < (size_t) n_threads) {
+ // Calculate step rows parameters for column-chunked dma pipelining
+ uint32_t dma_step_rows = 2;
+ uint32_t dma_step_rows_shift = 1;
+ if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) {
+ size_t thread_scratch_bytes = hex_align_down(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128);
+ size_t thread_scratch_elements = thread_scratch_bytes / sizeof(float);
+ size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div);
+ if (dma_step_rows_max >= 4) {
+ dma_step_rows = 4;
+ dma_step_rows_shift = 2;
+ }
+ }
+
+ activation_transfer_col_chunk_state_t col_state;
+ col_state.dst = dst;
+ col_state.src = src;
+ col_state.n_rows = n_rows;
+ col_state.k_block = k_block;
+ col_state.k_stride = k_stride;
+ col_state.k_valid = k_valid;
+ col_state.n_col_chunks = n_threads;
+ col_state.n_threads_div = *act_threads_div;
+ col_state.vtcm_f32_act = vtcm_f32_act;
+ col_state.vtcm_f32_act_bytes = vtcm_f32_act_bytes;
+ col_state.traces = ctx->trace;
+ col_state.ctx = ctx;
+ col_state.dma_step_rows = dma_step_rows;
+ col_state.dma_step_rows_shift = dma_step_rows_shift;
+
+ worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_col_chunk_worker_fn, &col_state, n_threads);
+ return;
+ }
+
assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0 && k_stride % HTP_MM_HMX_TILE_N_COLS == 0);
size_t n_tot_chunks = n_rows;
size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : 32; // must be multiple of 32 to ensure correct destination address
- uint32_t dma_step_rows = 2;
- uint32_t dma_step_rows_shift = 1;
- if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) {
- size_t thread_scratch_elements = vtcm_f32_act_bytes / (n_threads * sizeof(float));
- size_t dma_step_rows_max = (thread_scratch_elements / 2) / k_block;
- if (dma_step_rows_max >= 4) {
- dma_step_rows = 4;
- dma_step_rows_shift = 2;
- } else {
- dma_step_rows = 2;
- dma_step_rows_shift = 1;
- }
- }
-
activation_transfer_task_state_t state;
- state.n_tasks = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task;
+ state.n_tasks = (n_threads == 1) ? 1 : hmx_ceil_div(n_tot_chunks, 32);
state.n_tot_chunks = n_tot_chunks;
state.n_chunks_per_task = n_chunks_per_task;
state.dst = dst;
state.vtcm_f32_act = vtcm_f32_act;
int active_threads = hex_smin(n_threads, (int)state.n_tasks);
- state.vtcm_f32_act_bytes_per_thread = (vtcm_f32_act_bytes / active_threads) & ~127u;
+ state.vtcm_f32_act_bytes_per_thread = hex_align_down(vtcm_f32_act_bytes / active_threads, 128);
+
+ uint32_t dma_step_rows = 2;
+ uint32_t dma_step_rows_shift = 1;
+ if (vtcm_f32_act && state.vtcm_f32_act_bytes_per_thread > 0 && k_block > 0) {
+ size_t thread_scratch_elements = state.vtcm_f32_act_bytes_per_thread / sizeof(float);
+ size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div);
+ if (dma_step_rows_max >= 4) {
+ dma_step_rows = 4;
+ dma_step_rows_shift = 2;
+ }
+ }
state.dma_step_rows = dma_step_rows;
state.dma_step_rows_shift = dma_step_rows_shift;
int pipeline,
int n_threads,
int act_threads,
+ const struct fastdiv_values * act_threads_div,
+ const struct fastdiv_values * k_div,
int tile_size,
int aligned_tile_size,
int vtcm_size) {
+ struct htp_thread_trace * tr = &ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
if (k % 32 != 0 || n % 32 != 0) { return -1; }
if (!hex_is_aligned(dst, VLEN) || !hex_is_aligned(activation, VLEN)) { return -1; }
int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
if (pipeline) {
// --- Asynchronous Pipelined Loop ---
hmx_matmul_job_t job_slots[2]; // persistent double-buffered job descriptors
void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 };
void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 };
- transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes);
+ struct activation_transfer_params act_params = {
+ .ctx = ctx,
+ .dst = vtcm_f16_act,
+ .src = activation + mr * act_stride,
+ .n_rows = (int) n_rows,
+ .k_block = k,
+ .k_stride = act_stride,
+ .n_threads = act_threads,
+ .act_threads_div = act_threads_div,
+ .k_div = k_div,
+ .k_valid = k_valid,
+ .vtcm_f32_act = vtcm_f32_act,
+ .vtcm_f32_act_bytes = L.act_f32_bytes,
+ };
+ transfer_activation_chunk_threaded(&act_params);
// Prologue: push A0 and optionally A1 (if n_chunk_cnt > 1)
const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols);
for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) {
const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
- transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes);
+ struct activation_transfer_params act_params = {
+ .ctx = ctx,
+ .dst = vtcm_f16_act,
+ .src = activation + mr * act_stride,
+ .n_rows = (int) n_rows,
+ .k_block = k,
+ .k_stride = act_stride,
+ .n_threads = act_threads,
+ .act_threads_div = act_threads_div,
+ .k_div = k_div,
+ .k_valid = k_valid,
+ .vtcm_f32_act = vtcm_f32_act,
+ .vtcm_f32_act_bytes = L.act_f32_bytes,
+ };
+ transfer_activation_chunk_threaded(&act_params);
// A0: Pre-fetch the first weight chunk (nc = 0)
if (n > 0) {
static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx,
const hmx_mm_f16_f32_batched_params_t *params,
- int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size) {
+ int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size,
+ const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div) {
int ret = 0;
for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) {
for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) {
params->act_stride, params->weight_stride * (int)sizeof(__fp16),
HTP_TYPE_F16, params->k, params->dst_stride, params->src2_stride, params->n,
m_chunk, n_chunk, pipeline, n_threads, act_threads,
- 0, 0, vtcm_size);
+ act_threads_div, k_div, 0, 0, vtcm_size);
}
}
return ret;
}
static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_batched_params_t *params,
- int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size) {
+ int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads,
+ const struct fastdiv_values * act_threads_div,
+ const struct fastdiv_values * k_div,
+ int vtcm_size) {
if (params->act_stride < params->k || params->weight_stride < params->k || params->dst_stride < params->n) { return -1; }
if (params->ne02 <= 0 || params->ne03 <= 0 || params->ne12 <= 0 || params->ne13 <= 0) { return -1; }
if (params->ne12 % params->ne02 != 0 || params->ne13 % params->ne03 != 0) { return -1; }
// Grouped path is only valid if group_size > 1 and it fits within VTCM budget.
bool run_grouped = (group_size > 1 && (size_t)vtcm_size <= vtcm_budget);
if (!run_grouped) {
- return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size);
+ return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div);
}
+ struct htp_thread_trace * tr = &ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
const size_t vec_dot_size = params->k * sizeof(__fp16);
const bool use_dma_activation = (params->act_stride > params->k);
if (L.total_bytes > vtcm_budget) {
FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__);
- return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+ return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div);
}
uint8_t * const base = (uint8_t *) ctx->vtcm_base;
const size_t fp16_row_bytes = (size_t) params->k * sizeof(__fp16);
const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
hmx_matmul_job_t job;
for (int b3 = 0; b3 < params->ne13; ++b3) {
for (int g = 0; g < group_size; ++g) {
const float *activation_chunk = hmx_mm_activation_batch_ptr(params, b2_base + g, b3) + mr * params->act_stride;
__fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride;
- transfer_activation_chunk_threaded(ctx, vtcm_act_g,
- activation_chunk, (int) n_rows,
- params->k, params->act_stride, act_threads, params->k, vtcm_f32_act, L.act_f32_bytes);
+ struct activation_transfer_params act_params = {
+ .ctx = ctx,
+ .dst = vtcm_act_g,
+ .src = activation_chunk,
+ .n_rows = (int) n_rows,
+ .k_block = params->k,
+ .k_stride = params->act_stride,
+ .n_threads = act_threads,
+ .act_threads_div = act_threads_div,
+ .k_div = k_div,
+ .k_valid = params->k,
+ .vtcm_f32_act = vtcm_f32_act,
+ .vtcm_f32_act_bytes = L.act_f32_bytes,
+ };
+ transfer_activation_chunk_threaded(&act_params);
}
// Prologue: Push A0 and A1 (if exists)
const struct mmid_row_mapping *matrix_rows,
int cur_a,
int mapping_stride) {
+ struct htp_thread_trace * tr = &ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
const int cne1 = m;
const int m_padded = hex_align_up(m, 32);
hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
hmx_matmul_job_t job;
for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) {
const int act_stride = (int)(src1->nb[1] / sizeof(float));
const int wgt_stride = (int)(src0->nb[1] / sizeof(__fp16));
- if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
- return HTP_STATUS_OK;
- }
-
const float * src2_ptr = NULL;
uint32_t src2_stride = 0;
size_t src2_nb2 = 0;
kparams->m_chunk, kparams->n_chunk,
kparams->pipeline, n_threads,
kparams->n_act_threads,
+ &kparams->div_n_act_threads,
+ &kparams->div_ne00_padded,
kparams->vtcm_size);
} else {
ret = hmx_mm_2d_f32(
(int)(dst->nb[1] / sizeof(float)), src2_stride, (int)dst->ne[0],
kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads,
kparams->n_act_threads,
+ &kparams->div_n_act_threads,
+ &kparams->div_ne00_padded,
kparams->tile_size, kparams->aligned_tile_size, kparams->vtcm_size
);
}
static int hmx_mm_op_matmul_id(
struct htp_ops_context * octx,
- struct htp_mm_context * mmctx,
- const uint32_t * matrix_row_counts,
- const struct mmid_row_mapping * matrix_rows,
- void * mapping_buf,
- bool must_free_mapping
+ struct htp_mm_context * mmctx
) {
+ const uint32_t * matrix_row_counts = mmctx->matrix_row_counts;
+ const struct mmid_row_mapping * matrix_rows = mmctx->matrix_rows;
htp_matmul_tensors_preamble;
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
const int n_ids = octx->src[2]->ne[0];
nb11, nb12,
nb1, nb2,
(int) src0->nb[1], (int) src0->type,
- matrix_rows, cur_a, n_ids * octx->src[2]->ne[1]);
+ matrix_rows, cur_a, mmctx->mapping_stride);
if (ret != 0) {
FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret);
- if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_NO_SUPPORT;
}
}
- if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_OK;
}
static int hvx_mm_matmul_id(
struct htp_ops_context * octx,
struct htp_mm_context * mmctx,
- size_t src0_row_size_padded,
- uint32_t src1_nrows,
- worker_callback_t matmul_id_job_func,
- void * mapping_buf,
- bool must_free_mapping
+ work_queue_func_t hvx_mmid_task_func
) {
htp_matmul_tensors_preamble;
+ const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded;
+ const uint32_t src1_nrows = mmctx->src1_nrows;
+
+ struct htp_thread_trace * tr = &octx->ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
const struct htp_tensor * restrict ids = octx->src[2];
const size_t src0_row_size = nb01;
const uint32_t nb = (ne10 + qk - 1) / qk;
const uint32_t total_nb = src1_nrows * nb;
- worker_callback_t quant_task_func;
+ work_queue_func_t quant_task_func;
uint32_t n_quant_tasks = 1;
if (src1_nrows < octx->n_threads) {
n_quant_tasks = MIN(total_nb, octx->n_threads);
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
- 0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false);
+ 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false, false);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
// Make sure the reserved vtcm size is sufficient
if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-id-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size);
- if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_VTCM_TOO_SMALL;
}
mmctx->n_quant_tasks = n_quant_tasks;
atomic_init(&mmctx->quant_barrier, n_quant_tasks);
- const uint32_t n_matmul_jobs = octx->n_threads;
- worker_pool_run_func(octx->ctx->worker_pool, matmul_id_job_func, mmctx, n_matmul_jobs);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
+ worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads);
- if (must_free_mapping) free(mapping_buf);
return HTP_STATUS_OK;
}
+static inline void scan_expert_ids_n(
+ const struct htp_tensor * ids,
+ const uint32_t n_ids,
+ uint32_t n_as,
+ uint32_t * counts,
+ struct mmid_row_mapping * matrix_rows,
+ uint32_t mapping_stride
+) {
+ const size_t ids_nb1 = ids->nb[1];
+ const uint8_t * ids_data = (const uint8_t *) ids->data;
+
+ for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+ const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
+ for (uint32_t id = 0; id < n_ids; ++id) {
+ const int32_t i02 = row_ptr[id];
+ if (i02 < 0) {
+ continue;
+ }
+ assert(i02 < n_as);
+
+ if (matrix_rows) {
+ matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
+ }
+ counts[i02] += 1;
+ }
+ }
+}
+
+static inline void scan_expert_ids(
+ const struct htp_tensor * ids,
+ uint32_t n_ids,
+ uint32_t n_as,
+ uint32_t * counts,
+ struct mmid_row_mapping * matrix_rows,
+ uint32_t mapping_stride
+) {
+ const size_t ids_nb0 = ids->nb[0];
+
+ if (ids_nb0 == 4) {
+ switch (n_ids) {
+ case 8: scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride); break;
+ case 4: scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride); break;
+ case 2: scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride); break;
+ default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
+ }
+ } else {
+ // Strided fallback
+ const size_t ids_nb1 = ids->nb[1];
+ const uint8_t * ids_data = (const uint8_t *) ids->data;
+ for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+ const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
+ for (uint32_t id = 0; id < n_ids; ++id) {
+ const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
+ if (i02 < 0) {
+ continue;
+ }
+ assert(i02 < n_as);
+
+ if (matrix_rows) {
+ matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
+ }
+ counts[i02] += 1;
+ }
+ }
+ }
+}
+
int op_matmul_id(struct htp_ops_context * octx) {
htp_matmul_tensors_preamble;
+ struct htp_thread_trace * tr = &octx->ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
struct htp_mm_context mmctx_struct = {0};
struct htp_mm_context * mmctx = &mmctx_struct;
mmctx->octx = octx;
const uint32_t src0_nrows = ne01; // per expert
const uint32_t src1_nrows = ne11 * ne12 * ne13;
- worker_callback_t quant_task_func;
- worker_callback_t matmul_id_job_func = src1_nrows > 1 ? hvx_mm_id : hvx_mv_id;
-
- // Compute src0_nrows_per_thread
- mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
- mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
+ mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
+ mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
// row groups
const int n_ids = ids->ne[0]; // n_expert_used
const int n_as = ne02; // n_expert
- size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
- size_t matrix_row_map_size = n_as * ids->ne[0] * ids->ne[1] * sizeof(struct mmid_row_mapping);
- const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
+ uint8_t * mapping_buf = octx->ctx->ddr_spad_base;
+ uint32_t mapping_stride = 1;
+ uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+ struct mmid_row_mapping * matrix_rows = NULL;
- void * mapping_buf = NULL;
- bool must_free_mapping = false;
+ if (src1_nrows > 1) {
+ const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
+ assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
- if (octx->ctx->ddr_spad_base && total_map_size <= octx->ctx->ddr_spad_size) {
- mapping_buf = octx->ctx->ddr_spad_base;
- } else {
- mapping_buf = memalign(128, total_map_size);
- if (mapping_buf) {
- must_free_mapping = true;
- } else {
- return HTP_STATUS_INTERNAL_ERR;
- }
- }
+ hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
- uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
- struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) ((uint8_t *) mapping_buf + matrix_row_counts_size);
+ memset(matrix_row_counts, 0, matrix_row_counts_size);
+ scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
- mmctx->matrix_row_counts = matrix_row_counts;
- mmctx->matrix_rows = matrix_rows;
- mmctx->mm_div_ne11 = kparams->div_ne11;
+ uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
+ mapping_stride = max_count > 0 ? max_count : 1;
- if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) {
- if (must_free_mapping) free(mapping_buf);
- return HTP_STATUS_NO_SUPPORT;
- }
+ size_t matrix_row_map_size = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
+ const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
+
+ if (total_map_size > octx->ctx->ddr_spad_size) {
+ mapping_buf = memalign(128, total_map_size);
+ if (!mapping_buf) {
+ return HTP_STATUS_INTERNAL_ERR;
+ }
+ }
+
+ matrix_row_counts = (uint32_t *) mapping_buf;
+ matrix_rows = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
- if (src1_nrows > 1) {
- // initialize matrix_row_counts and map
memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
+ scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
+ }
- // group rows by src0 matrix
- for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { // token idx
- for (uint32_t id = 0; id < n_ids; ++id) { // expert idx
- const int32_t i02 = *(const int32_t *) ((const uint8_t *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
+ mmctx->matrix_row_counts = matrix_row_counts;
+ mmctx->matrix_rows = matrix_rows;
+ mmctx->mapping_stride = mapping_stride;
+ mmctx->mm_div_ne11 = kparams->div_ne11;
+ mmctx->src0_row_size_padded = src0_row_size_padded;
+ mmctx->src1_nrows = src1_nrows;
- if (i02 < 0) {
- continue;
- }
- assert(i02 < n_as);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
- matrix_rows[i02 * n_ids * ids->ne[1] + matrix_row_counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
- matrix_row_counts[i02] += 1;
- }
+ int s;
+ if (kparams->n_hmx) {
+ s = hmx_mm_op_matmul_id(octx, mmctx);
+ } else {
+ if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) {
+ s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id);
+ } else {
+ s = HTP_STATUS_NO_SUPPORT;
}
}
- if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
- if (must_free_mapping) free(mapping_buf);
- return HTP_STATUS_OK;
+ if (mapping_buf != octx->ctx->ddr_spad_base) {
+ free(mapping_buf);
}
- if (kparams->n_hmx) {
- return hmx_mm_op_matmul_id(octx, mmctx, matrix_row_counts, matrix_rows, mapping_buf, must_free_mapping);
- }
-
- return hvx_mm_matmul_id(octx, mmctx, src0_row_size_padded, src1_nrows, matmul_id_job_func, mapping_buf, must_free_mapping);
+ return s;
}
int op_matmul_qkv(struct htp_ops_context * octx) {
+ struct htp_thread_trace * tr = &octx->ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
const struct htp_tensor * restrict src0 = octx->src[0]; // Wk
const struct htp_tensor * restrict src1 = octx->src[1]; // x
const struct htp_tensor * restrict src2 = octx->src[2]; // Wv
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
- 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false);
+ 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true, false);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
mmctx->vtcm_src3_size_per_thread = L.src3_bytes / octx->n_threads;
mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads;
- if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
- return HTP_STATUS_OK;
-
mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
mmctx->quant_task_func = quant_task_func;
mmctx->n_quant_tasks = n_quant_tasks;
} else {
matmul_job_func = hvx_mm_qkv_2d;
}
+
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs);
return HTP_STATUS_OK;
}
int op_matmul_ffn(struct htp_ops_context * octx) {
+ struct htp_thread_trace * tr = &octx->ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
const struct htp_tensor * restrict src0 = octx->src[0]; // Wgate
const struct htp_tensor * restrict src1 = octx->src[1]; // y
const struct htp_tensor * restrict src2 = octx->src[2]; // Wup
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
- 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true);
+ 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, false, true);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads;
mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads;
- if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
- return HTP_STATUS_OK;
-
mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
mmctx->quant_task_func = quant_task_func;
mmctx->n_quant_tasks = n_quant_tasks;
} else {
matmul_job_func = hvx_mm_ffn_2d;
}
+
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs);
return HTP_STATUS_OK;
struct fastdiv_values div_r2;
struct fastdiv_values div_r3;
struct fastdiv_values div_ne11;
+ struct fastdiv_values div_n_act_threads;
+ struct fastdiv_values div_ne00_padded;
};
#if defined(__cplusplus)
size_t dst_row_size,
size_t src0_row_size,
size_t src1_row_size,
+ size_t src2_row_size,
uint32_t n_prefetch,
bool is_matmul_id,
bool is_fused_qkv,
) {
size_t src0_sz = 0;
size_t src1_sz = 0;
- size_t src2_sz = 0;
+ size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
size_t src3_sz = 0;
size_t dst_sz = 0;
return L.total_bytes;
}
+static inline bool htp_mm_hmx_solve_batched_params(
+ int wtype,
+ uint32_t k,
+ uint32_t ne01_padded,
+ uint32_t ne11,
+ uint32_t group_size,
+ bool use_dma_activation,
+ int n_threads,
+ bool pipeline,
+ size_t vtcm_budget,
+ size_t * m_chunk_out,
+ size_t * n_chunk_out,
+ int * act_threads_out,
+ size_t * vtcm_size_out
+) {
+ size_t best_mblocks = SIZE_MAX;
+ int best_act_threads = 0;
+ size_t best_m_chunk = 0;
+ size_t best_n_chunk = 0;
+ size_t best_vtcm_size = 0;
+
+ int act_threads = n_threads;
+ while (act_threads >= 1) {
+ size_t group_overhead = 256;
+ size_t group_size_per_n, group_size_per_m, group_size_per_mn;
+ htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);
+
+ size_t m_chunk_candidate = 0;
+ size_t n_chunk_candidate = 0;
+ size_t vtcm_size_candidate = 0;
+
+ if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded,
+ (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT,
+ &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
+ size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads);
+ if (exact_size <= vtcm_budget) {
+ size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate;
+ if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
+ best_mblocks = mblocks;
+ best_act_threads = act_threads;
+ best_m_chunk = m_chunk_candidate;
+ best_n_chunk = n_chunk_candidate;
+ best_vtcm_size = exact_size;
+ }
+ }
+ }
+ if (act_threads == 1) {
+ act_threads = 0;
+ } else {
+ act_threads /= 2;
+ }
+ }
+
+ if (best_act_threads > 0) {
+ *m_chunk_out = best_m_chunk;
+ *n_chunk_out = best_n_chunk;
+ *vtcm_size_out = best_vtcm_size;
+ *act_threads_out = best_act_threads;
+ return true;
+ }
+ return false;
+}
+
+static inline bool htp_mm_hmx_solve_2d_params(
+ int wtype,
+ uint32_t k,
+ uint32_t m_id_rows,
+ uint32_t ne01_padded,
+ uint32_t ne11_padded,
+ uint32_t m_for_cost,
+ int n_threads,
+ bool pipeline,
+ bool is_matmul_id,
+ uint32_t aligned_tile_size,
+ size_t vtcm_budget,
+ size_t * m_chunk_out,
+ size_t * n_chunk_out,
+ int * act_threads_out,
+ size_t * vtcm_size_out
+) {
+ size_t best_mblocks = SIZE_MAX;
+ int best_act_threads = 0;
+ size_t best_m_chunk = 0;
+ size_t best_n_chunk = 0;
+ size_t best_vtcm_size = 0;
+
+ const int m_for_chunks = is_matmul_id ? hex_align_up(m_id_rows, 32) : ne11_padded;
+
+ int act_threads = n_threads;
+ while (act_threads >= 1) {
+ size_t simple_2d_overhead = 256;
+ size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn;
+ htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);
+
+ size_t m_chunk_candidate = 0;
+ size_t n_chunk_candidate = 0;
+ size_t vtcm_size_candidate = 0;
+
+ if (htp_mm_hmx_compute_chunks(vtcm_budget, simple_2d_overhead, simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn, m_for_chunks, ne01_padded,
+ (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) m_for_cost * HTP_MM_HMX_COST_A_CONVERT,
+ &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
+ size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size);
+ if (exact_size <= vtcm_budget) {
+ size_t mblocks = ((size_t) m_for_cost + m_chunk_candidate - 1) / m_chunk_candidate;
+ if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
+ best_mblocks = mblocks;
+ best_act_threads = act_threads;
+ best_m_chunk = m_chunk_candidate;
+ best_n_chunk = n_chunk_candidate;
+ best_vtcm_size = exact_size;
+ }
+ }
+ }
+ if (act_threads == 1) {
+ act_threads = 0;
+ } else {
+ act_threads /= 2;
+ }
+ }
+
+ if (best_act_threads > 0) {
+ *m_chunk_out = best_m_chunk;
+ *n_chunk_out = best_n_chunk;
+ *vtcm_size_out = best_vtcm_size;
+ *act_threads_out = best_act_threads;
+ return true;
+ }
+ return false;
+}
+
#ifdef __cplusplus
}
#endif
#include "htp-ctx.h"
#include "htp-ops.h"
#include "htp-ops.h"
+#include "htp-tensor.h"
// Redefined the rope type constants as we can't include ggml.h
#define HTP_ROPE_TYPE_NORMAL 0
}
int op_rope(struct htp_ops_context * octx) {
- int err = HTP_STATUS_OK;
-
switch (octx->src[0]->type) {
case HTP_TYPE_F32:
- err = execute_op_rope_f32(octx);
- break;
+ return execute_op_rope_f32(octx);
default:
- err = HTP_STATUS_NO_SUPPORT;
- break;
+ return HTP_STATUS_NO_SUPPORT;
}
-
- return err;
}
#include <string.h>
#include "hex-dma.h"
+#include "hex-fastdiv.h"
#include "hvx-exp.h"
#include "hvx-sigmoid.h"
#include "hvx-utils.h"
+#include "unary-ops.h"
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "htp-ctx.h"
#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
+#include "hex-profile.h"
struct htp_unary_context {
struct htp_ops_context * octx;
+ const struct htp_unary_kernel_params * kparams;
- // Precomputed values
const uint8_t * data_src0;
const uint8_t * data_src1; // weight/scale tensor for RMS_NORM_MUL
uint8_t * data_dst;
size_t src1_row_size_aligned;
size_t dst_row_size_aligned;
- size_t src0_spad_half_size;
- size_t src1_spad_half_size;
- size_t dst_spad_half_size;
+ size_t src0_vtcm_half_size;
+ size_t src1_vtcm_half_size;
+ size_t dst_vtcm_half_size;
uint32_t block;
uint32_t src0_nrows;
uint32_t src0_nrows_per_thread;
uint32_t nc;
+ uint32_t col_tile; // tiled mode
bool broadcast_weight;
+
+ uint8_t * vtcm_src0;
+ uint8_t * vtcm_src1;
+ uint8_t * vtcm_dst;
+
+ size_t vtcm_src0_size_per_thread;
+ size_t vtcm_src1_size_per_thread;
+ size_t vtcm_dst_size_per_thread;
};
// Convert flat row index to DDR byte offset using the tensor's actual strides.
// ir = i1 + ne1*(i2 + ne2*i3) => offset = i1*nb1 + i2*nb2 + i3*nb3
static inline size_t unary_row_offset(uint32_t ir,
uint32_t ne1, uint32_t ne2,
+ const struct fastdiv_values * div_ne1,
+ const struct fastdiv_values * div_ne2,
+ const struct fastdiv_values * div_ne12,
size_t nb1, size_t nb2, size_t nb3) {
- const uint32_t i1 = ir % ne1;
- const uint32_t i2 = (ir / ne1) % ne2;
- const uint32_t i3 = ir / (ne1 * ne2);
+ const uint32_t i1 = fastmodulo(ir, ne1, div_ne1);
+ const uint32_t ir_div_ne1 = fastdiv(ir, div_ne1);
+ const uint32_t i2 = fastmodulo(ir_div_ne1, ne2, div_ne2);
+ const uint32_t i3 = fastdiv(ir, div_ne12);
return i1 * nb1 + i2 * nb2 + i3 * nb3;
}
+
// Safe DMA block size from row `ir`: clamp to the tighter dim-1 slice
// boundary of src and dst so the nb1 stride stays valid for all rows.
static inline uint32_t unary_block_size(uint32_t ir,
uint32_t block,
bool src_contig,
bool dst_contig,
- uint32_t src_ne1,
- uint32_t dst_ne1) {
+ uint32_t ne1,
+ const struct fastdiv_values * div_ne1) {
uint32_t limit = MIN(block, end_row - ir);
- if (!src_contig) {
- const uint32_t src_slice_end = (ir / src_ne1 + 1) * src_ne1;
- limit = MIN(limit, src_slice_end - ir);
- }
-
- if (!dst_contig) {
- const uint32_t dst_slice_end = (ir / dst_ne1 + 1) * dst_ne1;
- limit = MIN(limit, dst_slice_end - ir);
+ if (!src_contig || !dst_contig) {
+ const uint32_t slice_end = (fastdiv(ir, div_ne1) + 1) * ne1;
+ limit = MIN(limit, slice_end - ir);
}
return limit;
const uint32_t nb2 = dst->nb[2]; \
const uint32_t nb3 = dst->nb[3];
-static void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
- uint8_t * restrict dst,
- uint8_t * restrict pad,
- const int num_elems,
- float epsilon) {
- (void)pad;
-
- const HVX_Vector * restrict v_src = (HVX_Vector *) src;
- HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
-
- const int nvec = num_elems / VLEN_FP32; // number of full vectors
- const int nloe = num_elems % VLEN_FP32; // leftover elements
-
- // Compute sum of squares for full vectors
- HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
- HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
- }
-
- // Handle tail elements using vectorized ops with masking
- if (nloe > 0) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
- }
-
- // Reduce HVX sum
- sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
-
- HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
- HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
- HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
- HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
-
- // Scale full vectors
- HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
- v_dst[i] = Q6_Vsf_equals_Vqf32(v2);
- }
-
- // Handle tail elements using vectorized ops with masking
- if (nloe > 0) {
-
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
- HVX_Vector result = Q6_Vsf_equals_Vqf32(v2);
-
- // Store with masking to avoid overwriting memory beyond the tensor
- hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
- }
-}
-
-static void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
- const uint8_t * restrict weight,
- uint8_t * restrict dst,
- const int num_elems,
- float epsilon) {
- const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
- const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
- HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
-
- const int nvec = num_elems / VLEN_FP32; // number of full vectors
- const int nloe = num_elems % VLEN_FP32; // leftover elements
-
- // Compute sum of squares for full vectors
- HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
- HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
- }
-
- // Handle tail elements using vectorized ops with masking
- if (nloe > 0) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
- }
-
- // Reduce HVX sum
- sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
-
- HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
- HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
- HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
- HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
-
- // Scale and multiply
- HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
- HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
- HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[i]);
- v_dst[i] = Q6_Vsf_equals_Vqf32(result);
- }
-
- // Handle tail elements using vectorized ops with masking
- if (nloe > 0) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
- HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
- HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[nvec]);
- HVX_Vector res_v = Q6_Vsf_equals_Vqf32(result);
-
- // Store with masking to avoid overwriting memory beyond the tensor
- hvx_vec_store_a(&v_dst[nvec], nloe * 4, res_v);
- }
-}
-
-static void hvx_fast_norm_f32(const uint8_t * restrict src,
- uint8_t * restrict dst,
- uint8_t * restrict pad,
- const int num_elems,
- float epsilon) {
- (void)pad;
-
- const HVX_Vector * restrict v_src = (HVX_Vector *) src;
- HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
-
- const int nvec = num_elems / VLEN_FP32; // number of full vectors
- const int nloe = num_elems % VLEN_FP32; // leftover elements
-
- // Compute sum of squares and sum of values for full vectors
- HVX_Vector sum_sq_v = Q6_V_vsplat_R(0x00000000);
- HVX_Vector sum_x_v = Q6_V_vsplat_R(0x00000000);
- HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
- sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
- }
-
- // Handle tail elements using vectorized ops with masking
- if (nloe > 0) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
- sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
- }
-
- // Reduce HVX sums
- sum_sq_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_sq_v));
- sum_x_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_x_v));
-
- HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
- HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
- HVX_Vector mean_sq_v = Q6_Vqf32_vmpy_VsfVsf(sum_sq_v, denom_v);
- HVX_Vector mean_x_v = Q6_Vqf32_vmpy_VsfVsf(sum_x_v, denom_v);
- HVX_Vector mean_x_sq_v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(mean_x_v), Q6_Vsf_equals_Vqf32(mean_x_v));
- HVX_Vector var_v = Q6_Vqf32_vsub_Vqf32Vqf32(mean_sq_v, mean_x_sq_v);
- HVX_Vector var_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(var_v, epsilon_v);
-
- // scale = rsqrt(variance + epsilon), mean_x broadcast for subtraction
- HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(var_epsilon_v));
- HVX_Vector mean_x_b = hvx_vec_repl_f32(Q6_Vsf_equals_Vqf32(mean_x_v));
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
- HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
- v_dst[i] = Q6_Vsf_equals_Vqf32(v3);
- }
-
- // Handle tail elements using vectorized ops with masking
- if (nloe > 0) {
-
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
- HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
- HVX_Vector result = Q6_Vsf_equals_Vqf32(v3);
-
- // Store with masking to avoid overwriting memory beyond the tensor
- hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
- }
-}
+#define htp_unary_op_preamble \
+ int32_t * op_params = uctx->octx->op_params; \
+ const uint32_t ne0 = uctx->nc; \
+ const size_t src0_row_size_aligned = uctx->src0_row_size_aligned; \
+ const size_t dst_row_size_aligned = uctx->dst_row_size_aligned;
static void scale_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
float scale = 0.f;
float bias = 0.f;
memcpy(&scale, &op_params[0], sizeof(float));
memcpy(&bias, &op_params[1], sizeof(float));
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_scale_offset_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems, scale, bias);
+ hvx_scale_offset_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0, scale, bias);
+ }
+}
+
+static void clamp_f32(const float * restrict src,
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
+ float min = 0.f;
+ float max = 0.f;
+ memcpy(&min, &op_params[0], sizeof(float));
+ memcpy(&max, &op_params[1], sizeof(float));
+
+ for (uint32_t ir = 0; ir < num_rows; ir++) {
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+ hvx_clamp_scalar_f32(dst_local, src_local, min, max, ne0);
}
}
static void rms_norm_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
float epsilon = 0.f;
memcpy(&epsilon, op_params, sizeof(float));
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_fast_rms_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, spad, row_elems, epsilon);
+ hvx_fast_rms_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, ne0, epsilon);
}
}
const float * restrict weight,
float * restrict dst,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- const size_t weight_row_size,
- int32_t * op_params,
- bool broadcast_weight) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
float epsilon = 0.f;
memcpy(&epsilon, op_params, sizeof(float));
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- const uint8_t * restrict w_local = (const uint8_t *)weight + (broadcast_weight ? 0 : ir * weight_row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ const uint8_t * restrict w_local = (const uint8_t *)weight + (uctx->broadcast_weight ? 0 : ir * uctx->src1_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_fast_rms_norm_mul_f32(src_local, w_local, dst_local, row_elems, epsilon);
+ hvx_fast_rms_norm_mul_f32(src_local, w_local, dst_local, ne0, epsilon);
}
}
static void norm_f32(const float * restrict src,
- float * restrict dst,
- uint8_t * restrict spad,
- const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
float epsilon = 0.f;
memcpy(&epsilon, op_params, sizeof(float));
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_fast_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, spad, row_elems, epsilon);
+ hvx_fast_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, ne0, epsilon);
}
}
static void sqr_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_sqr_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems);
+ hvx_sqr_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
}
}
static void sqrt_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_sqrt_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems);
+ hvx_sqrt_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
}
}
static void neg_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_scale_f32_aa(dst_local, src_local, row_elems, -1.0f);
+ hvx_scale_f32_aa(dst_local, src_local, ne0, -1.0f);
}
}
static void exp_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_exp_f32(dst_local, src_local, row_elems, false);
+ hvx_exp_f32(dst_local, src_local, ne0, false);
}
}
static void sigmoid_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_sigmoid_f32_aa(dst_local, src_local, row_elems);
+ hvx_sigmoid_f32_aa(dst_local, src_local, ne0);
+ }
+}
+
+// silu(x) = x * sigmoid(x)
+static void silu_f32(const float * restrict src,
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
+
+ for (uint32_t ir = 0; ir < num_rows; ir++) {
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+ hvx_sigmoid_f32_aa(dst_local, src_local, ne0);
+ hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
+ }
+}
+
+// gelu(x) = x * sigmoid(1.702 * x) (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+static void gelu_f32(const float * restrict src,
+ float * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
+
+ for (uint32_t ir = 0; ir < num_rows; ir++) {
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+ hvx_mul_scalar_f32(dst_local, src_local, 1.702f, ne0);
+ hvx_sigmoid_f32_aa(dst_local, dst_local, ne0);
+ hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
}
}
static void tri_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params,
const uint32_t ir,
const struct htp_unary_context * uctx) {
-
+ htp_unary_op_preamble;
const int32_t ttype = op_params[0];
const HVX_Vector zero = hvx_vec_splat_f32(0.0f);
- const uint32_t nvec = row_elems / VLEN_FP32;
- const uint32_t nloe = row_elems % VLEN_FP32;
+ const uint32_t nvec = ne0 / VLEN_FP32;
+ const uint32_t nloe = ne0 % VLEN_FP32;
const uint32_t ne01 = uctx->octx->src[0]->ne[1];
const uint32_t abs_row = ir + b;
const uint32_t i01 = abs_row % ne01;
- const HVX_Vector * restrict v_src = (const HVX_Vector *) ((const uint8_t *) src + b * row_size);
- HVX_Vector * restrict v_dst = (HVX_Vector *) ((uint8_t *) dst + b * row_size);
+ const HVX_Vector * restrict v_src = (const HVX_Vector *) ((const uint8_t *) src + b * src0_row_size_aligned);
+ HVX_Vector * restrict v_dst = (HVX_Vector *) ((uint8_t *) dst + b * dst_row_size_aligned);
uint32_t boundary;
int keep_left;
case 3: boundary = i01; keep_left = 1; break; // keep col < row
default: boundary = 0; keep_left = 0; break;
}
- if (boundary > row_elems) boundary = row_elems;
+ if (boundary > ne0) boundary = ne0;
// Full HVX vectors — each starts at a 128-byte aligned offset
for (uint32_t i = 0; i < nvec; i++) {
// Tail elements (row_elems not a multiple of VLEN_FP32)
if (nloe > 0) {
- const uint32_t vec_start = nvec * VLEN_FP32;
- const uint32_t vec_end = vec_start + nloe;
+ const uint32_t abs_start = nvec * VLEN_FP32;
+ const uint32_t abs_end = abs_start + nloe;
HVX_Vector tail_val;
if (keep_left) {
- if (vec_end <= boundary) {
+ if (abs_end <= boundary) {
tail_val = v_src[nvec];
- } else if (vec_start >= boundary) {
+ } else if (abs_start >= boundary) {
tail_val = zero;
} else {
- HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - vec_start) * sizeof(float));
+ HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - abs_start) * sizeof(float));
tail_val = Q6_V_vmux_QVV(mask, v_src[nvec], zero);
}
} else {
- if (vec_end <= boundary) {
+ if (abs_end <= boundary) {
tail_val = zero;
- } else if (vec_start >= boundary) {
+ } else if (abs_start >= boundary) {
tail_val = v_src[nvec];
} else {
- HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - vec_start) * sizeof(float));
+ HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - abs_start) * sizeof(float));
tail_val = Q6_V_vmux_QVV(mask, zero, v_src[nvec]);
}
}
static void softplus_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
// softplus(x) = log(1 + exp(x))
// Match CPU reference: ggml_compute_softplus_f32() in ggml-impl.h
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * row_size));
- float * restrict dst_f = (float *)((uint8_t *)dst + (ir * row_size));
+ const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * src0_row_size_aligned));
+ float * restrict dst_f = (float *)((uint8_t *)dst + (ir * dst_row_size_aligned));
- for (uint32_t i = 0; i < row_elems; i++) {
+ for (uint32_t i = 0; i < ne0; i++) {
float x = src_f[i];
// For x > 20: softplus(x) ≈ x (avoids exp overflow)
dst_f[i] = (x > 20.0f) ? x : logf(1.0f + expf(x));
}
}
-// --- L2_NORM HVX kernel ---
-// Computes y[i] = x[i] / fmax(sqrt(sum(x[j]^2)), epsilon) for each row.
-// scale = 1/fmax(sqrt(sum), epsilon) is computed entirely in HVX registers
-// using rsqrt + inverse to avoid scalar extraction.
-static void hvx_fast_l2_norm_f32(const uint8_t * restrict src,
- uint8_t * restrict dst,
- uint8_t * restrict pad,
- const int num_elems,
- float epsilon) {
- (void)pad;
-
- const HVX_Vector * restrict v_src = (HVX_Vector *) src;
- HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
-
- HVX_Vector sum_v = hvx_vec_splat_f32(0.0f);
-
- const int nvec = num_elems / VLEN_FP32;
- const int nloe = num_elems % VLEN_FP32;
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
- }
-
- // Include tail elements in the sum-of-squares using a predicate mask
- if (nloe > 0) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
- sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
- }
-
- // Compute scale = 1/fmax(sqrt(sum), epsilon) entirely in HVX registers.
- // hvx_vec_rsqrt_f32 + hvx_vec_inverse_f32 avoids scalar extraction.
- HVX_Vector sum_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
- HVX_Vector rsqrt_v = hvx_vec_rsqrt_f32(sum_sf); // 1/sqrt(sum)
- HVX_Vector sqrt_v = hvx_vec_inverse_f32(rsqrt_v); // sqrt(sum)
- HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
- HVX_Vector denom_v = Q6_Vsf_vmax_VsfVsf(sqrt_v, epsilon_v); // fmax(sqrt(sum), epsilon)
- HVX_Vector scale_v = hvx_vec_inverse_f32(denom_v); // 1/fmax(sqrt(sum), epsilon)
-
- #pragma unroll(4)
- for (int i = 0; i < nvec; i++) {
- HVX_Vector v1 = v_src[i];
- v_dst[i] = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
- }
-
- if (nloe > 0) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
- HVX_Vector result = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
- hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
- }
-}
-
static void l2_norm_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
float epsilon = 0.f;
memcpy(&epsilon, op_params, sizeof(float));
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * row_size));
- float * restrict dst_f = (float *)((uint8_t *)dst + (ir * row_size));
+ const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * src0_row_size_aligned));
+ float * restrict dst_f = (float *)((uint8_t *)dst + (ir * dst_row_size_aligned));
- hvx_fast_l2_norm_f32((const uint8_t *)src_f, (uint8_t *)dst_f, spad, row_elems, epsilon);
+ hvx_fast_l2_norm_f32((const uint8_t *)src_f, (uint8_t *)dst_f, ne0, epsilon);
}
}
static void tanh_f32(const float * restrict src,
float * restrict dst,
- uint8_t * restrict spad,
const uint32_t num_rows,
- const uint32_t row_elems,
- const size_t row_size,
- int32_t * op_params) {
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
+
for (uint32_t ir = 0; ir < num_rows; ir++) {
- const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size);
- uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size);
+ const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned);
- hvx_tanh_f32_aa(dst_local, src_local, row_elems);
+ hvx_tanh_f32_aa(dst_local, src_local, ne0);
}
}
-static void unary_job_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
- const struct htp_unary_context * uctx = (const struct htp_unary_context *) data;
- struct htp_ops_context * octx = uctx->octx;
- const struct htp_tensor * src = octx->src[0];
- const struct htp_tensor * dst = octx->dst;
-
- htp_unary_preamble;
-
- int htp_op = octx->op;
- int32_t * op_params = octx->op_params;
- uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread;
-
- const size_t src0_data_row_size = uctx->src0_data_row_size;
- const size_t dst_data_row_size = uctx->dst_data_row_size;
-
- const size_t src0_row_size_aligned = uctx->src0_row_size_aligned;
- const size_t dst_row_size_aligned = uctx->dst_row_size_aligned;
-
- const uint32_t src0_nrows = uctx->src0_nrows;
- const uint32_t src0_start_row = src0_nrows_per_thread * ith;
- const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+#define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \
+static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * data) { \
+ const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; \
+ struct htp_ops_context * octx = uctx->octx; \
+ const struct htp_tensor * src = octx->src[0]; \
+ const struct htp_tensor * dst = octx->dst; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
+ \
+ htp_unary_preamble; \
+ \
+ int32_t * op_params = octx->op_params; \
+ uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread; \
+ \
+ const size_t src0_data_row_size = uctx->src0_data_row_size; \
+ const size_t dst_data_row_size = uctx->dst_data_row_size; \
+ \
+ const size_t src0_row_size_aligned = uctx->src0_row_size_aligned; \
+ const size_t dst_row_size_aligned = uctx->dst_row_size_aligned; \
+ \
+ const uint32_t src0_nrows = uctx->src0_nrows; \
+ const uint32_t src0_start_row = src0_nrows_per_thread * ith; \
+ const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); \
+ \
+ if (src0_start_row >= src0_end_row) { \
+ return; \
+ } \
+ \
+ const uint8_t * restrict data_src = uctx->data_src0; \
+ const uint8_t * restrict data_src1 = uctx->data_src1; \
+ uint8_t * restrict data_dst = uctx->data_dst; \
+ \
+ const struct htp_tensor * src1 = (IS_RMS_NORM_MUL) ? octx->src[1] : NULL; \
+ const uint32_t nb11 = src1 ? src1->nb[1] : 0; \
+ const uint32_t nb12 = src1 ? src1->nb[2] : 0; \
+ const uint32_t nb13 = src1 ? src1->nb[3] : 0; \
+ const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false; \
+ \
+ uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \
+ uint8_t * src1_vtcm_data = uctx->vtcm_src1 ? (uctx->vtcm_src1 + (ith * uctx->vtcm_src1_size_per_thread)) : NULL;\
+ uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); \
+ \
+ size_t src0_vtcm_half_size = uctx->src0_vtcm_half_size; \
+ size_t src1_vtcm_half_size = uctx->src1_vtcm_half_size; \
+ size_t dst_vtcm_half_size = uctx->dst_vtcm_half_size; \
+ \
+ const bool src0_contig = (nb02 == (size_t)ne01 * nb01) && \
+ (nb03 == (size_t)ne02 * nb02); \
+ const bool dst_contig = (nb2 == (size_t)ne1 * nb1) && \
+ (nb3 == (size_t)ne2 * nb2); \
+ \
+ const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; \
+ const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \
+ const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \
+ \
+ const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \
+ const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \
+ const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); \
+ if (BLOCK == 0) { \
+ FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small, needed at least %zu\n", \
+ uctx->vtcm_src0_size_per_thread, src0_row_size_aligned); \
+ return; \
+ } \
+ \
+ dma_queue * dma_queue = octx->ctx->dma[ith]; \
+ \
+ if ((IS_RMS_NORM_MUL) && uctx->broadcast_weight) { \
+ dma_queue_push(dma_queue, dma_make_ptr(src1_vtcm_data, data_src1), \
+ uctx->src1_row_size_aligned, 0, uctx->src1_data_row_size, 1); \
+ dma_queue_flush(dma_queue); \
+ } \
+ \
+ for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { \
+ const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \
+ div_ne01); \
+ \
+ dma_queue_push(dma_queue, \
+ dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), \
+ nb1, dst_row_size_aligned, dst_data_row_size, 0); \
+ \
+ const size_t src0_off = src0_contig ? (ir * nb01) : \
+ unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \
+ dma_queue_push(dma_queue, \
+ dma_make_ptr(src0_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src + src0_off), \
+ src0_row_size_aligned, nb01, src0_data_row_size, block_size); \
+ \
+ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \
+ const size_t src1_off = src1_contig ? (ir * nb11) : \
+ unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \
+ dma_queue_push(dma_queue, \
+ dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), \
+ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size); \
+ } \
+ \
+ ir += block_size; \
+ } \
+ \
+ for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { \
+ const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \
+ div_ne01); \
+ \
+ float * dst_vtcm = (float *) dma_queue_pop(dma_queue).src; \
+ float * src0_vtcm = (float *) dma_queue_pop(dma_queue).dst; \
+ float * src1_vtcm = NULL; \
+ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \
+ src1_vtcm = (float *) dma_queue_pop(dma_queue).dst; \
+ } \
+ \
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); \
+ CORE_EXPR; \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); \
+ \
+ const size_t dst_off = dst_contig ? (ir * nb1) : \
+ unary_row_offset(ir, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3); \
+ dma_queue_push(dma_queue, \
+ dma_make_ptr(data_dst + dst_off, dst_vtcm), \
+ nb1, dst_row_size_aligned, dst_data_row_size, block_size); \
+ \
+ const uint32_t next_ir = ir + block_size; \
+ if (next_ir < src0_end_row) { \
+ const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig,\
+ ne01, div_ne01); \
+ const uint32_t pref_ir = next_ir + next_block_size; \
+ if (pref_ir < src0_end_row) { \
+ const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, \
+ dst_contig, ne01, div_ne01); \
+ const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : \
+ unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \
+ dma_queue_push(dma_queue, \
+ dma_make_ptr(src0_vtcm, data_src + src0_pref_off), \
+ src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); \
+ \
+ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \
+ const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : \
+ unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \
+ dma_queue_push(dma_queue, \
+ dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off), \
+ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); \
+ } \
+ } \
+ } \
+ ir += block_size; \
+ } \
+ \
+ dma_queue_flush(dma_queue); \
+}
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
- }
+DEFINE_UNARY_TASK(norm, false, false, norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(rms_norm, false, false, rms_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(rms_norm_mul, true, false, rms_norm_mul_f32(src0_vtcm, uctx->broadcast_weight ? (const float *) src1_vtcm_data : src1_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(scale, false, false, scale_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(clamp, false, false, clamp_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(sqr, false, false, sqr_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(sqrt, false, false, sqrt_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_neg, false, false, neg_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_exp, false, false, exp_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_sigmoid, false, false, sigmoid_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_silu, false, false, silu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_gelu, false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_tanh, false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(l2_norm, false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(tri, false, true, tri_f32(src0_vtcm, dst_vtcm, block_size, ir, uctx))
+
+// Apply a pointwise unary op to one column tile that is already in VTCM.
+#define DEFINE_UNARY_TILED_TASK(NAME, IS_TRI, CORE_TILE_EXPR) \
+static void unary_task_f32_tiled_##NAME(unsigned int nth, unsigned int ith, void * data) { \
+ const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; \
+ struct htp_ops_context * octx = uctx->octx; \
+ const struct htp_tensor * src = octx->src[0]; \
+ const struct htp_tensor * dst = octx->dst; \
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
+ \
+ htp_unary_preamble; \
+ \
+ int32_t * op_params = octx->op_params; \
+ const uint32_t col_tile = uctx->col_tile; \
+ \
+ const uint32_t src0_nrows = uctx->src0_nrows; \
+ const uint32_t src0_start_row = uctx->src0_nrows_per_thread * ith; \
+ const uint32_t src0_end_row = MIN(src0_start_row + uctx->src0_nrows_per_thread, src0_nrows); \
+ \
+ if (src0_start_row >= src0_end_row) { \
+ return; \
+ } \
+ \
+ const uint8_t * restrict data_src = uctx->data_src0; \
+ uint8_t * restrict data_dst = uctx->data_dst; \
+ \
+ uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \
+ uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); \
+ \
+ const size_t src0_half = uctx->src0_vtcm_half_size; \
+ const size_t dst_half = uctx->dst_vtcm_half_size; \
+ \
+ dma_queue * dmaq = octx->ctx->dma[ith]; \
+ \
+ const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; \
+ const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \
+ const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \
+ const struct fastdiv_values * div_tpr = &uctx->kparams->div_tpr; \
+ \
+ const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile; \
+ const int32_t tri_ttype = (IS_TRI) ? op_params[0] : 0; \
+ \
+ const bool src0_contig = (nb02 == (size_t)ne01 * nb01) && \
+ (nb03 == (size_t)ne02 * nb02); \
+ const bool dst_contig = (nb2 == (size_t)ne1 * nb1) && \
+ (nb3 == (size_t)ne2 * nb2); \
+ \
+ const uint32_t total_tiles = (src0_end_row - src0_start_row) * tiles_per_row; \
+ \
+ for (uint32_t t = 0, vtcm_idx = 0; t < total_tiles && vtcm_idx < 2; t++, vtcm_idx++) { \
+ const uint32_t row = src0_start_row + t / tiles_per_row; \
+ const uint32_t col = (t % tiles_per_row) * col_tile; \
+ const uint32_t tw = MIN(col_tile, ne0 - col); \
+ const size_t tb = (size_t) tw * sizeof(float); \
+ const size_t soff = (src0_contig ? (row * nb01) : \
+ unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) +\
+ (size_t) col * sizeof(float); \
+ \
+ dma_queue_push(dmaq, dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_half)), 0, 0, 0, 0); \
+ dma_queue_push(dmaq, dma_make_ptr(src0_vtcm_data + (vtcm_idx * src0_half), data_src + soff), tb, tb, tb, 1);\
+ } \
+ \
+ uint32_t row = src0_start_row; \
+ uint32_t col = 0; \
+ uint32_t tile_in_row = 0; \
+ uint32_t i01 = fastmodulo(row, ne01, div_ne01); \
+ \
+ uint32_t prow = src0_start_row + fastdiv(2, div_tpr); \
+ uint32_t pcol = fastmodulo(2, tiles_per_row, div_tpr) * col_tile; \
+ uint32_t ptile_in_row = fastmodulo(2, tiles_per_row, div_tpr); \
+ \
+ for (uint32_t t = 0; t < total_tiles; t++) { \
+ uint8_t * dst_vtcm = (uint8_t *) dma_queue_pop(dmaq).src; \
+ uint8_t * src_vtcm = (uint8_t *) dma_queue_pop(dmaq).dst; \
+ \
+ const uint32_t tw = MIN(col_tile, ne0 - col); \
+ \
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, t); \
+ CORE_TILE_EXPR; \
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, t); \
+ \
+ const size_t doff = (dst_contig ? (row * nb1) : \
+ unary_row_offset(row, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3)) + \
+ (size_t) col * sizeof(float); \
+ const size_t tb = (size_t) tw * sizeof(float); \
+ dma_queue_push(dmaq, dma_make_ptr(data_dst + doff, dst_vtcm), tb, tb, tb, 1); \
+ \
+ const uint32_t pt = t + 2; \
+ if (pt < total_tiles) { \
+ const uint32_t ptw = MIN(col_tile, ne0 - pcol); \
+ const size_t ptb = (size_t) ptw * sizeof(float); \
+ const size_t psoff = (src0_contig ? (prow * nb01) : \
+ unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, \
+ nb03)) + \
+ (size_t) pcol * sizeof(float); \
+ dma_queue_push(dmaq, dma_make_ptr(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1); \
+ } \
+ \
+ tile_in_row++; \
+ col += col_tile; \
+ if (tile_in_row == tiles_per_row) { \
+ tile_in_row = 0; \
+ col = 0; \
+ row++; \
+ i01++; \
+ if (i01 == ne01) { \
+ i01 = 0; \
+ } \
+ } \
+ \
+ ptile_in_row++; \
+ pcol += col_tile; \
+ if (ptile_in_row == tiles_per_row) { \
+ ptile_in_row = 0; \
+ pcol = 0; \
+ prow++; \
+ } \
+ } \
+ \
+ dma_queue_flush(dmaq); \
+}
- uint64_t t1, t2;
- t1 = HAP_perf_get_qtimer_count();
-
- const uint8_t * restrict data_src = uctx->data_src0;
- const uint8_t * restrict data_src1 = uctx->data_src1;
- uint8_t * restrict data_dst = uctx->data_dst;
-
- const struct htp_tensor * src1 = (htp_op == HTP_OP_RMS_NORM_MUL) ? octx->src[1] : NULL;
- const uint32_t nb11 = src1 ? src1->nb[1] : 0;
- const uint32_t nb12 = src1 ? src1->nb[2] : 0;
- const uint32_t nb13 = src1 ? src1->nb[3] : 0;
-
- uint8_t * src0_spad_data = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
- uint8_t * src1_spad_data = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread);
- uint8_t * dst_spad_data = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread);
-
- size_t src0_spad_half_size = uctx->src0_spad_half_size;
- size_t src1_spad_half_size = uctx->src1_spad_half_size;
- size_t dst_spad_half_size = uctx->dst_spad_half_size;
-
- // Non-contiguous tensors have gaps at dim-2/3 boundaries that a single-stride
- // 2D DMA descriptor cannot span. Clamp BLOCK to ne1 (one dim-1 slice) so every
- // transfer stays within a nb1-uniform region. Skipped for contiguous tensors.
- const bool src0_contig = (nb02 == (size_t)ne01 * nb01) &&
- (nb03 == (size_t)ne02 * nb02);
- const bool dst_contig = (nb2 == (size_t)ne1 * nb1) &&
- (nb3 == (size_t)ne2 * nb2);
- const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01);
- const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1);
- const uint32_t BLOCK = MIN(src0_max_block, dst_max_block);
- if (BLOCK == 0) {
- FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",
- octx->src0_spad.size_per_thread, src0_row_size_aligned);
- return;
- }
+static inline void tile_scale_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) {
+ float scale = 0.f;
+ float bias = 0.f;
+ memcpy(&scale, &op_params[0], sizeof(float));
+ memcpy(&bias, &op_params[1], sizeof(float));
+ hvx_scale_offset_f32_aa(dst_vtcm, src_vtcm, tw, scale, bias);
+}
- dma_queue * dma_queue = octx->ctx->dma[ith];
+static inline void tile_clamp_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) {
+ float min = 0.f;
+ float max = 0.f;
+ memcpy(&min, &op_params[0], sizeof(float));
+ memcpy(&max, &op_params[1], sizeof(float));
+ hvx_clamp_scalar_f32(dst_vtcm, src_vtcm, min, max, tw);
+}
- // If weight is broadcasted, load it once per thread at the beginning of execution
- if (htp_op == HTP_OP_RMS_NORM_MUL && uctx->broadcast_weight) {
- dma_queue_push(dma_queue, dma_make_ptr(src1_spad_data, data_src1), uctx->src1_row_size_aligned, 0, uctx->src1_data_row_size, 1);
- dma_queue_flush(dma_queue);
+static inline void tile_unary_softplus_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+ const float * restrict sf = (const float *) src_vtcm;
+ float * restrict df = (float *) dst_vtcm;
+ for (uint32_t i = 0; i < tw; i++) {
+ float x = sf[i];
+ df[i] = (x > 20.0f) ? x : logf(1.0f + expf(x));
}
+}
- for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; spad_idx++) {
- const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1);
+// silu(x) = x * sigmoid(x)
+static inline void tile_silu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+ hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw);
+ hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw);
+}
- // Dummy DMA transation for sequencing (interleaving dst,src,dst,...)
- dma_queue_push(dma_queue,
- dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
- nb1, dst_row_size_aligned, dst_data_row_size, 0);
+// gelu(x) = x * sigmoid(1.702 * x) (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+static inline void tile_gelu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+ hvx_mul_scalar_f32(dst_vtcm, src_vtcm, 1.702f, tw);
+ hvx_sigmoid_f32_aa(dst_vtcm, dst_vtcm, tw);
+ hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw);
+}
- const size_t src0_off = unary_row_offset(ir, ne01, ne02, nb01, nb02, nb03);
- dma_queue_push(dma_queue,
- dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src + src0_off),
- src0_row_size_aligned, nb01, src0_data_row_size, block_size);
+// Triangular mask applied to one column tile. Boundary is an absolute column index, so
+// each vector compares against its absolute column position (col_start + i*VLEN_FP32).
+static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * restrict dst,
+ uint32_t tile_elems, uint32_t col_start, uint32_t i01,
+ uint32_t ne0, int32_t ttype) {
+ const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
+ HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
+ const HVX_Vector zero = hvx_vec_splat_f32(0.0f);
- if (htp_op == HTP_OP_RMS_NORM_MUL && !uctx->broadcast_weight) {
- const size_t src1_off = unary_row_offset(ir, ne01, ne02, nb11, nb12, nb13);
- dma_queue_push(dma_queue,
- dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + src1_off),
- uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size);
+ uint32_t boundary;
+ int keep_left;
+ switch (ttype) {
+ case 0: boundary = i01; keep_left = 0; break;
+ case 1: boundary = i01 + 1; keep_left = 0; break;
+ case 2: boundary = i01 + 1; keep_left = 1; break;
+ case 3: boundary = i01; keep_left = 1; break;
+ default: boundary = 0; keep_left = 0; break;
+ }
+ if (boundary > ne0) boundary = ne0;
+
+ const uint32_t nvec = tile_elems / VLEN_FP32;
+ const uint32_t nloe = tile_elems % VLEN_FP32;
+
+ for (uint32_t i = 0; i < nvec; i++) {
+ const uint32_t abs_start = col_start + i * VLEN_FP32;
+ const uint32_t abs_end = abs_start + VLEN_FP32;
+ if (keep_left) {
+ if (abs_end <= boundary) {
+ v_dst[i] = v_src[i];
+ } else if (abs_start >= boundary) {
+ v_dst[i] = zero;
+ } else {
+ HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - abs_start) * sizeof(float));
+ v_dst[i] = Q6_V_vmux_QVV(mask, v_src[i], zero);
+ }
+ } else {
+ if (abs_end <= boundary) {
+ v_dst[i] = zero;
+ } else if (abs_start >= boundary) {
+ v_dst[i] = v_src[i];
+ } else {
+ HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - abs_start) * sizeof(float));
+ v_dst[i] = Q6_V_vmux_QVV(mask, zero, v_src[i]);
+ }
}
-
- ir += block_size;
}
- for (uint32_t ir = src0_start_row; ir < src0_end_row; ) {
- const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1);
-
- float * dst_spad = (float *) dma_queue_pop(dma_queue).src;
- float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;
- float * src1_spad = NULL;
- if (htp_op == HTP_OP_RMS_NORM_MUL && !uctx->broadcast_weight) {
- src1_spad = (float *) dma_queue_pop(dma_queue).dst;
- }
-
- // Process block in VTCM
- switch (htp_op) {
- case HTP_OP_NORM:
- norm_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_RMS_NORM:
- rms_norm_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_RMS_NORM_MUL:
- {
- const float * w_ptr = uctx->broadcast_weight ? (const float *) src1_spad_data : src1_spad;
- rms_norm_mul_f32(src0_spad, w_ptr, dst_spad, block_size, ne0, src0_row_size_aligned, uctx->src1_row_size_aligned, op_params, uctx->broadcast_weight);
- }
- break;
- case HTP_OP_SCALE:
- scale_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_SQR:
- sqr_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_SQRT:
- sqrt_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_UNARY_NEG:
- neg_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_UNARY_EXP:
- exp_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_UNARY_SIGMOID:
- sigmoid_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_UNARY_SOFTPLUS:
- softplus_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_UNARY_TANH:
- tanh_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_L2_NORM:
- l2_norm_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params);
- break;
- case HTP_OP_TRI:
- tri_f32(src0_spad, dst_spad, NULL, block_size, ne00, src0_row_size_aligned, op_params, ir, uctx);
- break;
- default:
- break;
- }
-
- const size_t dst_off = unary_row_offset(ir, ne1, ne2, nb1, nb2, nb3);
- dma_queue_push(dma_queue,
- dma_make_ptr(data_dst + dst_off, dst_spad),
- nb1, dst_row_size_aligned, dst_data_row_size, block_size);
-
- // prefetch N+2 loop iteration if any
- const uint32_t next_ir = ir + block_size;
- if (next_ir < src0_end_row) {
- const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1);
- const uint32_t pref_ir = next_ir + next_block_size;
- if (pref_ir < src0_end_row) {
- const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1);
- const size_t src0_pref_off = unary_row_offset(pref_ir, ne01, ne02, nb01, nb02, nb03);
- dma_queue_push(dma_queue,
- dma_make_ptr(src0_spad, data_src + src0_pref_off),
- src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size);
-
- if (htp_op == HTP_OP_RMS_NORM_MUL && !uctx->broadcast_weight) {
- const size_t src1_pref_off = unary_row_offset(pref_ir, ne01, ne02, nb11, nb12, nb13);
- dma_queue_push(dma_queue,
- dma_make_ptr(src1_spad, data_src1 + src1_pref_off),
- uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size);
- }
+ if (nloe > 0) {
+ const uint32_t abs_start = col_start + nvec * VLEN_FP32;
+ const uint32_t abs_end = abs_start + nloe;
+ HVX_Vector tail_val;
+ if (keep_left) {
+ if (abs_end <= boundary) {
+ tail_val = v_src[nvec];
+ } else if (abs_start >= boundary) {
+ tail_val = zero;
+ } else {
+ HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - abs_start) * sizeof(float));
+ tail_val = Q6_V_vmux_QVV(mask, v_src[nvec], zero);
+ }
+ } else {
+ if (abs_end <= boundary) {
+ tail_val = zero;
+ } else if (abs_start >= boundary) {
+ tail_val = v_src[nvec];
+ } else {
+ HVX_VectorPred mask = Q6_Q_vsetq_R((boundary - abs_start) * sizeof(float));
+ tail_val = Q6_V_vmux_QVV(mask, zero, v_src[nvec]);
}
}
- ir += block_size;
+ hvx_vec_store_a(&v_dst[nvec], nloe * sizeof(float), tail_val);
}
-
- dma_queue_flush(dma_queue);
-
- t2 = HAP_perf_get_qtimer_count();
-
- FARF(HIGH, "unary-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", ith, nth, src->ne[0],
- src->ne[1], src->ne[2], src->ne[3], src0_start_row, src0_end_row, dst->ne[0], dst->ne[1], dst->ne[2],
- dst->ne[3], (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
}
+DEFINE_UNARY_TILED_TASK(scale, false, tile_scale_f32(dst_vtcm, src_vtcm, tw, op_params))
+DEFINE_UNARY_TILED_TASK(clamp, false, tile_clamp_f32(dst_vtcm, src_vtcm, tw, op_params))
+DEFINE_UNARY_TILED_TASK(sqr, false, hvx_sqr_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(sqrt, false, hvx_sqrt_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_neg, false, hvx_scale_f32_aa(dst_vtcm, src_vtcm, tw, -1.0f))
+DEFINE_UNARY_TILED_TASK(unary_exp, false, hvx_exp_f32(dst_vtcm, src_vtcm, tw, false))
+DEFINE_UNARY_TILED_TASK(unary_sigmoid, false, hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_silu, false, tile_silu_f32(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_gelu, false, tile_gelu_f32(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_tanh, false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(tri, true, tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype))
+
static int execute_op_unary_f32(struct htp_ops_context * octx) {
int err = HTP_STATUS_OK;
const char * op_type = NULL;
switch (octx->op) {
- case HTP_OP_NORM:
- op_type = "norm-f32";
- break;
- case HTP_OP_RMS_NORM:
- op_type = "rmsnorm-f32";
- break;
- case HTP_OP_RMS_NORM_MUL:
- op_type = "rmsnorm-mul-f32";
- break;
- case HTP_OP_SCALE:
- op_type = "scale-f32";
- break;
- case HTP_OP_SQR:
- op_type = "sqr-f32";
- break;
- case HTP_OP_SQRT:
- op_type = "sqrt-f32";
- break;
- case HTP_OP_UNARY_NEG:
- op_type = "neg-f32";
- break;
- case HTP_OP_UNARY_EXP:
- op_type = "exp-f32";
- break;
- case HTP_OP_UNARY_SIGMOID:
- op_type = "sigmoid-f32";
- break;
- case HTP_OP_UNARY_SOFTPLUS:
- op_type = "softplus-f32";
- break;
- case HTP_OP_UNARY_TANH:
- op_type = "tanh-f32";
- break;
- case HTP_OP_L2_NORM:
- op_type = "l2norm-f32";
- break;
- case HTP_OP_TRI:
- op_type = "tri-f32";
- break;
+ case HTP_OP_NORM: op_type = "norm-f32"; break;
+ case HTP_OP_RMS_NORM: op_type = "rmsnorm-f32"; break;
+ case HTP_OP_RMS_NORM_MUL: op_type = "rmsnorm-mul-f32"; break;
+ case HTP_OP_SCALE: op_type = "scale-f32"; break;
+ case HTP_OP_CLAMP: op_type = "clamp-f32"; break;
+ case HTP_OP_SQR: op_type = "sqr-f32"; break;
+ case HTP_OP_SQRT: op_type = "sqrt-f32"; break;
+ case HTP_OP_UNARY_NEG: op_type = "neg-f32"; break;
+ case HTP_OP_UNARY_EXP: op_type = "exp-f32"; break;
+ case HTP_OP_UNARY_SIGMOID: op_type = "sigmoid-f32"; break;
+ case HTP_OP_UNARY_SILU: op_type = "silu-f32"; break;
+ case HTP_OP_UNARY_GELU: op_type = "gelu-f32"; break;
+ case HTP_OP_UNARY_SOFTPLUS: op_type = "softplus-f32"; break;
+ case HTP_OP_UNARY_TANH: op_type = "tanh-f32"; break;
+ case HTP_OP_L2_NORM: op_type = "l2norm-f32"; break;
+ case HTP_OP_TRI: op_type = "tri-f32"; break;
default:
FARF(ERROR, "Unsupported unary Op %u\n", octx->op);
return HTP_STATUS_NO_SUPPORT;
}
+ const struct htp_unary_kernel_params * kparams = (const struct htp_unary_kernel_params *) octx->kernel_params;
+
const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
- const uint32_t n_threads = MIN(octx->n_threads, src0_nrows);
+ const uint32_t n_threads = kparams->n_threads;
const size_t src0_data_row_size = src0->ne[0] * sizeof(float);
const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
- const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, VLEN);
- const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, VLEN);
+ const size_t src0_row_size_aligned = kparams->src0_row_size_aligned;
+ const size_t dst_row_size_aligned = kparams->dst_row_size_aligned;
+
+ const uint32_t col_tile = kparams->col_tile;
size_t src1_data_row_size = 0;
- size_t src1_row_size_aligned = 0;
- bool broadcast_weight = false;
+ size_t src1_row_size_aligned = kparams->src1_row_size_aligned;
+ bool broadcast_weight = kparams->broadcast_weight;
const struct htp_tensor * src1 = NULL;
if (octx->op == HTP_OP_RMS_NORM_MUL) {
src1 = octx->src[1];
src1_data_row_size = src1->ne[0] * sizeof(float);
- src1_row_size_aligned = hex_round_up(src1_data_row_size, VLEN);
- broadcast_weight = (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
- }
-
- // VTCM scratchpads for all tensors
- // N rows per thread, padded to HVX vector size
- // Double buffering requires 2x size per buffer
-
- size_t spad_size_per_row = 0;
- size_t vtcm_row_per_thread = 0;
-
- if (octx->op == HTP_OP_RMS_NORM_MUL) {
- if (broadcast_weight) {
- size_t available_vtcm = octx->ctx->vtcm_size;
- size_t src1_spad_total = n_threads * src1_row_size_aligned;
- if (available_vtcm > src1_spad_total) {
- available_vtcm -= src1_spad_total;
- } else {
- available_vtcm = 0;
- }
- spad_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
- vtcm_row_per_thread = available_vtcm / (n_threads * spad_size_per_row);
- } else {
- spad_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned + src1_row_size_aligned);
- vtcm_row_per_thread = (octx->ctx->vtcm_size) / (n_threads * spad_size_per_row);
- }
- } else {
- spad_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
- vtcm_row_per_thread = (octx->ctx->vtcm_size)/ (n_threads * spad_size_per_row);
}
- // Make sure the reserved vtcm size is sufficient
- if (vtcm_row_per_thread == 0) {
- FARF(ERROR, "unary-%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size,
- spad_size_per_row * n_threads);
+ if (octx->ctx->vtcm_size < (size_t)kparams->vtcm_size) {
+ FARF(ERROR, "unary-%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, (size_t)kparams->vtcm_size);
return HTP_STATUS_VTCM_TOO_SMALL;
}
- octx->src0_spad.size_per_thread = src0_row_size_aligned * vtcm_row_per_thread * 2;
- octx->dst_spad.size_per_thread = dst_row_size_aligned * vtcm_row_per_thread * 2;
-
- octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
- octx->dst_spad.size = n_threads * octx->dst_spad.size_per_thread;
-
- if (octx->op == HTP_OP_RMS_NORM_MUL) {
- if (broadcast_weight) {
- octx->src1_spad.size_per_thread = src1_row_size_aligned;
- } else {
- octx->src1_spad.size_per_thread = src1_row_size_aligned * vtcm_row_per_thread * 2;
- }
- octx->src1_spad.size = n_threads * octx->src1_spad.size_per_thread;
- } else {
- octx->src1_spad.size = 0;
- octx->src1_spad.size_per_thread = 0;
- }
-
- octx->src0_spad.data = octx->ctx->vtcm_base;
- if (octx->op == HTP_OP_RMS_NORM_MUL) {
- octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
- octx->dst_spad.data = octx->src1_spad.data + octx->src1_spad.size;
- } else {
- octx->dst_spad.data = octx->src0_spad.data + octx->src0_spad.size;
- }
-
octx->src0_spad.src = NULL;
octx->src1_spad.src = NULL;
octx->dst_spad.src = NULL;
- FARF(HIGH, "%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n", op_type,
+ FARF(HIGH, "%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %u src1-vtcm-size %u dst-vtcm-size %u\n", op_type,
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
- octx->src0_spad.size, octx->src1_spad.size, octx->dst_spad.size);
+ kparams->vtcm_src0_size, kparams->vtcm_src1_size, kparams->vtcm_dst_size);
if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
+ uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
struct htp_unary_context uctx = {
.octx = octx,
+ .kparams = kparams,
.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads,
.src0_nrows = src0_nrows,
.src1_row_size_aligned = src1_row_size_aligned,
.dst_row_size_aligned = dst_row_size_aligned,
- .src0_spad_half_size = octx->src0_spad.size_per_thread / 2,
- .src1_spad_half_size = (octx->op == HTP_OP_RMS_NORM_MUL) ? (octx->src1_spad.size_per_thread / (broadcast_weight ? 1 : 2)) : 0,
- .dst_spad_half_size = octx->dst_spad.size_per_thread / 2,
+ .src0_vtcm_half_size = kparams->vtcm_src0_size_per_thread / 2,
+ .src1_vtcm_half_size = (octx->op == HTP_OP_RMS_NORM_MUL) ? (kparams->vtcm_src1_size_per_thread / (broadcast_weight ? 1 : 2)) : 0,
+ .dst_vtcm_half_size = kparams->vtcm_dst_size_per_thread / 2,
- .block = (octx->src0_spad.size_per_thread / 2) / src0_row_size_aligned,
+ .block = kparams->block,
.nc = src0->ne[0],
+ .col_tile = (uint32_t) kparams->col_tile,
.broadcast_weight = broadcast_weight,
- };
- worker_pool_run_func(octx->ctx->worker_pool, unary_job_f32_per_thread, &uctx, n_threads);
- }
-
- return err;
-}
+ .vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, 0),
+ .vtcm_src1 = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, kparams->vtcm_src0_size, kparams->vtcm_src1_size > 0),
+ .vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, kparams->vtcm_src0_size + kparams->vtcm_src1_size),
-int op_tri(struct htp_ops_context * octx) {
- int err = HTP_STATUS_OK;
+ .vtcm_src0_size_per_thread = kparams->vtcm_src0_size_per_thread,
+ .vtcm_src1_size_per_thread = kparams->vtcm_src1_size_per_thread,
+ .vtcm_dst_size_per_thread = kparams->vtcm_dst_size_per_thread,
+ };
- switch (octx->src[0]->type) {
- case HTP_TYPE_F32:
- err = execute_op_unary_f32(octx);
- break;
+ FARF(HIGH, "%s: %s mode (col_tile %u)\n", op_type, col_tile ? "tiled" : "row-block", col_tile);
+
+ worker_callback_t task_func = NULL;
+ if (col_tile) {
+ switch (octx->op) {
+ case HTP_OP_SCALE: task_func = unary_task_f32_tiled_scale; break;
+ case HTP_OP_CLAMP: task_func = unary_task_f32_tiled_clamp; break;
+ case HTP_OP_SQR: task_func = unary_task_f32_tiled_sqr; break;
+ case HTP_OP_SQRT: task_func = unary_task_f32_tiled_sqrt; break;
+ case HTP_OP_UNARY_NEG: task_func = unary_task_f32_tiled_unary_neg; break;
+ case HTP_OP_UNARY_EXP: task_func = unary_task_f32_tiled_unary_exp; break;
+ case HTP_OP_UNARY_SIGMOID: task_func = unary_task_f32_tiled_unary_sigmoid; break;
+ case HTP_OP_UNARY_SILU: task_func = unary_task_f32_tiled_unary_silu; break;
+ case HTP_OP_UNARY_GELU: task_func = unary_task_f32_tiled_unary_gelu; break;
+ case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_tiled_unary_softplus; break;
+ case HTP_OP_UNARY_TANH: task_func = unary_task_f32_tiled_unary_tanh; break;
+ case HTP_OP_TRI: task_func = unary_task_f32_tiled_tri; break;
+ default: break;
+ }
+ } else {
+ switch (octx->op) {
+ case HTP_OP_NORM: task_func = unary_task_f32_norm; break;
+ case HTP_OP_RMS_NORM: task_func = unary_task_f32_rms_norm; break;
+ case HTP_OP_RMS_NORM_MUL: task_func = unary_task_f32_rms_norm_mul; break;
+ case HTP_OP_SCALE: task_func = unary_task_f32_scale; break;
+ case HTP_OP_CLAMP: task_func = unary_task_f32_clamp; break;
+ case HTP_OP_SQR: task_func = unary_task_f32_sqr; break;
+ case HTP_OP_SQRT: task_func = unary_task_f32_sqrt; break;
+ case HTP_OP_UNARY_NEG: task_func = unary_task_f32_unary_neg; break;
+ case HTP_OP_UNARY_EXP: task_func = unary_task_f32_unary_exp; break;
+ case HTP_OP_UNARY_SIGMOID: task_func = unary_task_f32_unary_sigmoid; break;
+ case HTP_OP_UNARY_SILU: task_func = unary_task_f32_unary_silu; break;
+ case HTP_OP_UNARY_GELU: task_func = unary_task_f32_unary_gelu; break;
+ case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_unary_softplus; break;
+ case HTP_OP_UNARY_TANH: task_func = unary_task_f32_unary_tanh; break;
+ case HTP_OP_L2_NORM: task_func = unary_task_f32_l2_norm; break;
+ case HTP_OP_TRI: task_func = unary_task_f32_tri; break;
+ default: break;
+ }
+ }
- default:
+ if (task_func) {
+ worker_pool_run_func(octx->ctx->worker_pool, task_func, &uctx, n_threads);
+ } else {
+ FARF(ERROR, "execute_op_unary_f32: task function is NULL for op %d\n", octx->op);
err = HTP_STATUS_NO_SUPPORT;
- break;
+ }
}
return err;
}
int op_unary(struct htp_ops_context * octx) {
- int err = HTP_STATUS_OK;
-
switch (octx->src[0]->type) {
case HTP_TYPE_F32:
- err = execute_op_unary_f32(octx);
- break;
+ return execute_op_unary_f32(octx);
default:
- err = HTP_STATUS_NO_SUPPORT;
- break;
+ return HTP_STATUS_NO_SUPPORT;
}
-
- return err;
}
--- /dev/null
+#ifndef HTP_UNARY_OPS_H
+#define HTP_UNARY_OPS_H
+
+#include "hex-common.h"
+#include "htp-ops.h"
+
+// Op-specific struct for precomputed unary params
+struct htp_unary_kernel_params {
+ uint32_t n_threads;
+ uint32_t col_tile;
+ uint32_t vtcm_row_per_thread;
+ uint32_t block;
+ uint32_t broadcast_weight;
+
+ uint32_t vtcm_src0_size_per_thread;
+ uint32_t vtcm_src1_size_per_thread;
+ uint32_t vtcm_dst_size_per_thread;
+
+ uint32_t vtcm_src0_size;
+ uint32_t vtcm_src1_size;
+ uint32_t vtcm_dst_size;
+
+ uint32_t src0_row_size_aligned;
+ uint32_t src1_row_size_aligned;
+ uint32_t dst_row_size_aligned;
+
+ uint32_t vtcm_size;
+
+ // Fastdiv helpers
+ struct fastdiv_values div_ne01;
+ struct fastdiv_values div_ne02;
+ struct fastdiv_values div_ne012;
+ struct fastdiv_values div_tpr;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
+#endif
+
+static inline bool htp_op_is_unary(uint32_t opcode) {
+ switch (opcode) {
+ case HTP_OP_CLAMP:
+ case HTP_OP_NORM:
+ case HTP_OP_RMS_NORM:
+ case HTP_OP_RMS_NORM_MUL:
+ case HTP_OP_SCALE:
+ case HTP_OP_SQR:
+ case HTP_OP_SQRT:
+ case HTP_OP_UNARY_NEG:
+ case HTP_OP_UNARY_EXP:
+ case HTP_OP_UNARY_SIGMOID:
+ case HTP_OP_UNARY_SILU:
+ case HTP_OP_UNARY_GELU:
+ case HTP_OP_UNARY_SOFTPLUS:
+ case HTP_OP_UNARY_TANH:
+ case HTP_OP_L2_NORM:
+ case HTP_OP_TRI:
+ return true;
+ default:
+ return false;
+ }
+}
+
+struct htp_unary_vtcm_layout {
+ size_t total_bytes;
+ size_t off_src0;
+ size_t off_src1;
+ size_t off_dst;
+
+ size_t src0_bytes;
+ size_t src1_bytes;
+ size_t dst_bytes;
+};
+
+static inline void htp_unary_vtcm_layout_build(
+ struct htp_unary_vtcm_layout * L,
+ uint32_t op,
+ uint32_t ne00,
+ uint32_t ne10,
+ uint32_t ne11,
+ bool broadcast_weight,
+ uint32_t n_threads,
+ size_t vtcm_size,
+ uint32_t * out_col_tile,
+ uint32_t * out_vtcm_row_per_thread
+) {
+ const size_t src0_data_row_size = ne00 * sizeof(float);
+ const size_t dst_data_row_size = ne10 * sizeof(float);
+
+ const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
+ const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
+
+ size_t src1_row_size_aligned = 0;
+ if (op == HTP_OP_RMS_NORM_MUL) {
+ const size_t src1_data_row_size = ne11 * sizeof(float);
+ src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
+ }
+
+ size_t vtcm_size_per_row = 0;
+ size_t vtcm_row_per_thread = 0;
+
+ if (op == HTP_OP_RMS_NORM_MUL) {
+ if (broadcast_weight) {
+ size_t available_vtcm = vtcm_size;
+ size_t src1_vtcm_total = n_threads * src1_row_size_aligned;
+ if (available_vtcm > src1_vtcm_total) {
+ available_vtcm -= src1_vtcm_total;
+ } else {
+ available_vtcm = 0;
+ }
+ vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
+ vtcm_row_per_thread = available_vtcm / (n_threads * vtcm_size_per_row);
+ } else {
+ vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned + src1_row_size_aligned);
+ vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
+ }
+ } else {
+ vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
+ vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
+ }
+
+ const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
+ op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM);
+ uint32_t col_tile = 0;
+
+ if (vtcm_row_per_thread == 0 && !is_reduction) {
+ const size_t per_thread_budget = vtcm_size / n_threads;
+ const size_t col_tile_bytes = hex_align_down(per_thread_budget / 4, 128);
+ col_tile = (uint32_t) (col_tile_bytes / sizeof(float));
+
+ L->src0_bytes = col_tile_bytes * 2;
+ L->dst_bytes = col_tile_bytes * 2;
+ L->src1_bytes = 0;
+ } else {
+ L->src0_bytes = src0_row_size_aligned * vtcm_row_per_thread * 2;
+ L->dst_bytes = dst_row_size_aligned * vtcm_row_per_thread * 2;
+ if (op == HTP_OP_RMS_NORM_MUL) {
+ if (broadcast_weight) {
+ L->src1_bytes = src1_row_size_aligned;
+ } else {
+ L->src1_bytes = src1_row_size_aligned * vtcm_row_per_thread * 2;
+ }
+ } else {
+ L->src1_bytes = 0;
+ }
+ }
+
+ L->off_src0 = 0;
+ if (op == HTP_OP_RMS_NORM_MUL) {
+ L->off_src1 = L->off_src0 + L->src0_bytes * n_threads;
+ L->off_dst = L->off_src1 + L->src1_bytes * n_threads;
+ } else {
+ L->off_src1 = 0;
+ L->off_dst = L->off_src0 + L->src0_bytes * n_threads;
+ }
+
+ L->total_bytes = L->off_dst + L->dst_bytes * n_threads;
+
+ *out_col_tile = col_tile;
+ *out_vtcm_row_per_thread = vtcm_row_per_thread;
+}
+
+#endif /* HTP_UNARY_OPS_H */
--- /dev/null
+#include "work-queue.h"
+#include "hex-utils.h"
+
+#include <qurt.h>
+#include <qurt_hvx.h>
+
+#include <stdatomic.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "HAP_farf.h"
+
+#define LOWEST_USABLE_QURT_PRIO (254)
+
+// internal structure kept in thread-local storage per instance of work queue
+typedef struct {
+ work_queue_t queue;
+ unsigned int id;
+} worker_context_t;
+
+struct work_queue_task_s {
+ work_queue_func_t func;
+ void * data;
+ unsigned int n_threads;
+ atomic_uint barrier;
+};
+
+// internal structure kept in thread-local storage per instance of work queue
+struct work_queue_s {
+ atomic_uint seqn; // seqno used to detect new jobs
+ atomic_uint idx_read; // Updated by producer (pop/reclaim)
+ unsigned int idx_write; // Updated by producer (push)
+ uint32_t idx_mask;
+ uint32_t capacity;
+
+ qurt_thread_t thread[WORK_QUEUE_MAX_N_THREADS]; // thread ID's of the workers
+ worker_context_t context[WORK_QUEUE_MAX_N_THREADS]; // worker contexts
+ void * stack[WORK_QUEUE_MAX_N_THREADS]; // thread stack pointers
+ unsigned int n_threads; // total threads (workers + main)
+ unsigned int n_workers; // number of active threads (just workers)
+
+ atomic_bool active; // workers are polling/active
+ atomic_bool killed; // threads need to exit
+ bool external_mem; // memory owned externally
+
+ struct work_queue_task_s queue[] __attribute__((aligned(HEX_L2_LINE_SIZE)));
+};
+
+static void work_queue_thread(void * context) {
+ worker_context_t * me = (worker_context_t *) context;
+ work_queue_t q = me->queue;
+
+ FARF(HIGH, "work-queue: thread %u started", me->id);
+
+ unsigned int prev_seqn = 0;
+
+ while (!atomic_load_explicit(&q->killed, memory_order_relaxed)) {
+ unsigned int seqn = atomic_load_explicit(&q->seqn, memory_order_acquire);
+ if (seqn == prev_seqn) {
+ if (atomic_load_explicit(&q->active, memory_order_relaxed)) {
+ hex_pause();
+ } else {
+ qurt_futex_wait(&q->seqn, prev_seqn);
+ }
+ continue;
+ }
+
+ prev_seqn = seqn;
+
+ // Process all active tasks in the queue
+ unsigned int ir = atomic_load_explicit(&q->idx_read, memory_order_relaxed);
+ unsigned int iw = q->idx_write;
+
+ while (ir != iw) {
+ struct work_queue_task_s * task = &q->queue[ir];
+
+ unsigned int n = task->n_threads;
+ unsigned int i = me->id;
+ if (i < n) {
+ task->func(n, i, task->data);
+
+ atomic_fetch_sub_explicit(&task->barrier, 1, memory_order_release);
+ } else {
+ while (atomic_load_explicit(&task->barrier, memory_order_relaxed) > 0) {
+ hex_pause();
+ }
+ }
+
+ ir = (ir + 1) & q->idx_mask;
+ }
+ }
+
+ FARF(HIGH, "work-queue: thread %u stopped", me->id);
+}
+
+bool work_queue_run_async(work_queue_t q, work_queue_func_t func, void * data, unsigned int n) {
+ if (n > q->n_threads) {
+ FARF(ERROR, "work-queue: invalid number of jobs %u for n-threads %u", n, q->n_threads);
+ return false;
+ }
+
+ unsigned int ir = atomic_load_explicit(&q->idx_read, memory_order_relaxed);
+ unsigned int iw = q->idx_write;
+
+ if (((iw + 1) & q->idx_mask) == ir) {
+ FARF(ERROR, "work-queue-push: queue is full\n");
+ return false;
+ }
+
+ struct work_queue_task_s * task = &q->queue[iw];
+ task->func = func;
+ task->data = data;
+ task->n_threads = n;
+ atomic_store_explicit(&task->barrier, n, memory_order_relaxed);
+
+ q->idx_write = (iw + 1) & q->idx_mask;
+
+ // publish job to workers (already awake and polling)
+ atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
+
+ // main thread runs job #0
+ func(n, 0, data);
+
+ atomic_fetch_sub_explicit(&task->barrier, 1, memory_order_release);
+
+ while (atomic_load_explicit(&task->barrier, memory_order_relaxed) > 0) {
+ hex_pause();
+ }
+
+ atomic_thread_fence(memory_order_acquire);
+
+ atomic_store_explicit(&q->idx_read, (ir + 1) & q->idx_mask, memory_order_relaxed);
+
+ return true;
+}
+
+size_t work_queue_sizeof(uint32_t n_threads, uint32_t capacity, uint32_t stack_size) {
+ capacity = hex_ceil_pow2(capacity);
+ uint32_t n_workers = n_threads > 1 ? n_threads - 1 : 0;
+ size_t size_stacks = stack_size * n_workers;
+ size_t size_q = hex_align_up(sizeof(struct work_queue_s) + capacity * sizeof(struct work_queue_task_s), HEX_L2_LINE_SIZE);
+ return size_stacks + size_q;
+}
+
+size_t work_queue_alignof(void) {
+ return 4096;
+}
+
+work_queue_t work_queue_init(void * ptr, uint32_t n_threads, uint32_t capacity, uint32_t stack_size) {
+ capacity = hex_ceil_pow2(capacity);
+ uint32_t n_workers = n_threads > 1 ? n_threads - 1 : 0;
+ unsigned char * mem_blob = (unsigned char *) ptr;
+
+ work_queue_t q = (work_queue_t) (mem_blob + stack_size * n_workers);
+ memset(q, 0, sizeof(struct work_queue_s) + capacity * sizeof(struct work_queue_task_s));
+
+ q->n_threads = n_threads;
+ q->n_workers = n_workers;
+ q->external_mem = true;
+ q->capacity = capacity;
+
+ for (unsigned int i = 0; i < n_workers; i++) {
+ q->stack[i] = mem_blob; mem_blob += stack_size;
+ q->thread[i] = 0;
+ q->context[i].id = i + 1;
+ q->context[i].queue = q;
+ }
+
+ atomic_init(&q->idx_read, 0);
+ atomic_init(&q->seqn, 0);
+ atomic_init(&q->active, false);
+ q->idx_write = 0;
+ q->idx_mask = capacity - 1;
+ q->killed = 0;
+ for (int i = 0; i < (int) capacity; i++) {
+ atomic_init(&q->queue[i].barrier, 0);
+ q->queue[i].func = NULL;
+ q->queue[i].data = NULL;
+ q->queue[i].n_threads = 0;
+ }
+
+ // launch the workers
+ qurt_thread_attr_t attr;
+ qurt_thread_attr_init(&attr);
+
+ for (unsigned int i = 0; i < n_workers; i++) {
+ qurt_thread_attr_set_stack_addr(&attr, q->stack[i]);
+ qurt_thread_attr_set_stack_size(&attr, stack_size);
+
+ char thread_name[32];
+ snprintf(thread_name, sizeof(thread_name), "work-queue:%u", i);
+ qurt_thread_attr_set_name(&attr, thread_name);
+
+ // set up priority - by default, match the creating thread's prio
+ int prio = qurt_thread_get_priority(qurt_thread_get_id());
+ if (prio < 1) {
+ prio = 1;
+ }
+ if (prio > LOWEST_USABLE_QURT_PRIO) {
+ prio = LOWEST_USABLE_QURT_PRIO;
+ }
+
+ qurt_thread_attr_set_priority(&attr, prio);
+
+ int err = qurt_thread_create(&q->thread[i], &attr, work_queue_thread, (void *) &q->context[i]);
+ if (err) {
+ FARF(ERROR, "Could not launch worker threads!");
+ work_queue_free(q);
+ return NULL;
+ }
+ }
+
+ return q;
+}
+
+void work_queue_free(work_queue_t q) {
+ if (!q) { return; }
+
+ atomic_store_explicit(&q->killed, 1, memory_order_relaxed);
+ atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
+ qurt_futex_wake(&q->seqn, q->n_workers);
+
+ for (unsigned int i = 0; i < q->n_workers; i++) {
+ if (q->thread[i]) {
+ int status;
+ (void) qurt_thread_join(q->thread[i], &status);
+ }
+ }
+}
+
+void work_queue_wakeup(work_queue_t q) {
+ if (!atomic_load_explicit(&q->active, memory_order_relaxed)) {
+ atomic_store_explicit(&q->active, true, memory_order_release);
+ // Increment seqn and wake workers to transition them out of sleep
+ atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
+ qurt_futex_wake(&q->seqn, q->n_workers);
+ }
+}
+
+void work_queue_suspend(work_queue_t q) {
+ atomic_store_explicit(&q->active, false, memory_order_release);
+}
--- /dev/null
+#ifndef HTP_WORK_QUEUE_H
+#define HTP_WORK_QUEUE_H
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <stddef.h>
+
+typedef void (*work_queue_func_t)(unsigned int n, unsigned int i, void *);
+
+struct work_queue_s;
+typedef struct work_queue_s * work_queue_t;
+
+#define WORK_QUEUE_MAX_N_THREADS 10
+
+size_t work_queue_sizeof(uint32_t n_threads, uint32_t capacity, uint32_t stack_size);
+size_t work_queue_alignof(void);
+work_queue_t work_queue_init(void * ptr, uint32_t n_threads, uint32_t capacity, uint32_t stack_size);
+void work_queue_free(work_queue_t q);
+
+void work_queue_wakeup(work_queue_t q);
+void work_queue_suspend(work_queue_t q);
+
+bool work_queue_run_async(work_queue_t q, work_queue_func_t func, void * data, unsigned int n);
+
+static inline bool work_queue_run(work_queue_t q, work_queue_func_t func, void * data, unsigned int n) {
+ if (n <= 1) {
+ func(n, 0, data);
+ return true;
+ }
+ return work_queue_run_async(q, func, data, n);
+}
+
+// Legacy compatibility
+typedef work_queue_func_t worker_callback_t;
+#define worker_pool_run_func work_queue_run
+#define worker_pool work_queue
+
+#endif // #ifndef HTP_WORK_QUEUE_H
+++ /dev/null
-#include "worker-pool.h"
-#include "hex-utils.h"
-
-#include <qurt.h>
-#include <qurt_hvx.h>
-
-#include <stdatomic.h>
-#include <stdint.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include "HAP_farf.h"
-
-#define LOWEST_USABLE_QURT_PRIO (254)
-
-struct worker_pool_s;
-
-// internal structure kept in thread-local storage per instance of worker pool
-typedef struct {
- struct worker_pool_s * pool;
- unsigned int id;
-} worker_context_t;
-
-// internal structure kept in thread-local storage per instance of worker pool
-typedef struct worker_pool_s {
- worker_pool_job_t job[MAX_NUM_WORKERS]; // list of job descriptors
- qurt_thread_t thread[MAX_NUM_WORKERS]; // thread ID's of the workers
- worker_context_t context[MAX_NUM_WORKERS]; // worker contexts
- void * stack[MAX_NUM_WORKERS]; // thread stack pointers
- unsigned int n_threads; // number of workers in this pool
-
- atomic_uint seqn; // seqno used to detect new jobs
- atomic_uint next_job; // next job index
- atomic_uint n_pending; // number of pending jobs
- atomic_uint n_jobs; // number of current jobs
- atomic_bool killed; // threads need to exit
-} worker_pool_t;
-
-static void worker_pool_main(void * context) {
- worker_context_t * me = (worker_context_t *) context;
- worker_pool_t * pool = me->pool;
-
- FARF(HIGH, "worker-pool: thread %u started", me->id);
-
- unsigned int prev_seqn = 0;
- unsigned int poll_cnt = WORKER_POOL_POLL_COUNT;
- while (!atomic_load(&pool->killed)) {
- unsigned int seqn = atomic_load(&pool->seqn);
- if (seqn == prev_seqn) {
- // drop HVX context while spinning
- if (poll_cnt > 1 && poll_cnt == WORKER_POOL_POLL_COUNT) {
- qurt_hvx_unlock();
- }
- if (--poll_cnt) {
- hex_pause();
- continue;
- }
- qurt_futex_wait(&pool->seqn, prev_seqn);
- poll_cnt = WORKER_POOL_POLL_COUNT;
- continue;
- }
-
- prev_seqn = seqn;
- poll_cnt = WORKER_POOL_POLL_COUNT;
-
- // New job
- unsigned int n = atomic_load(&pool->n_jobs);
- unsigned int i = atomic_fetch_add(&pool->next_job, 1);
- if (i >= n) {
- // Spurious wakeup
- continue;
- }
-
- pool->job[i].func(n, i, pool->job[i].data);
-
- atomic_fetch_sub(&pool->n_pending, 1);
- }
-
- FARF(HIGH, "worker-pool: thread %u stopped", me->id);
-}
-
-AEEResult worker_pool_init_with_stack_size(worker_pool_context_t * context, uint32_t n_threads, uint32_t stack_size) {
- int err = 0;
-
- if (NULL == context) {
- FARF(ERROR, "NULL context passed to worker_pool_init().");
- return AEE_EBADPARM;
- }
-
- // Allocations
- int size = (stack_size * n_threads) + (sizeof(worker_pool_t));
-
- unsigned char * mem_blob = (unsigned char *) malloc(size);
- if (!mem_blob) {
- FARF(ERROR, "Could not allocate memory for worker pool!!");
- return AEE_ENOMEMORY;
- }
-
- worker_pool_t * me = (worker_pool_t *) (mem_blob + stack_size * n_threads);
-
- // name for the first worker, useful in debugging threads
- char name[19];
- snprintf(name, 12, "0x%8x:", (int) me);
- strcat(name, "worker0");
- me->n_threads = n_threads;
-
- // initializations
- for (unsigned int i = 0; i < me->n_threads; i++) {
- me->stack[i] = NULL;
- me->thread[i] = 0;
-
- me->context[i].id = i;
- me->context[i].pool = me;
- }
-
- // initialize job queue
- me->n_pending = 0;
- me->n_jobs = 0;
- me->next_job = 0;
- me->seqn = 0;
- me->killed = 0;
-
- // launch the workers
- qurt_thread_attr_t attr;
- qurt_thread_attr_init(&attr);
-
- for (unsigned int i = 0; i < me->n_threads; i++) {
- // set up stack
- me->stack[i] = mem_blob;
- mem_blob += stack_size;
- qurt_thread_attr_set_stack_addr(&attr, me->stack[i]);
- qurt_thread_attr_set_stack_size(&attr, stack_size);
-
- // set up name
- qurt_thread_attr_set_name(&attr, name);
- name[17] = (name[17] + 1);
- // name threads context:worker0, context:worker1, .. (recycle at 9, but num threads should be less than that anyway)
- if (name[17] > '9') {
- name[17] = '0';
- }
-
- // set up priority - by default, match the creating thread's prio
- int prio = qurt_thread_get_priority(qurt_thread_get_id());
-
- if (prio < 1) {
- prio = 1;
- }
- if (prio > LOWEST_USABLE_QURT_PRIO) {
- prio = LOWEST_USABLE_QURT_PRIO;
- }
-
- qurt_thread_attr_set_priority(&attr, prio);
-
- // launch
- err = qurt_thread_create(&me->thread[i], &attr, worker_pool_main, (void *) &me->context[i]);
- if (err) {
- FARF(ERROR, "Could not launch worker threads!");
- worker_pool_release((worker_pool_context_t *) &me);
- return AEE_EQURTTHREADCREATE;
- }
- }
- *context = (worker_pool_context_t *) me;
- return AEE_SUCCESS;
-}
-
-AEEResult worker_pool_init(worker_pool_context_t * context, uint32_t n_threads) {
- return worker_pool_init_with_stack_size(context, n_threads, WORKER_THREAD_STACK_SZ);
-}
-
-// clean up worker pool
-void worker_pool_release(worker_pool_context_t * context) {
- worker_pool_t * me = (worker_pool_t *) *context;
-
- // if no worker pool exists, return error.
- if (NULL == me) {
- return;
- }
-
- atomic_store(&me->killed, 1);
- atomic_fetch_add(&me->seqn, 1);
- qurt_futex_wake(&me->seqn, me->n_threads);
-
- // de-initializations
- for (unsigned int i = 0; i < me->n_threads; i++) {
- if (me->thread[i]) {
- int status;
- (void) qurt_thread_join(me->thread[i], &status);
- }
- }
-
- // free allocated memory (were allocated as a single buffer starting at stack[0])
- if (me->stack[0]) {
- free(me->stack[0]);
- }
-
- *context = NULL;
-}
-
-// run jobs
-AEEResult worker_pool_run_jobs(worker_pool_context_t context, worker_pool_job_t * job, unsigned int n) {
- worker_pool_t * me = (worker_pool_t *) context;
- if (NULL == me) {
- FARF(ERROR, "worker-pool: invalid context");
- return AEE_EBADPARM;
- }
-
- if (n > me->n_threads) {
- FARF(ERROR, "worker-pool: invalid number of jobs %u for n-threads %u", n, me->n_threads);
- return AEE_EBADPARM;
- }
-
- memcpy(me->job, job, sizeof(worker_pool_job_t) * n);
-
- if (n > 1) {
- atomic_store(&me->next_job, 1);
- atomic_store(&me->n_jobs, n);
- atomic_store(&me->n_pending, n - 1);
-
- // wake up workers
- atomic_fetch_add(&me->seqn, 1);
- qurt_futex_wake(&me->seqn, n - 1);
- }
-
- // main thread runs job #0
- me->job[0].func(n, 0, me->job[0].data);
-
- if (n > 1) {
- while (atomic_load(&me->n_pending))
- ;
- }
-
- return 0;
-}
-
-// run func
-AEEResult worker_pool_run_func(worker_pool_context_t context, worker_callback_t func, void * data, unsigned int n) {
- worker_pool_job_t job[n];
-
- for (unsigned int i = 0; i < n; i++) {
- job[i].func = func;
- job[i].data = data;
- }
-
- return worker_pool_run_jobs(context, job, n);
-}
-
-AEEResult worker_pool_set_thread_priority(worker_pool_context_t context, unsigned int prio) {
- worker_pool_t * me = (worker_pool_t *) context;
-
- // if no worker pool exists, return error.
- if (!me) {
- return AEE_ENOMORE;
- }
-
- int result = AEE_SUCCESS;
- if (prio < 1) {
- prio = 1;
- }
- if (prio > LOWEST_USABLE_QURT_PRIO) {
- prio = LOWEST_USABLE_QURT_PRIO;
- }
-
- for (unsigned int i = 0; i < me->n_threads; i++) {
- int res = qurt_thread_set_priority(me->thread[i], (unsigned short) prio);
- if (0 != res) {
- result = AEE_EBADPARM;
- FARF(ERROR, "QURT failed to set priority of thread %d, ERROR = %d", me->thread[i], res);
- }
- }
-
- return result;
-}
-
-AEEResult worker_pool_retrieve_thread_id(worker_pool_context_t context, unsigned int * tids) {
- worker_pool_t * me = (worker_pool_t *) context;
- if (!me) {
- FARF(ERROR, "worker-pool: invalid context");
- return AEE_EBADPARM;
- ;
- }
-
- for (int i = 0; i < me->n_threads; i++) {
- tids[i] = me->thread[i];
- }
-
- return AEE_SUCCESS;
-}
-
-AEEResult worker_pool_get_thread_priority(worker_pool_context_t context, unsigned int * prio) {
- worker_pool_t * me = (worker_pool_t *) context;
- if (!me) {
- FARF(ERROR, "worker-pool: invalid context");
- return AEE_EBADPARM;
- }
-
- int priority = qurt_thread_get_priority(me->thread[0]);
- if (priority > 0) {
- *prio = priority;
- return 0;
- } else {
- *prio = 0;
- return AEE_EBADSTATE;
- }
-}
+++ /dev/null
-#ifndef HTP_WORKER_POOL_H
-#define HTP_WORKER_POOL_H
-
-// MACRO enables function to be visible in shared-library case.
-#define WORKERPOOL_API __attribute__((visibility("default")))
-
-#include <AEEStdDef.h>
-#include <AEEStdErr.h>
-#include <stdint.h>
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-/// signature of callbacks to be invoked by worker threads
-typedef void (*worker_callback_t)(unsigned int n, unsigned int i, void *);
-
-/// Typedef of worker_pool context
-typedef void * worker_pool_context_t;
-
-/// descriptor for requested callback
-typedef struct {
- worker_callback_t func;
- void * data;
-} worker_pool_job_t;
-
-#define WORKER_THREAD_STACK_SZ (2 * 16384)
-
-/// Maximum supported number of worker threads.
-#define MAX_NUM_WORKERS 10
-
-#if __HVX_ARCH__ > 79
-#define WORKER_POOL_POLL_COUNT 2000
-#else
-#define WORKER_POOL_POLL_COUNT 1
-#endif
-
-// Initialize worker pool.
-WORKERPOOL_API AEEResult worker_pool_init(worker_pool_context_t * context, uint32_t n_threads);
-
-// Initialize worker pool with custom stack size
-WORKERPOOL_API AEEResult worker_pool_init_with_stack_size(worker_pool_context_t * context,
- uint32_t n_threads,
- uint32_t stack_size);
-
-// Kill worker threads and release worker pool resources
-WORKERPOOL_API void worker_pool_release(worker_pool_context_t * context);
-
-// Run jobs with the worker pool.
-WORKERPOOL_API AEEResult worker_pool_run_jobs(worker_pool_context_t context, worker_pool_job_t * job, unsigned int n);
-
-WORKERPOOL_API AEEResult worker_pool_run_func(worker_pool_context_t context,
- worker_callback_t func,
- void * data,
- unsigned int n);
-
-WORKERPOOL_API AEEResult worker_pool_set_thread_priority(worker_pool_context_t context, unsigned int prio);
-WORKERPOOL_API AEEResult worker_pool_get_thread_priority(worker_pool_context_t context, unsigned int * prio);
-WORKERPOOL_API AEEResult worker_pool_retrieve_thread_id(worker_pool_context_t context, unsigned int * tids);
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif // #ifndef HTP_WORKER_POOL_H
add_compile_definitions(GGML_HIP_NO_VMM)
endif()
-if (GGML_HIP_ROCWMMA_FATTN)
- add_compile_definitions(GGML_HIP_ROCWMMA_FATTN)
-endif()
-
if (NOT GGML_HIP_MMQ_MFMA)
add_compile_definitions(GGML_HIP_NO_MMQ_MFMA)
endif()
endif()
target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas)
-
-target_compile_options(ggml-hip PRIVATE "$<$<COMPILE_LANGUAGE:HIP>:-ffast-math;-fno-finite-math-only>")
nsg = N_SG_Q1_0;
nr0 = N_R0_Q1_0;
} break;
+ case GGML_TYPE_Q2_0:
+ {
+ nsg = N_SG_Q2_0;
+ nr0 = N_R0_Q2_0;
+ } break;
case GGML_TYPE_Q4_0:
{
nsg = N_SG_Q4_0;
nsg = N_SG_Q1_0;
nr0 = N_R0_Q1_0;
} break;
+ case GGML_TYPE_Q2_0:
+ {
+ nsg = N_SG_Q2_0;
+ nr0 = N_R0_Q2_0;
+ } break;
case GGML_TYPE_Q4_0:
{
nsg = N_SG_Q4_0;
return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_library_t lib, int n) {
+ char base[256];
+ char name[256];
+
+ snprintf(base, 256, "kernel_fwht_f32_%d", n);
+ snprintf(name, 256, "%s", base);
+
+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+ }
+
+ return res;
+}
+
// note: reuse the argsort kernel for top_k
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_TOP_K);
return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake(ggml_metal_library_t lib, enum ggml_type type) {
+ GGML_ASSERT(type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16);
+
+ char base[256];
+ char name[256];
+
+ snprintf(base, 256, "kernel_snake_%s", ggml_type_name(type));
+ snprintf(name, 256, "%s", base);
+
+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+ }
+
+ return res;
+}
+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_CONV_TRANSPOSE_2D);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake (ggml_metal_library_t lib, enum ggml_type type);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tiled);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d (ggml_metal_library_t lib, const struct ggml_tensor * op);
// a collection of residency sets (non-owning)
typedef struct ggml_metal_rsets * ggml_metal_rsets_t;
-ggml_metal_rsets_t ggml_metal_rsets_init(void);
+ggml_metal_rsets_t ggml_metal_rsets_init(ggml_metal_device_t dev);
void ggml_metal_rsets_free(ggml_metal_rsets_t rsets);
//
dispatch_group_t d_group;
};
-ggml_metal_rsets_t ggml_metal_rsets_init(void) {
+#if defined(GGML_METAL_HAS_RESIDENCY_SETS)
+static void ggml_metal_dummy_work(ggml_metal_device_t dev) {
+ if (dev->mtl_queue == nil) {
+ return;
+ }
+
+ @autoreleasepool {
+ // perform a minimal dummy operation on the GPU
+ id<MTLBuffer> buf = [dev->mtl_device newBufferWithLength:1 options:MTLResourceStorageModePrivate];
+ id<MTLCommandBuffer> cmd_buf = [dev->mtl_queue commandBuffer];
+
+ {
+ id<MTLBlitCommandEncoder> encoder = [cmd_buf blitCommandEncoder];
+
+ [encoder fillBuffer:buf range:NSMakeRange(0, 1) value:0];
+
+ [encoder endEncoding];
+ }
+
+ [cmd_buf commit];
+ [buf release];
+ }
+}
+#endif
+
+ggml_metal_rsets_t ggml_metal_rsets_init(ggml_metal_device_t dev) {
ggml_metal_rsets_t res = calloc(1, sizeof(struct ggml_metal_rsets));
res->lock = [[NSLock alloc] init];
#endif
});
+#if defined(GGML_METAL_HAS_RESIDENCY_SETS)
+ if (@available(macOS 15.0, iOS 18.0, tvOS 18.0, visionOS 2.0, *)) {
+ // workaround for residency set memory not being released if no GPU operation occurs
+ // https://developer.apple.com/forums/thread/839089
+ // https://github.com/ggml-org/llama.cpp/issues/25937
+ ggml_metal_dummy_work(dev);
+ }
+#endif
+
return res;
}
}
if (dev->props.use_residency_sets) {
- dev->rsets = ggml_metal_rsets_init();
+ dev->rsets = ggml_metal_rsets_init(dev);
} else {
dev->rsets = nil;
}
(ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0);
case GGML_OP_PAD_REFLECT_1D:
case GGML_OP_TIMESTEP_EMBEDDING:
- case GGML_OP_LEAKY_RELU:
return op->src[0]->type == GGML_TYPE_F32;
+ case GGML_OP_LEAKY_RELU:
+ return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
case GGML_OP_ARGSORT:
case GGML_OP_TOP_K:
case GGML_OP_ARANGE:
case GGML_TYPE_BF16:
case GGML_TYPE_Q8_0:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
return false;
}
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
return op->src[0]->type != GGML_TYPE_NVFP4;
case GGML_OP_SET_ROWS:
{
- if (op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) {
+ if (op->src[0]->type == GGML_TYPE_F16) {
+ return op->type == GGML_TYPE_F16;
+ }
+
+ if (op->src[0]->type != GGML_TYPE_F32) {
return false;
}
if (buf->rset) {
[buf->rset endResidency];
[buf->rset removeAllAllocations];
+ [buf->rset commit];
[buf->rset release];
}
}
#define N_R0_Q1_0 8
#define N_SG_Q1_0 2
+#define N_R0_Q2_0 8
+#define N_SG_Q2_0 2
+
#define N_R0_Q4_0 4
#define N_SG_Q4_0 2
int32_t p0;
} ggml_metal_kargs_col2im_1d;
+typedef struct {
+ int32_t T;
+ int32_t C;
+} ggml_metal_kargs_snake;
+
typedef struct {
int32_t IC;
int32_t IH;
int32_t len;
} ggml_metal_kargs_argsort_merge;
+typedef struct {
+ int32_t nrows;
+} ggml_metal_kargs_fwht;
+
typedef struct {
int64_t ne0;
float start;
return 1;
}
+// supported FWHT sizes, must stay in sync with the
+// kernel_fwht_f32_<N> templates in ggml-metal.metal
+static bool ggml_metal_fwht_supported_size(int64_t n) {
+ return n == 64 || n == 128 || n == 256 || n == 512;
+}
+
+int ggml_metal_op_fwht(ggml_metal_op_t ctx, int idx) {
+ ggml_tensor * op = ctx->node(idx);
+
+ ggml_metal_library_t lib = ctx->lib;
+ ggml_metal_encoder_t enc = ctx->enc;
+
+ ggml_tensor * src1 = op->src[1];
+
+ const int64_t n = src1->ne[0];
+ const int64_t nrows = ggml_nrows(src1);
+
+ ggml_metal_kargs_fwht args = {
+ /*.nrows = */ (int32_t) nrows,
+ };
+
+ auto pipeline = ggml_metal_library_get_pipeline_fwht(lib, n);
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline);
+ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(src1), 1);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2);
+
+ const int th_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
+ const int simd_size = 32;
+
+ int sg_per_tg = 2;
+ sg_per_tg = std::min(sg_per_tg, th_max/simd_size);
+ sg_per_tg = std::max(sg_per_tg, 1);
+
+ const int64_t n_tg = (nrows + sg_per_tg - 1) / sg_per_tg;
+ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, 32*sg_per_tg, 1, 1);
+
+ return 1;
+}
int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
ggml_metal_library_t lib = ctx->lib;
ggml_metal_encoder_t enc = ctx->enc;
+ const int32_t hint = ggml_get_op_params_i32(op, 1);
+
+ if (hint == GGML_HINT_SRC0_IS_HADAMARD) {
+ if (op->src[1]->type == GGML_TYPE_F32 &&
+ op->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op->src[1]) &&
+ ggml_is_contiguous(op) &&
+ ggml_are_same_shape(op->src[1], op) &&
+ ggml_metal_fwht_supported_size(op->src[1]->ne[0])) {
+ return ggml_metal_op_fwht(ctx, idx);
+ }
+ }
const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
op->src[0]->type == GGML_TYPE_F16 ||
op->src[0]->type == GGML_TYPE_BF16 ||
op->src[0]->type == GGML_TYPE_Q1_0 ||
+ op->src[0]->type == GGML_TYPE_Q2_0 ||
op->src[0]->type == GGML_TYPE_Q4_0 ||
op->src[0]->type == GGML_TYPE_Q4_1 ||
op->src[0]->type == GGML_TYPE_Q5_0 ||
return 1;
}
+// Snake activation autofuse: mul -> sin -> sqr -> mul -> add
+static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) {
+ static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
+
+ if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) {
+ return false;
+ }
+
+ const ggml_tensor * mul0 = ctx->node(idx + 0);
+ const ggml_tensor * sin_node = ctx->node(idx + 1);
+ const ggml_tensor * sqr = ctx->node(idx + 2);
+ const ggml_tensor * mul1 = ctx->node(idx + 3);
+ const ggml_tensor * add = ctx->node(idx + 4);
+
+ // x carries the full activation shape, a is the broadcast operand
+ const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+ const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+
+ // mul1 reads sqr and inv_b in either operand order
+ const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+ // closure check: the trailing add reads the same x as the leading mul
+ const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
+
+ // x is in the supported whitelist and every chain intermediate shares x's type.
+ // a and inv_b bind as device const float * in the kernel, so they stay F32.
+ const bool types_ok =
+ (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
+ (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) &&
+ (mul0->type == x->type) && (sin_node->type == x->type) &&
+ (sqr->type == x->type) && (mul1->type == x->type) &&
+ (add->type == x->type);
+ // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
+ const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
+ const bool dim_ok =
+ (x->ne[2] == 1) && (x->ne[3] == 1) &&
+ (add->ne[2] == 1) && (add->ne[3] == 1) &&
+ (a->ne[2] == 1) && (a->ne[3] == 1) &&
+ (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
+ // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
+ const bool contig_ok =
+ ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
+ ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
+
+ return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;
+}
+
int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
+ if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) {
+ return ggml_metal_op_snake_fused(ctx, idx);
+ }
+
ggml_tensor * op = ctx->node(idx);
ggml_metal_library_t lib = ctx->lib;
return 1;
}
+// Dispatch the fused snake kernel from the matched mul -> sin -> sqr -> mul -> add chain.
+// idx points at the leading mul. The caller has validated the chain.
+int ggml_metal_op_snake_fused(ggml_metal_op_t ctx, int idx) {
+ ggml_metal_library_t lib = ctx->lib;
+ ggml_metal_encoder_t enc = ctx->enc;
+
+ const ggml_tensor * mul0 = ctx->node(idx + 0);
+ const ggml_tensor * sqr = ctx->node(idx + 2);
+ const ggml_tensor * mul1 = ctx->node(idx + 3);
+ ggml_tensor * add = ctx->node(idx + 4);
+
+ const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+ const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+ const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+ const int T = (int) x->ne[0];
+ const int C = (int) x->ne[1];
+ const int total = T * C;
+
+ // the encode loop pre-checked the leading mul only, check the rest of the chain
+ for (int i = 1; i < 5; ++i) {
+ if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) {
+ ggml_metal_op_concurrency_reset(ctx);
+
+ break;
+ }
+ }
+
+ auto pipeline = ggml_metal_library_get_pipeline_snake(lib, x->type);
+
+ ggml_metal_kargs_snake args = {
+ /*.T =*/ T,
+ /*.C =*/ C,
+ };
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline);
+ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(x), 1);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(a), 2);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(inv_b), 3);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(add), 4);
+
+ const int nth = 256;
+ const int ntg = (total + nth - 1) / nth;
+ ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
+
+ return 5;
+}
+
int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
int ggml_metal_op_cpy (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_pool_1d (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_pool_2d (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_fwht (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_mul_mat (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_col2im_1d (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_snake_fused (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_upscale (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_pad (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_pad_reflect_1d (ggml_metal_op_t ctx, int idx);
reg = (type4) reg_f;
}
+template <typename type4x4>
+void dequantize_q2_0(device const block_q2_0 * xb, short il, thread type4x4 & reg) {
+ device const uint8_t * qs = xb->qs;
+ const float d = xb->d;
+
+ const int byte_offset = il * 4; // il*16 elements = il*4 bytes (4 elements per byte)
+ float4x4 reg_f;
+
+ for (int i = 0; i < 4; i++) {
+ const uint8_t b = qs[byte_offset + i];
+ reg_f[i][0] = ((float)((b >> 0) & 3) - 1.0f) * d;
+ reg_f[i][1] = ((float)((b >> 2) & 3) - 1.0f) * d;
+ reg_f[i][2] = ((float)((b >> 4) & 3) - 1.0f) * d;
+ reg_f[i][3] = ((float)((b >> 6) & 3) - 1.0f) * d;
+ }
+
+ reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q2_0_t4(device const block_q2_0 * xb, short il, thread type4 & reg) {
+ const float d = xb->d;
+ const uint8_t b = xb->qs[il];
+
+ float4 reg_f;
+ reg_f[0] = ((float)((b >> 0) & 3) - 1.0f) * d;
+ reg_f[1] = ((float)((b >> 2) & 3) - 1.0f) * d;
+ reg_f[2] = ((float)((b >> 4) & 3) - 1.0f) * d;
+ reg_f[3] = ((float)((b >> 6) & 3) - 1.0f) * d;
+
+ reg = (type4) reg_f;
+}
+
template <typename type4x4>
void dequantize_q4_0(device const block_q4_0 * xb, short il, thread type4x4 & reg) {
device const uint16_t * qs = ((device const uint16_t *)xb + 1);
}
}
+void quantize_q2_0(device const float * src, device block_q2_0 & dst) {
+ float amax = 0.0f;
+ for (int j = 0; j < QK2_0; j++) {
+ float a = fabs(src[j]);
+ if (a > amax) amax = a;
+ }
+ const float d = amax;
+ dst.d = d;
+
+ const float id = d > 0.0f ? 1.0f / d : 0.0f;
+
+ for (int j = 0; j < QK2_0 / 4; j++) {
+ dst.qs[j] = 0;
+ }
+ for (int j = 0; j < QK2_0; j++) {
+ int q = (int)round(src[j] * id) + 1;
+ q = max(0, min(3, q));
+ dst.qs[j / 4] |= (q << (2 * (j % 4)));
+ }
+}
+
void quantize_q4_0(device const float * src, device block_q4_0 & dst) {
#pragma METAL fp math_mode(safe)
float amax = 0.0f; // absolute max
return qb_curr->d * (2.0f * acc - sumy);
}
+// Q2_0 dot: d * (sum_lo(y) + 2*sum_hi(y) - sumy) via per-bit conditional adds
+inline float block_q_n_dot_y(device const block_q2_0 * qb_curr, float sumy, thread float * yl, int il) {
+ device const uint8_t * qs = qb_curr->qs + (il / 4);
+ const uint8_t b0 = qs[0];
+ const uint8_t b1 = qs[1];
+ const uint8_t b2 = qs[2];
+ const uint8_t b3 = qs[3];
+
+ // Accumulate where low bit is set (bits 0,2,4,6 of each byte)
+ float acc_lo = 0.0f;
+ acc_lo += select(0.0f, yl[ 0], bool(b0 & 0x01));
+ acc_lo += select(0.0f, yl[ 1], bool(b0 & 0x04));
+ acc_lo += select(0.0f, yl[ 2], bool(b0 & 0x10));
+ acc_lo += select(0.0f, yl[ 3], bool(b0 & 0x40));
+ acc_lo += select(0.0f, yl[ 4], bool(b1 & 0x01));
+ acc_lo += select(0.0f, yl[ 5], bool(b1 & 0x04));
+ acc_lo += select(0.0f, yl[ 6], bool(b1 & 0x10));
+ acc_lo += select(0.0f, yl[ 7], bool(b1 & 0x40));
+ acc_lo += select(0.0f, yl[ 8], bool(b2 & 0x01));
+ acc_lo += select(0.0f, yl[ 9], bool(b2 & 0x04));
+ acc_lo += select(0.0f, yl[10], bool(b2 & 0x10));
+ acc_lo += select(0.0f, yl[11], bool(b2 & 0x40));
+ acc_lo += select(0.0f, yl[12], bool(b3 & 0x01));
+ acc_lo += select(0.0f, yl[13], bool(b3 & 0x04));
+ acc_lo += select(0.0f, yl[14], bool(b3 & 0x10));
+ acc_lo += select(0.0f, yl[15], bool(b3 & 0x40));
+
+ // Accumulate where high bit is set (bits 1,3,5,7 of each byte)
+ float acc_hi = 0.0f;
+ acc_hi += select(0.0f, yl[ 0], bool(b0 & 0x02));
+ acc_hi += select(0.0f, yl[ 1], bool(b0 & 0x08));
+ acc_hi += select(0.0f, yl[ 2], bool(b0 & 0x20));
+ acc_hi += select(0.0f, yl[ 3], bool(b0 & 0x80));
+ acc_hi += select(0.0f, yl[ 4], bool(b1 & 0x02));
+ acc_hi += select(0.0f, yl[ 5], bool(b1 & 0x08));
+ acc_hi += select(0.0f, yl[ 6], bool(b1 & 0x20));
+ acc_hi += select(0.0f, yl[ 7], bool(b1 & 0x80));
+ acc_hi += select(0.0f, yl[ 8], bool(b2 & 0x02));
+ acc_hi += select(0.0f, yl[ 9], bool(b2 & 0x08));
+ acc_hi += select(0.0f, yl[10], bool(b2 & 0x20));
+ acc_hi += select(0.0f, yl[11], bool(b2 & 0x80));
+ acc_hi += select(0.0f, yl[12], bool(b3 & 0x02));
+ acc_hi += select(0.0f, yl[13], bool(b3 & 0x08));
+ acc_hi += select(0.0f, yl[14], bool(b3 & 0x20));
+ acc_hi += select(0.0f, yl[15], bool(b3 & 0x80));
+
+ return qb_curr->d * (acc_lo + 2.0f * acc_hi - sumy);
+}
+
// function for calculate inner product between half a q4_0 block and 16 floats (yl), sumy is SUM(yl[i])
// il indicates where the q4 quants begin (0 or QK4_0/4)
// we assume that the yl's have been multiplied with the appropriate scale factor
kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
}
+template<int nr0, typename args_t>
+void kernel_mul_mv_q2_0_f32_impl(
+ args_t args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ threadgroup char * shmem,
+ uint3 tgpig,
+ ushort tiisg,
+ ushort sgitg) {
+ const short NSG = FC_mul_mv_nsg;
+
+ const int nb = args.ne00/QK2_0;
+
+ const int r0 = tgpig.x;
+ const int r1 = tgpig.y;
+ const int im = tgpig.z;
+
+ const int first_row = (r0 * NSG + sgitg) * nr0;
+
+ const uint i12 = im%FC_mul_mv_ne12;
+ const uint i13 = im/FC_mul_mv_ne12;
+
+ const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
+
+ device const float * y = (device const float *) (src1 + offset1);
+
+ device const block_q2_0 * ax[nr0];
+ for (int row = 0; row < nr0; ++row) {
+ const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+ ax[row] = (device const block_q2_0 *) ((device char *) src0 + offset0);
+ }
+
+ float yl[16];
+ float sumf[nr0] = {0.f};
+
+ // group 64: 4 sub-blocks of 16 weights per Q2_0 block
+ const short ix = (tiisg/4);
+ const short il = (tiisg%4)*16;
+
+ device const float * yb = y + ix*QK2_0 + il;
+
+ for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/4) {
+ float sumy = 0.f;
+
+ FOR_UNROLL (short i = 0; i < 16; i++) {
+ yl[i] = yb[i];
+ sumy += yb[i];
+ }
+
+ FOR_UNROLL (short row = 0; row < nr0; row++) {
+ sumf[row] += block_q_n_dot_y(ax[row] + ib, sumy, yl, il);
+ }
+
+ yb += QK2_0 * (N_SIMDWIDTH/4);
+ }
+
+ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
+
+ for (int row = 0; row < nr0; ++row) {
+ const float tot = simd_sum(sumf[row]);
+
+ if (tiisg == 0 && first_row + row < args.ne01) {
+ dst_f32[first_row + row] = tot;
+ }
+ }
+}
+
+[[host_name("kernel_mul_mv_q2_0_f32")]]
+kernel void kernel_mul_mv_q2_0_f32(
+ constant ggml_metal_kargs_mul_mv & args,
+ device const char * src0,
+ device const char * src1,
+ device char * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort tiisg[[thread_index_in_simdgroup]],
+ ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+ kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+}
+
kernel void kernel_mul_mv_q4_0_f32(
constant ggml_metal_kargs_mul_mv & args,
device const char * src0,
template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q1_0, 128, dequantize_q1_0_t4>;
template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q1_0, 128, dequantize_q1_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q2_0, 64, dequantize_q2_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q2_0, 64, dequantize_q2_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q2_0, 64, dequantize_q2_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q2_0, 64, dequantize_q2_0_t4>;
+
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q4_0, 32, dequantize_q4_0_t4>;
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q4_0, 32, dequantize_q4_0_t4>;
template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q4_0, 32, dequantize_q4_0_t4>;
#endif
+template <typename T>
+kernel void kernel_snake(
+ constant ggml_metal_kargs_snake & args,
+ device const T * x,
+ device const float * a,
+ device const float * inv_b,
+ device T * dst,
+ uint tgpig [[threadgroup_position_in_grid]],
+ uint tpitg [[thread_position_in_threadgroup]],
+ uint ntg [[threads_per_threadgroup]]) {
+
+ const int idx = tgpig * ntg + tpitg;
+ if (idx >= args.T * args.C) {
+ return;
+ }
+
+ const int c = idx / args.T; // x is [T, C], a / inv_b collapse to [1, C]
+ const float xi = float(x[idx]);
+ const float si = sin(a[c] * xi);
+ dst[idx] = T(xi + si * si * inv_b[c]);
+}
+
+template [[host_name("kernel_snake_f32")]] kernel void kernel_snake<float>(constant ggml_metal_kargs_snake &, device const float *, device const float *, device const float *, device float *, uint, uint, uint);
+template [[host_name("kernel_snake_f16")]] kernel void kernel_snake<half>(constant ggml_metal_kargs_snake &, device const half *, device const float *, device const float *, device half *, uint, uint, uint);
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(constant ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
+#endif
+
+
typedef void (conv_transpose_2d_t)(
constant ggml_metal_kargs_conv_transpose_2d & args,
device const float * src0,
const float w_y2 = bicubic_weight1(1.0f - fd1);
const float w_y3 = bicubic_weight2(2.0f - fd1);
- const device const char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
+ const device char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1);
template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
+template<int N>
+kernel void kernel_fwht_f32(
+ constant ggml_metal_kargs_fwht & args,
+ device const float * src,
+ device float * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort sgitg[[simdgroup_index_in_threadgroup]],
+ ushort tiisg[[thread_index_in_simdgroup]],
+ ushort3 ntg[[threads_per_threadgroup]]) {
+
+ constexpr int NW = N_SIMDWIDTH;
+ constexpr int NE = N / NW;
+
+ const float scale = 1.0f / sqrt((float) N);
+
+ const int sg_per_tg = ntg.x / NW;
+ const int64_t r = tgpig.x * sg_per_tg + sgitg;
+ if (r >= args.nrows) {
+ return;
+ }
+
+ src += r * N;
+ dst += r * N;
+
+ const int lane = tiisg;
+
+ float reg[NE];
+ for (int i = 0; i < NE; i++) {
+ reg[i] = src[i*NW + lane]*scale;
+ }
+ for (int i = 1; i < NW; i *= 2) {
+ for (int j = 0; j < NE; j++) {
+ const float val = reg[j];
+ const float val2 = simd_shuffle_xor(val, i);
+ reg[j] = (lane & i) == 0 ? val2 + val : val2 - val;
+ }
+ }
+
+ for (int i = NW; i < N; i *= 2) {
+ const int step = i / NW;
+ for (int j = 0; j < NE; j += (2 * step)) {
+ for (int k = 0; k < step; k++) {
+ const float x = reg[j + k ];
+ const float y = reg[j + k + step];
+ reg[j + k] = x + y;
+ reg[j + k + step] = x - y;
+ }
+ }
+ }
+
+ for (int i = 0; i < NE; i++) {
+ dst[i*NW + lane] = reg[i];
+ }
+}
+
+typedef decltype(kernel_fwht_f32<64>) kernel_fwht_t;
+
+template [[host_name("kernel_fwht_f32_64")]] kernel kernel_fwht_t kernel_fwht_f32<64>;
+template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f32<128>;
+template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>;
+template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>;
+
constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]];
constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];
template [[host_name("kernel_cpy_f32_q8_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK8_0, block_q8_0, quantize_q8_0>;
template [[host_name("kernel_cpy_f32_q1_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK1_0, block_q1_0, quantize_q1_0>;
+template [[host_name("kernel_cpy_f32_q2_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK2_0, block_q2_0, quantize_q2_0>;
template [[host_name("kernel_cpy_f32_q4_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_0, block_q4_0, quantize_q4_0>;
template [[host_name("kernel_cpy_f32_q4_1")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_1, block_q4_1, quantize_q4_1>;
template [[host_name("kernel_cpy_f32_q5_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0, block_q5_0, quantize_q5_0>;
typedef decltype(kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>) cpy_q_f_t;
template [[host_name("kernel_cpy_q1_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_cpy_q2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q2_0, 4, dequantize_q2_0>;
template [[host_name("kernel_cpy_q4_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>;
template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_1, 2, dequantize_q4_1>;
template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_0, 2, dequantize_q5_0>;
template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_cpy_q2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q2_0, 4, dequantize_q2_0>;
template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
template [[host_name("kernel_cpy_q4_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_1, 2, dequantize_q4_1>;
template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_0, 2, dequantize_q5_0>;
typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_get_rows_q2_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q2_0, 4, dequantize_q2_0>;
template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>;
template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_1, 2, dequantize_q4_1>;
template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_0, 2, dequantize_q5_0>;
template [[host_name("kernel_mul_mm_bf16_f32")]] kernel mul_mm_t kernel_mul_mm<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4, 1, dequantize_bf16, bfloat, bfloat4x4, float, float2x4>;
#endif
template [[host_name("kernel_mul_mm_q1_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, float, float2x4>;
+template [[host_name("kernel_mul_mm_q2_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_q4_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_q4_1_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_q5_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_f32_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_f16_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_q1_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, half, half2x4>;
+template [[host_name("kernel_mul_mm_q2_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_q4_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_q4_1_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_q5_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_bf16_f32")]] kernel mul_mm_id kernel_mul_mm_id<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4, 1, dequantize_bf16, bfloat, bfloat4x4, float, float2x4>;
#endif
template [[host_name("kernel_mul_mm_id_q1_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q2_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_q4_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_q4_1_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_q5_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_f32_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_f16_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_q1_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q2_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_q4_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_q4_1_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_q5_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mv_id_q8_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0>>>;
template [[host_name("kernel_mul_mv_id_q1_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0>>>;
+template [[host_name("kernel_mul_mv_id_q2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0>>>;
template [[host_name("kernel_mul_mv_id_q4_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_0, N_R0_Q4_0>>>;
template [[host_name("kernel_mul_mv_id_q4_1_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_1, N_R0_Q4_1>>>;
template [[host_name("kernel_mul_mv_id_q5_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q5_0, N_R0_Q5_0>>>;
ggml_add_backend_library(${TARGET_NAME}
ggml-opencl.cpp
+ cl-program-cache.cpp
+ cl-program-cache.h
../../include/ggml-opencl.h)
target_link_libraries(${TARGET_NAME} PRIVATE ${OpenCL_LIBRARIES})
target_include_directories(${TARGET_NAME} PRIVATE ${OpenCL_INCLUDE_DIRS})
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
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
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
tanh
exp
expm1
+ abs
softplus
pad
repeat
--- /dev/null
+// Match the version setup ggml-opencl.cpp uses, so any cl.h declarations we
+// touch are consistent across this backend's translation units.
+#define CL_TARGET_OPENCL_VERSION GGML_OPENCL_TARGET_VERSION
+#define CL_USE_DEPRECATED_OPENCL_1_2_APIS
+
+#include "cl-program-cache.h"
+
+#include "ggml-impl.h" // GGML_LOG_INFO / WARN
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <system_error>
+#include <vector>
+
+#if defined(_WIN32)
+# ifndef WIN32_LEAN_AND_MEAN
+# define WIN32_LEAN_AND_MEAN
+# endif
+# ifndef NOMINMAX
+# define NOMINMAX
+# endif
+# include <windows.h>
+# include <process.h>
+# define ggml_getpid() ((int) GetCurrentProcessId())
+#else
+# include <unistd.h>
+# define ggml_getpid() ((int) getpid())
+#endif
+
+namespace fs = std::filesystem;
+
+// ----------------------------------------------------------------------------
+// SHA-256 (FIPS 180-4). Self-contained, ~80 lines, public-domain reference.
+// Hot path is a few KB of source per kernel ⇒ <1 ms total per process init.
+// ----------------------------------------------------------------------------
+
+namespace {
+
+struct sha256_ctx {
+ uint32_t state[8];
+ uint64_t bitlen;
+ uint8_t buf[64];
+ size_t buf_len;
+};
+
+const uint32_t K256[64] = {
+ 0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
+ 0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
+ 0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
+ 0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
+ 0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
+ 0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
+ 0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
+ 0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2,
+};
+
+inline uint32_t rotr32(uint32_t x, unsigned n) { return (x >> n) | (x << (32 - n)); }
+
+void sha256_compress(uint32_t state[8], const uint8_t block[64]) {
+ uint32_t w[64];
+ for (int i = 0; i < 16; ++i) {
+ w[i] = ((uint32_t)block[i*4 ] << 24) |
+ ((uint32_t)block[i*4 + 1] << 16) |
+ ((uint32_t)block[i*4 + 2] << 8) |
+ ((uint32_t)block[i*4 + 3] );
+ }
+ for (int i = 16; i < 64; ++i) {
+ uint32_t s0 = rotr32(w[i-15], 7) ^ rotr32(w[i-15], 18) ^ (w[i-15] >> 3);
+ uint32_t s1 = rotr32(w[i-2], 17) ^ rotr32(w[i-2], 19) ^ (w[i-2] >> 10);
+ w[i] = w[i-16] + s0 + w[i-7] + s1;
+ }
+
+ uint32_t a = state[0],b = state[1],c = state[2],d = state[3],e = state[4],f = state[5],g = state[6],h = state[7];
+
+ for (int i = 0; i < 64; ++i) {
+ uint32_t S1 = rotr32(e, 6) ^ rotr32(e, 11) ^ rotr32(e, 25);
+ uint32_t ch = (e & f) ^ ((~e) & g);
+ uint32_t t1 = h + S1 + ch + K256[i] + w[i];
+ uint32_t S0 = rotr32(a, 2) ^ rotr32(a, 13) ^ rotr32(a, 22);
+ uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
+ uint32_t t2 = S0 + maj;
+ h = g; g = f; f = e; e = d + t1;
+ d = c; c = b; b = a; a = t1 + t2;
+ }
+ state[0]+=a; state[1]+=b; state[2]+=c; state[3]+=d;
+ state[4]+=e; state[5]+=f; state[6]+=g; state[7]+=h;
+}
+
+void sha256_init(sha256_ctx & c) {
+ c.state[0]=0x6a09e667; c.state[1]=0xbb67ae85; c.state[2]=0x3c6ef372; c.state[3]=0xa54ff53a;
+ c.state[4]=0x510e527f; c.state[5]=0x9b05688c; c.state[6]=0x1f83d9ab; c.state[7]=0x5be0cd19;
+ c.bitlen = 0;
+ c.buf_len = 0;
+}
+
+void sha256_update(sha256_ctx & c, const void * data, size_t len) {
+ const uint8_t * p = (const uint8_t *) data;
+ c.bitlen += (uint64_t) len * 8;
+ if (c.buf_len > 0) {
+ size_t n = 64 - c.buf_len;
+ if (n > len) { n = len; }
+ memcpy(c.buf + c.buf_len, p, n);
+ c.buf_len += n;
+ p += n;
+ len -= n;
+ if (c.buf_len == 64) {
+ sha256_compress(c.state, c.buf);
+ c.buf_len = 0;
+ }
+ }
+ while (len >= 64) {
+ sha256_compress(c.state, p);
+ p += 64;
+ len -= 64;
+ }
+ if (len > 0) {
+ memcpy(c.buf, p, len);
+ c.buf_len = len;
+ }
+}
+
+void sha256_final(sha256_ctx & c, uint8_t out[32]) {
+ uint64_t bitlen = c.bitlen;
+ c.buf[c.buf_len++] = 0x80;
+ if (c.buf_len > 56) {
+ while (c.buf_len < 64) { c.buf[c.buf_len++] = 0; }
+ sha256_compress(c.state, c.buf);
+ c.buf_len = 0;
+ }
+ while (c.buf_len < 56) { c.buf[c.buf_len++] = 0; }
+ for (int i = 7; i >= 0; --i) { c.buf[c.buf_len++] = (uint8_t) (bitlen >> (i * 8)); }
+ sha256_compress(c.state, c.buf);
+ for (int i = 0; i < 8; ++i) {
+ out[i*4 ] = (uint8_t) (c.state[i] >> 24);
+ out[i*4 + 1] = (uint8_t) (c.state[i] >> 16);
+ out[i*4 + 2] = (uint8_t) (c.state[i] >> 8);
+ out[i*4 + 3] = (uint8_t) (c.state[i] );
+ }
+}
+
+std::string sha256_hex(const uint8_t digest[32]) {
+ static const char hex[] = "0123456789abcdef";
+ std::string s(64, '0');
+ for (int i = 0; i < 32; ++i) {
+ s[i*2 ] = hex[digest[i] >> 4];
+ s[i*2 + 1] = hex[digest[i] & 0xf];
+ }
+ return s;
+}
+
+std::string compute_key(const std::string & key_suffix,
+ const char * source,
+ const std::string & compile_opts) {
+ sha256_ctx c;
+ sha256_init(c);
+
+ static const uint8_t sep = 0;
+ sha256_update(c, source, strlen(source));
+ sha256_update(c, &sep, 1);
+ sha256_update(c, compile_opts.data(), compile_opts.size());
+ sha256_update(c, &sep, 1);
+ sha256_update(c, key_suffix.data(), key_suffix.size());
+
+ uint8_t digest[32];
+ sha256_final(c, digest);
+ return sha256_hex(digest);
+}
+
+bool make_dir_recursive(const std::string & path) {
+ if (path.empty()) { return false; }
+ // create_directories() already creates missing parents. It returns false
+ // (with ec clear) when the directory is already there, so re-check.
+ const fs::path p = fs::u8path(path);
+ std::error_code ec;
+ if (fs::create_directories(p, ec)) { return true; }
+ std::error_code ec_stat;
+ return fs::is_directory(p, ec_stat);
+}
+
+std::string default_cache_dir() {
+#if defined(_WIN32)
+ const char * base = std::getenv("LOCALAPPDATA");
+ if (!base || !*base) { base = std::getenv("APPDATA"); }
+ if (!base || !*base) { base = std::getenv("TEMP"); }
+ if (!base || !*base) { base = "."; }
+ return std::string(base) + "\\llama.cpp\\cl-cache";
+#elif defined(__APPLE__)
+ const char * home = std::getenv("HOME");
+ if (!home || !*home) { home = "."; }
+ return std::string(home) + "/Library/Caches/llama.cpp/cl-cache";
+#else
+ // The throwing overload aborts the process when no usable temp directory
+ // exists (e.g. Android app contexts with TMPDIR unset); an empty return
+ // here just disables the cache instead.
+ std::error_code ec;
+ const fs::path tmp_path = fs::temp_directory_path(ec);
+ if (ec || tmp_path.empty()) { return {}; }
+ return tmp_path.string() + "/llama.cpp/cl-cache";
+#endif
+}
+
+// Query a NUL-terminated string from clGetDeviceInfo / clGetPlatformInfo.
+template <typename GetInfoFn, typename Object>
+std::string query_string(GetInfoFn fn, Object obj, cl_uint name) {
+ size_t sz = 0;
+ if (fn(obj, name, 0, nullptr, &sz) != CL_SUCCESS || sz == 0) {
+ return {};
+ }
+ std::string s(sz, '\0');
+ if (fn(obj, name, sz, &s[0], nullptr) != CL_SUCCESS) {
+ return {};
+ }
+ if (!s.empty() && s.back() == '\0') {
+ s.pop_back();
+ }
+ return s;
+}
+
+std::string compute_key_suffix(cl_device_id device) {
+ cl_platform_id platform = nullptr;
+ clGetDeviceInfo(device, CL_DEVICE_PLATFORM, sizeof(platform), &platform, nullptr);
+
+ std::string s;
+ s.reserve(512);
+ s += query_string(clGetDeviceInfo, device, CL_DEVICE_NAME); s.push_back('\0');
+ s += query_string(clGetDeviceInfo, device, CL_DRIVER_VERSION); s.push_back('\0');
+ s += query_string(clGetDeviceInfo, device, CL_DEVICE_VERSION); s.push_back('\0');
+ if (platform) {
+ s += query_string(clGetPlatformInfo, platform, CL_PLATFORM_VERSION); s.push_back('\0');
+ }
+ s += "fmt=" + std::to_string(CL_PROGRAM_CACHE_FORMAT_VERSION);
+ return s;
+}
+
+const uint8_t MAGIC[8] = { 'G','G','M','L','C','L','B','C' };
+
+bool read_all(const std::string & path, std::vector<uint8_t> & out) {
+ std::ifstream f(fs::u8path(path), std::ios::binary);
+ if (!f) { return false; }
+ f.seekg(0, std::ios::end);
+ std::streamsize sz = f.tellg();
+ if (sz < 0) { return false; }
+ f.seekg(0, std::ios::beg);
+ out.resize((size_t) sz);
+ if (sz > 0) { f.read((char *) out.data(), sz); }
+ return f.good() || f.eof();
+}
+
+bool write_atomic(const std::string & path, const uint8_t * data, size_t len) {
+ const fs::path dst = fs::u8path(path);
+ const fs::path tmp = fs::u8path(path + ".tmp." + std::to_string(ggml_getpid()));
+ {
+ std::ofstream f(tmp, std::ios::binary | std::ios::trunc);
+ if (!f) { return false; }
+ f.write((const char *) data, (std::streamsize) len);
+ if (!f.good()) {
+ std::error_code ec_rm;
+ fs::remove(tmp, ec_rm);
+ return false;
+ }
+ }
+
+ std::error_code ec;
+ fs::rename(tmp, dst, ec);
+ if (ec) {
+ std::error_code ec_rm;
+ fs::remove(tmp, ec_rm);
+ return false;
+ }
+ return true;
+}
+
+} // namespace
+
+static bool cache_debug_enabled() {
+ static int cached = -1;
+ if (cached < 0) {
+ const char * e = std::getenv("GGML_OPENCL_KERNEL_CACHE_DEBUG");
+ cached = (e && *e) ? 1 : 0;
+ }
+ return cached != 0;
+}
+
+static std::string opts_preview(const std::string & opts, size_t n = 120) {
+ if (opts.size() <= n) { return opts; }
+ return opts.substr(0, n) + "...";
+}
+
+// Running cache tally (diagnostic; plain ints — a benign race in the rare
+// multi-threaded lazy-compile case at worst miscounts by one).
+static int g_cache_hits = 0, g_cache_misses = 0, g_cache_saves = 0;
+
+// Debug trace directly to stderr
+static void cache_debug_line(const char * kind, const std::string & key,
+ const char * source, const std::string & opts) {
+ if (!cache_debug_enabled()) { return; }
+ fprintf(stderr, "ggml_opencl: cache %-4s [h=%d m=%d s=%d] key=%s src=%zuB opts='%s'\n",
+ kind, g_cache_hits, g_cache_misses, g_cache_saves,
+ key.substr(0, 16).c_str(), strlen(source), opts_preview(opts).c_str());
+ fflush(stderr);
+}
+
+cl_program_cache_state cl_program_cache_init(cl_device_id device) {
+ cl_program_cache_state st;
+
+ const char * env = std::getenv("GGML_OPENCL_KERNEL_CACHE_DIR");
+ if (env && (!std::strcmp(env, "0") || !std::strcmp(env, "off") ||
+ !std::strcmp(env, "none") || !std::strcmp(env, "disable") ||
+ !std::strcmp(env, "disabled"))) {
+ if (cache_debug_enabled()) {
+ fprintf(stderr, "ggml_opencl: kernel cache disabled by GGML_OPENCL_KERNEL_CACHE_DIR=%s\n", env);
+ fflush(stderr);
+ }
+ return st;
+ }
+
+ std::string dir;
+ if (!env || !*env || !std::strcmp(env, "1") || !std::strcmp(env, "default")) {
+ dir = default_cache_dir();
+ if (dir.empty()) {
+ GGML_LOG_INFO("ggml_opencl: kernel cache disabled (no usable default cache directory)\n");
+ return st;
+ }
+ } else {
+ dir = env;
+ }
+
+ if (!make_dir_recursive(dir)) {
+ GGML_LOG_INFO("ggml_opencl: kernel cache disabled (cannot create directory '%s')\n", dir.c_str());
+ return st;
+ }
+
+ st.dir = dir;
+ st.key_suffix = compute_key_suffix(device);
+ GGML_LOG_INFO("ggml_opencl: kernel cache enabled at '%s'\n", st.dir.c_str());
+ if (cache_debug_enabled()) {
+ fprintf(stderr, "ggml_opencl: kernel cache enabled at '%s' "
+ "(GGML_OPENCL_KERNEL_CACHE_DIR=off to disable)\n", st.dir.c_str());
+ fflush(stderr);
+ }
+ return st;
+}
+
+cl_program cl_program_cache_try_load(
+ const cl_program_cache_state & state,
+ cl_context context,
+ cl_device_id device,
+ const char * source,
+ const std::string & compile_opts) {
+
+ if (state.dir.empty() || !source) { return nullptr; }
+
+ const std::string key = compute_key(state.key_suffix, source, compile_opts);
+ const std::string path = state.dir + "/" + key + ".clbin";
+
+ std::vector<uint8_t> file;
+ if (!read_all(path, file)) {
+ ++g_cache_misses;
+ cache_debug_line("MISS", key, source, compile_opts);
+ return nullptr;
+ }
+ if (file.size() < 16 || std::memcmp(file.data(), MAGIC, 8) != 0) { return nullptr; }
+
+ uint32_t fmt =
+ ((uint32_t) file[ 8]) | ((uint32_t) file[ 9] << 8) |
+ ((uint32_t) file[10] << 16) | ((uint32_t) file[11] << 24);
+ if (fmt != CL_PROGRAM_CACHE_FORMAT_VERSION) { return nullptr; }
+
+ const size_t hdr_len = 16;
+ const unsigned char * bin = file.data() + hdr_len;
+ const size_t bin_len = file.size() - hdr_len;
+
+ cl_int err = CL_SUCCESS;
+ cl_int bin_err = CL_SUCCESS;
+ cl_program p = clCreateProgramWithBinary(context, 1, &device, &bin_len, &bin, &bin_err, &err);
+ if (err != CL_SUCCESS || bin_err != CL_SUCCESS || p == nullptr) {
+ if (p) { clReleaseProgram(p); }
+ return nullptr;
+ }
+
+ err = clBuildProgram(p, 0, nullptr, compile_opts.c_str(), nullptr, nullptr);
+ if (err != CL_SUCCESS) {
+ clReleaseProgram(p);
+ return nullptr;
+ }
+ ++g_cache_hits;
+ cache_debug_line("HIT", key, source, compile_opts);
+ return p;
+}
+
+void cl_program_cache_try_save(
+ const cl_program_cache_state & state,
+ cl_program program,
+ cl_device_id /*device*/,
+ const char * source,
+ const std::string & compile_opts) {
+
+ if (state.dir.empty() || !program || !source) {
+ return;
+ }
+
+ cl_uint n_dev = 0;
+ if (clGetProgramInfo(program, CL_PROGRAM_NUM_DEVICES, sizeof(n_dev), &n_dev, nullptr) != CL_SUCCESS || n_dev == 0) {
+ return;
+ }
+
+ std::vector<size_t> sizes(n_dev);
+ if (clGetProgramInfo(program, CL_PROGRAM_BINARY_SIZES, sizeof(size_t) * n_dev, sizes.data(), nullptr) != CL_SUCCESS) {
+ return;
+ }
+ if (sizes.empty() || sizes[0] == 0) {
+ return;
+ }
+
+ std::vector<std::vector<uint8_t>> binaries(n_dev);
+ std::vector<unsigned char *> bin_ptrs(n_dev);
+ for (cl_uint i = 0; i < n_dev; ++i) {
+ binaries[i].resize(sizes[i]);
+ bin_ptrs[i] = binaries[i].data();
+ }
+ if (clGetProgramInfo(program, CL_PROGRAM_BINARIES, sizeof(unsigned char *) * n_dev, bin_ptrs.data(), nullptr) != CL_SUCCESS) {
+ return;
+ }
+
+ // We only care about the first device's binary — that's the one we'd
+ // re-load with on a future cache hit. Multi-device contexts aren't a
+ // pattern this backend uses today.
+ const std::vector<uint8_t> & bin = binaries[0];
+
+ std::vector<uint8_t> file;
+ file.reserve(16 + bin.size());
+ file.insert(file.end(), MAGIC, MAGIC + 8);
+ uint32_t fmt = CL_PROGRAM_CACHE_FORMAT_VERSION;
+ file.push_back((uint8_t) (fmt & 0xff));
+ file.push_back((uint8_t) ((fmt >> 8) & 0xff));
+ file.push_back((uint8_t) ((fmt >> 16) & 0xff));
+ file.push_back((uint8_t) ((fmt >> 24) & 0xff));
+ file.push_back(0); file.push_back(0); file.push_back(0); file.push_back(0); // reserved
+ file.insert(file.end(), bin.begin(), bin.end());
+
+ const std::string key = compute_key(state.key_suffix, source, compile_opts);
+ const std::string path = state.dir + "/" + key + ".clbin";
+ if (!write_atomic(path, file.data(), file.size())) {
+ GGML_LOG_INFO("ggml_opencl: kernel cache: failed to write '%s'\n", path.c_str());
+ } else {
+ ++g_cache_saves;
+ cache_debug_line("SAVE", key, source, compile_opts);
+ }
+}
--- /dev/null
+// On-disk cache for OpenCL cl_program binaries. Lets a fresh process skip the
+// expensive clBuildProgram-from-source step when a binary for the exact same
+// (source, compile options, device, driver, platform) was previously saved.
+//
+// Activation: default on via GGML_OPENCL_KERNEL_CACHE_DIR:
+// unset / empty / "1" / "default" : platform default cache dir
+// (%LOCALAPPDATA%\llama.cpp\cl-cache,
+// ~/Library/Caches/llama.cpp/cl-cache,
+// <temp dir>/llama.cpp/cl-cache elsewhere)
+// "0" / "off" / "none" / "disable(d)" : disabled (all functions no-op)
+// any other value : used verbatim as the cache path
+// If the chosen directory cannot be created/used, the cache silently disables
+// itself for the process and falls back to source compile.
+// GGML_OPENCL_KERNEL_CACHE_DEBUG=1 prints a HIT/MISS/SAVE trace (with a running
+// tally) straight to stderr — visible even in tools that filter INFO/WARN logs;
+// redirect stderr to record it.
+//
+// Cache key (SHA-256 hex):
+// sha256(source_bytes || '\x00' ||
+// compile_opts || '\x00' ||
+// CL_DEVICE_NAME || '\x00' ||
+// CL_DRIVER_VERSION || '\x00' ||
+// CL_PLATFORM_VERSION || '\x00' ||
+// CL_PROGRAM_CACHE_FORMAT_VERSION)
+//
+// The key fully captures everything that can affect the produced binary,
+// without needing the host source revision (a kernel source change shows up
+// in source_bytes; a compile-option change shows up in compile_opts).
+//
+// File layout per cache entry: <cache_dir>/<sha256-hex>.clbin
+// bytes [0..7] : magic "GGMLCLBC"
+// bytes [8..11] : uint32_t format version (CL_PROGRAM_CACHE_FORMAT_VERSION)
+// bytes [12..15] : uint32_t reserved (0)
+// bytes [16..] : raw cl_program binary as returned by
+// clGetProgramInfo(CL_PROGRAM_BINARIES)
+//
+// Concurrency: writes go to <name>.tmp.<pid> then atomic rename. On race,
+// last-writer-wins. No locks.
+
+#pragma once
+
+#include <CL/cl.h>
+#include <string>
+
+// Bumped manually if host-side OpenCL API usage changes in a way that
+// affects compile semantics but does not show up in source_bytes /
+// compile_opts (e.g. switching from clCreateProgramWithSource to
+// clCompileProgram + clLinkProgram, or changing how multiple sources
+// are concatenated). Most commits — including kernel changes — do NOT
+// require bumping this; the source bytes already capture those.
+#define CL_PROGRAM_CACHE_FORMAT_VERSION 1u
+
+struct cl_program_cache_state {
+ // Empty string means cache is disabled.
+ std::string dir;
+ // Concatenated device/driver/platform identity + cache format version,
+ // computed once at init and folded into every key.
+ std::string key_suffix;
+};
+
+cl_program_cache_state cl_program_cache_init(cl_device_id device);
+
+cl_program cl_program_cache_try_load(
+ const cl_program_cache_state & state,
+ cl_context context,
+ cl_device_id device,
+ const char * source,
+ const std::string & compile_opts);
+
+void cl_program_cache_try_save(
+ const cl_program_cache_state & state,
+ cl_program program,
+ cl_device_id device,
+ const char * source,
+ const std::string & compile_opts);
#include "ggml-backend-impl.h"
#include "ggml.h"
+#include "cl-program-cache.h"
+
#ifdef GGML_OPENCL_USE_ADRENO_BIN_KERNELS
#include "libdl.h"
#ifdef _WIN32
enum ADRENO_GPU_GEN {
ADRENO_UNKNOWN,
+ A6X,
A7X,
A8X,
X1E,
enum ADRENO_CL_COMPILER_TYPE {
E031,
+ E17,
DX,
};
}
static ADRENO_GPU_GEN get_adreno_gpu_gen(const char *device_name) {
+ if (strstr(device_name, "610") || strstr(device_name, "612") ||
+ strstr(device_name, "613") || strstr(device_name, "615") ||
+ strstr(device_name, "616") || strstr(device_name, "618") ||
+ strstr(device_name, "619") || strstr(device_name, "620") ||
+ strstr(device_name, "630") || strstr(device_name, "640") ||
+ strstr(device_name, "642") || strstr(device_name, "643") ||
+ strstr(device_name, "644") || strstr(device_name, "650") ||
+ strstr(device_name, "660") || strstr(device_name, "663") ||
+ strstr(device_name, "680") || strstr(device_name, "685") ||
+ strstr(device_name, "690")) {
+ return ADRENO_GPU_GEN::A6X;
+ }
+
if (strstr(device_name, "730") ||
strstr(device_name, "740") ||
strstr(device_name, "750")) {
return ADRENO_GPU_GEN::A7X;
}
- if (strstr(device_name, "830") ||
- strstr(device_name, "840")) {
+ if (strstr(device_name, "810") ||
+ strstr(device_name, "830") ||
+ strstr(device_name, "840") ||
+ strstr(device_name, "850")) {
return ADRENO_GPU_GEN::A8X;
}
size_t compiler_minor_offset = 8;
size_t compiler_patch_offset = 11;
+ if (compiler_ver_pos == std::string::npos) {
+ compiler_ver_pos = driver_ver_str.find("E17");
+ if (compiler_ver_pos != std::string::npos) {
+ type = ADRENO_CL_COMPILER_TYPE::E17;
+ compiler_ver_len = 12;
+ compiler_major_offset = 4;
+ compiler_minor_offset = 7;
+ compiler_patch_offset = 10;
+ }
+ }
+
if (compiler_ver_pos == std::string::npos) {
compiler_ver_pos = driver_ver_str.find("DX");
if (compiler_ver_pos == std::string::npos) {
type = ADRENO_CL_COMPILER_TYPE::DX;
compiler_ver_len = 11;
compiler_major_offset = 3;
+ compiler_minor_offset = 6;
+ compiler_patch_offset = 9;
}
std::string compiler_ver_str = driver_ver_str.substr(compiler_ver_pos, compiler_ver_len);
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_k_img;
// k-image variant of MQ_GQA=4 vec_mq_split
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_k_img;
+ // Cluster-parallel decode
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_c8;
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c8;
+ // NSG_SPLIT=2 specializations (WG=128): the c8 kernel's register footprint
+ // caps its per-kernel WG at 128 on X2, below the stock 256/192 requirement.
+ // 2 subgroups × FA_CL_NCL streams still gives 16 in-flight rows per WG.
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_c8_ns2;
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c8_ns2;
+ // FA_CL_C=32 / MQ_GQA=8 / NSG_SPLIT=2 specialization for the DK=DV=256
+ // GQA=8 class (Qwen3.5/3.6-35B-A3B: 16 Q heads, 2 KV heads). o_acc =
+ // DV_VEC/32 × 8 = 128B/lane (in budget); the baseline fa1 path for this
+ // shape has NO MQ/FD at all and pays an 8× KV re-read per Q head.
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c32;
// alternative decode
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_local_tile;
// hybrid local-tile + MQ + FD-split kernel for DK=DV=128 only
// KV-head-coalesced + flash-decoding split for q8_0 KV
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split;
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split_g8;
+ // Cluster-parallel q8_0 decode
+ std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split_c8;
std::map<std::pair<int, int>, cl_kernel> f32_q8_0; // prefill (baseline)
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_split; // N_SPLIT>1 variant
std::map<std::pair<int, int>, int> f32_q8_0_split_wg_size; // wg_size = bm*n_split
// kv-head-coalesced + flash-decoding split for q4_0 kv (dp4a K dot)
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split;
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_g8;
+ // Cluster-parallel q4_0 decode
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_g8_c8;
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_c8;
std::map<std::pair<int, int>, cl_kernel> f32_q4_0;
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_split;
std::map<std::pair<int, int>, int> f32_q4_0_split_wg_size;
bool fp16_support;
bool has_vector_subgroup_broadcast;
bool has_subgroup_shuffle = false; // cl_khr_subgroup_shuffle or cl_qcom_subgroup_shuffle
+ bool has_integer_dot = false; // cl_khr_integer_dot_product or cl_qcom_dot_product8
bool has_qcom_subgroup_shuffle = false; // specifically cl_qcom_subgroup_shuffle
bool disable_fusion;
// ragged moe, use int to directly pass to kernel
cl_uint adreno_use_moe_ragged;
cl_uint adreno_moe_ragged_skip_gran;
+ cl_uint adreno_use_moe_ragged_dp4;
+
+ // whether fuse moe combine
+ cl_uint fuse_moe_combine;
bool adreno_has_large_buffer;
bool adreno_use_large_buffer;
cl_context context;
cl_command_queue queue;
+ // On-disk compiled-program cache (see GGML_OPENCL_KERNEL_CACHE_DIR).
+ cl_program_cache_state program_cache;
+ bool program_cache_initialized = false;
+
// prealloc buffers for transposing weights and activations
ggml_cl_buffer prealloc_quant_trans;
ggml_cl_buffer prealloc_scales_trans;
ggml_cl_buffer prealloc_act_trans;
+ // q8_1-quantized reordered MoE activations for the dp4a prefill GEMM.
+ ggml_cl_buffer prealloc_moe_qa; // int8 quants [tok_slots * ne00]
+ ggml_cl_buffer prealloc_moe_da; // per-block d [tok_slots * ne00/32] (half)
+ ggml_cl_buffer prealloc_moe_sa; // per-block s [tok_slots * ne00/32] (half)
+ // scratch copy of the router weights to avoid dst aliasing
+ ggml_cl_buffer prealloc_moe_combine_w;
// pool of persistent image1d_buffer views over kv-cache layers, keyed by
// (parent buffer, offset within parent)
cl_kernel kernel_exp_f16, kernel_exp_f16_4, kernel_exp_f16_nc;
cl_kernel kernel_expm1_f32, kernel_expm1_f32_4, kernel_expm1_f32_nc;
cl_kernel kernel_expm1_f16, kernel_expm1_f16_4, kernel_expm1_f16_nc;
+ cl_kernel kernel_abs_f32, kernel_abs_f32_4, kernel_abs_f32_nc;
+ cl_kernel kernel_abs_f16, kernel_abs_f16_4, kernel_abs_f16_nc;
cl_kernel kernel_softplus_f32, kernel_softplus_f32_4, kernel_softplus_f32_nc;
cl_kernel kernel_softplus_f16, kernel_softplus_f16_4, kernel_softplus_f16_nc;
cl_kernel kernel_upscale;
// [size_idx][kda][tgpp] where size_idx: 0=S_V=16, 1=32, 2=64, 3=128; kda: 0 or 1.
// tgpp 0 = TG variant (COLS_PER_LANE_GROUP=1), tgpp 1 = prefill variant (COLS_PER_LANE_GROUP=4).
cl_kernel kernel_gated_delta_net_f32[4][2][2] = {};
-
cl_kernel kernel_timestep_embedding;
cl_kernel kernel_gemv_moe_q4_0_f32_ns, kernel_gemm_moe_q4_0_f32_ns, kernel_gemm_moe_q4_0_f32_ns_bin;
+ cl_kernel kernel_gemm_moe_q8_0_f32_ns;
cl_kernel kernel_gemv_moe_q4_1_f32_ns, kernel_gemm_moe_q4_1_f32_ns, kernel_gemm_moe_q4_1_f32_ns_bin;
cl_kernel kernel_gemv_moe_q5_0_f32_ns, kernel_gemm_moe_q5_0_f32_ns;
cl_kernel kernel_gemv_moe_q5_1_f32_ns, kernel_gemm_moe_q5_1_f32_ns;
cl_kernel kernel_gemv_moe_q4_k_f32_ns, kernel_gemm_moe_q4_k_f32_ns, kernel_gemm_moe_q4_k_f32_ns_bin;
+ cl_kernel kernel_gemv_moe_q4_k_f32_ns_wimg = nullptr; // weight-as-texture MoE decode GEMV (opt-in)
+ cl_kernel kernel_gemm_moe_q4_k_q8_1_dp4a = nullptr; // dp4a (int8) prefill GEMM variant
+ cl_kernel kernel_moe_reorder_quant_a_q8_1; // fused reorder + q8_1 quant for the dp4a GEMM
+ cl_kernel kernel_gemm_moe_q8_1_dp4a_q80 = nullptr; // generic dp4a MoE GEMM (MOE_QT=80), opt-in
+ cl_kernel kernel_moe_expand_scale_q8_0 = nullptr; // q8_0 per-block d -> uniform scale[16]
+ cl_kernel kernel_gemm_moe_q8_1_dp4a_q50 = nullptr; // generic dp4a MoE GEMM (MOE_QT=50, q5_0), opt-in
+ cl_kernel kernel_moe_expand_scale_q5_0 = nullptr; // q5_0 d -> uniform scale[2]/min[1] per 32-block
+ cl_kernel kernel_gemm_moe_q8_1_dp4a_q5k = nullptr; // generic dp4a MoE GEMM (MOE_QT=5, q5_K), opt-in
+ cl_kernel kernel_moe_expand_scale_q5_K = nullptr; // q5_K 6-bit s[] -> uniform scale[16]/min[8]
cl_kernel kernel_gemv_moe_q5_k_f32_ns, kernel_gemm_moe_q5_k_f32_ns;
- cl_kernel kernel_gemv_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns;
+ cl_kernel kernel_gemv_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns_bin;
+ cl_kernel kernel_gemm_moe_q6_k_q8_1_dp4a = nullptr; // dp4a (int8) q6_K MoE prefill GEMM
cl_kernel kernel_gemv_moe_mxfp4_f32, kernel_gemm_moe_mxfp4_f32;
cl_kernel kernel_gemv_moe_mxfp4_f32_ns, kernel_gemm_moe_mxfp4_f32_ns, kernel_gemm_moe_mxfp4_f32_ns_bin;
+ cl_kernel kernel_gemv_moe_mxfp4_f32_ns_wimg = nullptr; // weight-as-texture MoE decode GEMV
+ cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr; // dp4a (int8) mxfp4 MoE prefill GEMM
+ cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr; // dp4a (int8) q4_0 MoE prefill GEMM
cl_kernel kernel_moe_reorder_b;
cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter;
+ cl_kernel kernel_moe_combine_f32 = nullptr; // fused router-weight mul + cross-expert sum
cl_kernel kernel_mul_mv_id_q4_0_f32_8x_flat;
cl_kernel kernel_mul_mv_id_q8_0_f32, kernel_mul_mv_id_q8_0_f32_flat;
cl_kernel kernel_mul_mv_id_mxfp4_f32;
cl_kernel kernel_gemv_noshuffle_q4_1_f32;
cl_kernel kernel_gemm_noshuffle_q4_1_f32;
cl_kernel kernel_gemm_noshuffle_q8_0_f32, kernel_gemm_noshuffle_q8_0_f32_bin;
+ cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a = nullptr; // dp4a (int8) dense q8_0 prefill GEMM (opt-in)
+ cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg = nullptr; // q8_0 dense dp4a, weights via texture (opt-in)
cl_kernel kernel_gemv_noshuffle_q8_0_f32;
cl_kernel kernel_gemm_noshuffle_q1_0_f32;
cl_kernel kernel_gemv_noshuffle_q1_0_f32;
cl_kernel kernel_gemv_noshuffle_q4_k_f32;
cl_kernel kernel_gemm_noshuffle_q4_k_f32;
+ cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a = nullptr; // dp4a (int8) dense prefill GEMM
+ cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = nullptr; // dp4a dense prefill GEMM, weights via texture (X1 opt-in)
+ cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a = nullptr; // dp4a (int8) dense q5_K prefill GEMM
+ cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a = nullptr; // dp4a (int8) dense q6_K prefill GEMM
+ cl_kernel kernel_quant_a_q8_1; // plain activation q8_1 pre-pass
cl_kernel kernel_gemv_noshuffle_q6_K_f32;
cl_kernel kernel_gemm_noshuffle_q6_K_f32;
cl_kernel kernel_gemv_noshuffle_q5_k_f32;
cl_kernel kernel_gemm_noshuffle_q5_k_f32;
cl_kernel kernel_gemv_noshuffle_q5_0_f32;
cl_kernel kernel_gemm_noshuffle_q5_0_f32;
+ cl_kernel kernel_gemm_noshuffle_q5_0_q8_1_dp4a = nullptr; // dp4a (int8) dense q5_0 prefill GEMM
+ cl_kernel kernel_gemm_noshuffle_q5_0_q8_1_dp4a_wimg = nullptr; // q5_0 dense dp4a, qs plane via texture (opt-in)
cl_kernel kernel_gemv_noshuffle_q5_1_f32;
cl_kernel kernel_gemm_noshuffle_q5_1_f32;
cl_kernel kernel_gemv_noshuffle_iq4_nl_f32;
cl_kernel kernel_gemm_noshuffle_iq4_nl_f32;
+ cl_kernel kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a = nullptr; // dp4a (int8) dense IQ4_NL prefill GEMM
+ cl_kernel kernel_gemm_noshuffle_q4_0_q8_1_dp4a = nullptr; // dp4a (int8) dense q4_0 prefill GEMM
#endif // GGML_OPENCL_USE_ADRENO_KERNELS
void free() {
return NULL;
}
-static cl_program build_program_from_source(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts) {
- return build_program_from_source_ex(ctx, dev, program_buffer, compile_opts, /*fatal=*/true);
+static cl_program build_program_from_source(ggml_backend_opencl_context * backend_ctx, const char* program_buffer, const std::string &compile_opts) {
+ cl_context ctx = backend_ctx->context;
+ cl_device_id dev = backend_ctx->device;
+
+ // Try the on-disk binary cache first. Falls through silently on miss or
+ // any failure; never blocks the build path. Disabled cache => nullptr.
+ cl_program p_cached = cl_program_cache_try_load(
+ backend_ctx->program_cache, ctx, dev, program_buffer, compile_opts);
+ if (p_cached != nullptr) {
+ return p_cached;
+ }
+
+ cl_program p = build_program_from_source_ex(ctx, dev, program_buffer, compile_opts, /*fatal=*/true);
+
+ // Best-effort save of the freshly-built binary (no-op if cache disabled).
+ if (p != nullptr) {
+ cl_program_cache_try_save(backend_ctx->program_cache, p, dev, program_buffer, compile_opts);
+ }
+ return p;
}
static cl_program build_program_from_binary(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts, size_t bin_size = 0) {
const std::string kernel_src = read_file("argsort.cl");
#endif
backend_ctx->program_argsort_f32_i32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_argsort_f32_i32 = clCreateKernel(backend_ctx->program_argsort_f32_i32, "kernel_argsort_f32_i32", &err), err));
backend_ctx->kernels_loaded_argsort = true;
const std::string kernel_src = read_file("add.cl");
#endif
backend_ctx->program_add =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_add = clCreateKernel(backend_ctx->program_add, "kernel_add", &err), err));
CL_CHECK((backend_ctx->kernel_add_row = clCreateKernel(backend_ctx->program_add, "kernel_add_row", &err), err));
const std::string kernel_src = read_file("add_id.cl");
#endif
backend_ctx->program_add_id =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_add_id = clCreateKernel(backend_ctx->program_add_id, "kernel_add_id", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("tri.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_tri = clCreateKernel(prog, "kernel_tri_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("fill.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_fill = clCreateKernel(prog, "kernel_fill_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("clamp.cl");
#endif
backend_ctx->program_clamp =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_clamp = clCreateKernel(backend_ctx->program_clamp, "kernel_clamp", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("cpy.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_cpy_f16_f16 = clCreateKernel(prog, "kernel_cpy_f16_f16", &err), err));
CL_CHECK((backend_ctx->kernel_cpy_f16_f32 = clCreateKernel(prog, "kernel_cpy_f16_f32", &err), err));
const std::string kernel_src = read_file("cvt.cl");
#endif
backend_ctx->program_cvt =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_convert_block_q1_0 = clCreateKernel(backend_ctx->program_cvt, "kernel_convert_block_q1_0", &err), err));
CL_CHECK((backend_ctx->kernel_restore_block_q1_0 = clCreateKernel(backend_ctx->program_cvt, "kernel_restore_block_q1_0", &err), err));
CL_CHECK((backend_ctx->kernel_restore_block_iq4_nl_noshuffle = clCreateKernel(backend_ctx->program_cvt, "kernel_restore_block_iq4_nl_noshuffle", &err), err));
CL_CHECK((backend_ctx->kernel_convert_bf16_to_f16 = clCreateKernel(backend_ctx->program_cvt, "kernel_convert_bf16_to_f16", &err), err));
CL_CHECK((backend_ctx->kernel_convert_f16_to_bf16 = clCreateKernel(backend_ctx->program_cvt, "kernel_convert_f16_to_bf16", &err), err));
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ CL_CHECK((backend_ctx->kernel_moe_expand_scale_q8_0 = clCreateKernel(backend_ctx->program_cvt, "kernel_moe_expand_scale_q8_0", &err), err));
+ CL_CHECK((backend_ctx->kernel_moe_expand_scale_q5_0 = clCreateKernel(backend_ctx->program_cvt, "kernel_moe_expand_scale_q5_0", &err), err));
+ CL_CHECK((backend_ctx->kernel_moe_expand_scale_q5_K = clCreateKernel(backend_ctx->program_cvt, "kernel_moe_expand_scale_q5_K", &err), err));
+#endif
GGML_LOG_CONT(".");
}
const std::string kernel_src = read_file("diag_mask_inf.cl");
#endif
backend_ctx->program_diag_mask_inf =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_diag_mask_inf_8 = clCreateKernel(backend_ctx->program_diag_mask_inf, "kernel_diag_mask_inf_8", &err), err));
CL_CHECK((backend_ctx->kernel_diag_mask_inf = clCreateKernel(backend_ctx->program_diag_mask_inf, "kernel_diag_mask_inf", &err), err));
const std::string kernel_src = read_file("diag.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_diag_f32 = clCreateKernel(prog, "kernel_diag_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gelu.cl");
#endif
backend_ctx->program_gelu =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gelu = clCreateKernel(backend_ctx->program_gelu, "kernel_gelu", &err), err));
CL_CHECK((backend_ctx->kernel_gelu_4 = clCreateKernel(backend_ctx->program_gelu, "kernel_gelu_4", &err), err));
const std::string kernel_src = read_file("glu.cl");
#endif
backend_ctx->program_glu =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_geglu = clCreateKernel(backend_ctx->program_glu, "kernel_geglu", &err), err));
CL_CHECK((backend_ctx->kernel_reglu = clCreateKernel(backend_ctx->program_glu, "kernel_reglu", &err), err));
const std::string kernel_src = read_file("get_rows.cl");
#endif
backend_ctx->program_get_rows =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_get_rows_f32 = clCreateKernel(backend_ctx->program_get_rows, "kernel_get_rows_f32", &err), err));
CL_CHECK((backend_ctx->kernel_get_rows_f16 = clCreateKernel(backend_ctx->program_get_rows, "kernel_get_rows_f16", &err), err));
const std::string kernel_src = read_file("solve_tri.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_solve_tri_f32 = clCreateKernel(prog, "kernel_solve_tri_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("im2col_f32.cl");
#endif
backend_ctx->program_im2col_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_im2col_f32 = clCreateKernel(backend_ctx->program_im2col_f32, "kernel_im2col_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("im2col_f16.cl");
#endif
backend_ctx->program_im2col_f16 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_im2col_f16 = clCreateKernel(backend_ctx->program_im2col_f16, "kernel_im2col_f16", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q4_0_f32.cl");
#endif
backend_ctx->program_mul_mv_q4_0_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_q4_0_f32 = clCreateKernel(backend_ctx->program_mul_mv_q4_0_f32, "kernel_mul_mat_q4_0_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q4_0_f32_v.cl");
#endif
backend_ctx->program_mul_mv_q4_0_f32_v =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_q4_0_f32_v = clCreateKernel(backend_ctx->program_mul_mv_q4_0_f32_v, "kernel_mul_mat_q4_0_f32_v", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q4_0_f32_8x_flat.cl");
#endif
backend_ctx->program_mul_mv_q4_0_f32_8x_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_q4_0_f32_8x_flat = clCreateKernel(backend_ctx->program_mul_mv_q4_0_f32_8x_flat, "kernel_mul_mat_q4_0_f32_8x_flat", &err), err));
GGML_LOG_CONT(".");
// those compiler versions since it is anyway not used for Adreno.
if (backend_ctx->gpu_family != ADRENO ||
backend_ctx->adreno_cl_compiler_version.newer_than_or_same(E031, 38, 11, 0) ||
+ backend_ctx->adreno_cl_compiler_version.type == E17 ||
backend_ctx->adreno_cl_compiler_version.type == DX) {
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src {
const std::string kernel_src = read_file("mul_mv_q4_0_f32_1d_8x_flat.cl");
#endif
backend_ctx->program_mul_mv_q4_0_f32_1d_8x_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_q4_0_f32_1d_8x_flat = clCreateKernel(backend_ctx->program_mul_mv_q4_0_f32_1d_8x_flat, "kernel_mul_mat_q4_0_f32_1d_8x_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q4_0_f32_1d_16x_flat.cl");
#endif
backend_ctx->program_mul_mv_q4_0_f32_1d_16x_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_q4_0_f32_1d_16x_flat = clCreateKernel(backend_ctx->program_mul_mv_q4_0_f32_1d_16x_flat, "kernel_mul_mat_q4_0_f32_1d_16x_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q4_1_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q4_1_f32 = clCreateKernel(prog, "kernel_mul_mv_q4_1_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q4_1_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q4_1_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q4_1_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q4_k_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q4_K_f32 = clCreateKernel(prog, "kernel_mul_mv_q4_K_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q4_k_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q4_K_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q4_K_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q5_0_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q5_0_f32 = clCreateKernel(prog, "kernel_mul_mv_q5_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q5_0_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q5_0_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q5_0_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q5_1_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q5_1_f32 = clCreateKernel(prog, "kernel_mul_mv_q5_1_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q5_1_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q5_1_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q5_1_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q5_k_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q5_K_f32 = clCreateKernel(prog, "kernel_mul_mv_q5_K_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q5_k_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q5_K_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q5_K_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q6_k_f32.cl");
#endif
backend_ctx->program_mul_mv_q6_K =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q6_K_f32 = clCreateKernel(backend_ctx->program_mul_mv_q6_K, "kernel_mul_mv_q6_K_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q6_k_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q6_K_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q6_K_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q8_0_f32.cl");
#endif
backend_ctx->program_mul_mv_q8_0_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q8_0_f32 = clCreateKernel(backend_ctx->program_mul_mv_q8_0_f32, "kernel_mul_mv_q8_0_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q8_0_f32_flat.cl");
#endif
backend_ctx->program_mul_mv_q8_0_f32_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q8_0_f32_flat = clCreateKernel(backend_ctx->program_mul_mv_q8_0_f32_flat, "kernel_mul_mv_q8_0_f32_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_q1_0_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q1_0_f32 = clCreateKernel(prog, "kernel_mul_mv_q1_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_q1_0_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_q1_0_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q1_0_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_iq4_nl_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_iq4_nl_f32 = clCreateKernel(prog, "kernel_mul_mv_iq4_nl_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_iq4_nl_f32_flat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_iq4_nl_f32_flat = clCreateKernel(prog, "kernel_mul_mv_iq4_nl_f32_flat", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mv_mxfp4_f32.cl");
#endif
backend_ctx->program_mul_mv_mxfp4_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_mxfp4_f32 = clCreateKernel(backend_ctx->program_mul_mv_mxfp4_f32, "kernel_mul_mv_mxfp4_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_mxfp4_f32_flat.cl");
#endif
backend_ctx->program_mul_mv_mxfp4_f32_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_mxfp4_f32_flat = clCreateKernel(backend_ctx->program_mul_mv_mxfp4_f32_flat, "kernel_mul_mv_mxfp4_f32_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_f16_f16.cl");
#endif
backend_ctx->program_mul_mv_f16_f16 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f16 = clCreateKernel(backend_ctx->program_mul_mv_f16_f16, "kernel_mul_mat_f16_f16", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_f16_f32_1row.cl");
#endif
backend_ctx->program_mul_mv_f16_f32_1row =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_1row = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_1row, "kernel_mul_mat_f16_f32_1row", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_f16_f32_l4.cl");
#endif
backend_ctx->program_mul_mv_f16_f32_l4 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4 = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4", &err), err));
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4_dr = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_dr", &err), err));
const std::string kernel_src = read_file("mul_mv_f16_f32.cl");
#endif
backend_ctx->program_mul_mv_f16_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32 = clCreateKernel(backend_ctx->program_mul_mv_f16_f32, "kernel_mul_mat_f16_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_f32_f32.cl");
#endif
backend_ctx->program_mul_mv_f32_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_f32_f32 = clCreateKernel(backend_ctx->program_mul_mv_f32_f32, "kernel_mul_mat_f32_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mat_f16_f32.cl");
#endif
backend_ctx->program_mul_mat_f16_f32_tiled =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_tiled = clCreateKernel(backend_ctx->program_mul_mat_f16_f32_tiled, "mul_mat_f16_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("gemm_xmem_f16_f32_os8.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_adreno_xmem_pack_src_f32 =
clCreateKernel(prog, "adreno_xmem_pack_src_f32", &err), err));
const std::string kernel_src = read_file("mul_mm_f32_f32_l4_lm.cl");
#endif
backend_ctx->program_mul_mm_f32_f32_l4_lm =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_f32_f32_l4_lm = clCreateKernel(backend_ctx->program_mul_mm_f32_f32_l4_lm, "kernel_mul_mm_f32_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_f16_f32_l4_lm.cl");
#endif
backend_ctx->program_mul_mm_f16_f32_l4_lm =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_l4_lm = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_l4_lm, "kernel_mul_mm_f16_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_q4_0_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q4_0_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q4_0_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_q4_1_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q4_1_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q4_1_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_q5_0_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q5_0_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q5_0_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_q5_1_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q5_1_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q5_1_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_q8_0_f32_l4_lm.cl");
#endif
backend_ctx->program_mul_mm_q8_0_f32_l4_lm =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q8_0_f32_l4_lm = clCreateKernel(backend_ctx->program_mul_mm_q8_0_f32_l4_lm, "kernel_mul_mm_q8_0_f32_l4_lm", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mm_q1_0_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q1_0_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q1_0_f32_l4_lm", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mm_iq4_nl_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_iq4_nl_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_iq4_nl_f32_l4_lm", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mm_q4_k_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q4_k_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q4_k_f32_l4_lm", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mm_q6_k_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q6_k_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q6_k_f32_l4_lm", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mm_q5_k_f32_l4_lm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_q5_k_f32_l4_lm = clCreateKernel(prog, "kernel_mul_mm_q5_k_f32_l4_lm", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("mul_mm_f16_f32_kq_kqv.cl");
#endif
backend_ctx->program_mul_mm_f16_f32_kqv =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts+" -DKQV ");
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts+" -DKQV ");
backend_ctx->program_mul_mm_f16_f32_kq =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_kqv = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_kqv, "mul_mm_f16_f32_kqv", &err), err));
CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_kq = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_kq, "mul_mm_f16_f32_kq", &err), err));
const std::string kernel_src = read_file("mul.cl");
#endif
backend_ctx->program_mul =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul = clCreateKernel(backend_ctx->program_mul, "kernel_mul", &err), err));
CL_CHECK((backend_ctx->kernel_mul_row = clCreateKernel(backend_ctx->program_mul, "kernel_mul_row", &err), err));
const std::string kernel_src = read_file("norm.cl");
#endif
backend_ctx->program_norm =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_norm = clCreateKernel(backend_ctx->program_norm, "kernel_norm", &err), err));
CL_CHECK((backend_ctx->kernel_norm_mul_add = clCreateKernel(backend_ctx->program_norm, "kernel_norm_mul_add", &err), err));
const std::string kernel_src = read_file("relu.cl");
#endif
backend_ctx->program_relu =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_relu = clCreateKernel(backend_ctx->program_relu, "kernel_relu", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("rms_norm.cl");
#endif
backend_ctx->program_rms_norm =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_rms_norm = clCreateKernel(backend_ctx->program_rms_norm, "kernel_rms_norm", &err), err));
CL_CHECK((backend_ctx->kernel_rms_norm_mul = clCreateKernel(backend_ctx->program_rms_norm, "kernel_rms_norm_mul", &err), err));
const std::string kernel_src = read_file("l2_norm.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_l2_norm_f32 = clCreateKernel(prog, "kernel_l2_norm_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("rope.cl");
#endif
backend_ctx->program_rope =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_rope_norm_f32 = clCreateKernel(backend_ctx->program_rope, "kernel_rope_norm_f32", &err), err));
CL_CHECK((backend_ctx->kernel_rope_norm_f16 = clCreateKernel(backend_ctx->program_rope, "kernel_rope_norm_f16", &err), err));
const std::string kernel_src = read_file("scale.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_scale_f32 = clCreateKernel(prog, "kernel_scale_f32", &err), err));
CL_CHECK((backend_ctx->kernel_scale_f32_4 = clCreateKernel(prog, "kernel_scale_f32_4", &err), err));
const std::string kernel_src = read_file("silu.cl");
#endif
backend_ctx->program_silu =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_silu = clCreateKernel(backend_ctx->program_silu, "kernel_silu", &err), err));
CL_CHECK((backend_ctx->kernel_silu_4 = clCreateKernel(backend_ctx->program_silu, "kernel_silu_4", &err), err));
const std::string kernel_src = read_file("softmax_f32.cl");
#endif
backend_ctx->program_softmax_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_soft_max = clCreateKernel(backend_ctx->program_softmax_f32, "kernel_soft_max", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("softmax_f16.cl");
#endif
backend_ctx->program_softmax_f16 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_soft_max_f16 = clCreateKernel(backend_ctx->program_softmax_f16, "kernel_soft_max_f16", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("softmax_4_f32.cl");
#endif
backend_ctx->program_softmax_4_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_soft_max_4 = clCreateKernel(backend_ctx->program_softmax_4_f32, "kernel_soft_max_4", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("softmax_4_f16.cl");
#endif
backend_ctx->program_softmax_4_f16 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_soft_max_4_f16 = clCreateKernel(backend_ctx->program_softmax_4_f16, "kernel_soft_max_4_f16", &err), err));
GGML_LOG_CONT(".");
" -cl-mad-enable -cl-finite-math-only ";
backend_ctx->program_div =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_div = clCreateKernel(backend_ctx->program_div, "kernel_div", &err), err));
CL_CHECK((backend_ctx->kernel_div_row = clCreateKernel(backend_ctx->program_div, "kernel_div_row", &err), err));
const std::string kernel_src = read_file("sqr.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_sqr_cont_f32 = clCreateKernel(prog, "kernel_sqr_cont_f32", &err), err));
CL_CHECK((backend_ctx->kernel_sqr_cont_f32_4 = clCreateKernel(prog, "kernel_sqr_cont_f32_4", &err), err));
const std::string kernel_src = read_file("sqrt.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_sqrt_cont_f32 = clCreateKernel(prog, "kernel_sqrt_cont_f32", &err), err));
CL_CHECK((backend_ctx->kernel_sqrt_cont_f32_4 = clCreateKernel(prog, "kernel_sqrt_cont_f32_4", &err), err));
const std::string kernel_src = read_file("mean.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mean_f32 = clCreateKernel(prog, "kernel_mean_f32", &err), err));
CL_CHECK((backend_ctx->kernel_mean_f32_4 = clCreateKernel(prog, "kernel_mean_f32_4", &err), err));
const std::string kernel_src = read_file("sub.cl");
#endif
backend_ctx->program_sub =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_sub = clCreateKernel(backend_ctx->program_sub, "kernel_sub", &err), err));
CL_CHECK((backend_ctx->kernel_sub_row = clCreateKernel(backend_ctx->program_sub, "kernel_sub_row", &err), err));
const std::string kernel_src = read_file("sum_rows.cl");
#endif
backend_ctx->program_sum_rows_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_sum_rows_f32 = clCreateKernel(backend_ctx->program_sum_rows_f32, "kernel_sum_rows_f32", &err), err));
CL_CHECK((backend_ctx->kernel_sum_rows_f32_4 = clCreateKernel(backend_ctx->program_sum_rows_f32, "kernel_sum_rows_f32_4", &err), err));
const std::string kernel_src = read_file("cumsum.cl");
#endif
cl_program prog;
- prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_cumsum_blk = clCreateKernel(prog, "kernel_cumsum_blk", &err), err));
CL_CHECK((backend_ctx->kernel_cumsum_add = clCreateKernel(prog, "kernel_cumsum_add", &err), err));
const std::string kernel_src = read_file("sigmoid.cl");
#endif
backend_ctx->program_sigmoid =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_sigmoid_f32 = clCreateKernel(backend_ctx->program_sigmoid, "kernel_sigmoid_f32", &err), err));
CL_CHECK((backend_ctx->kernel_sigmoid_f16 = clCreateKernel(backend_ctx->program_sigmoid, "kernel_sigmoid_f16", &err), err));
const std::string kernel_src = read_file("group_norm.cl");
#endif
backend_ctx->program_group_norm =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_group_norm = clCreateKernel(backend_ctx->program_group_norm, "kernel_group_norm", &err), err));
CL_CHECK((backend_ctx->kernel_group_norm_mul_add = clCreateKernel(backend_ctx->program_group_norm, "kernel_group_norm_mul_add", &err), err));
const std::string kernel_src = read_file("repeat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_repeat_f32 = clCreateKernel(prog, "kernel_repeat_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
#endif
if (!kernel_src.empty()) {
backend_ctx->program_pad =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_pad = clCreateKernel(backend_ctx->program_pad, "kernel_pad", &err), err));
GGML_LOG_CONT(".");
} else {
const std::string kernel_src = read_file("tanh.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_tanh_f32 = clCreateKernel(prog, "kernel_tanh_f32", &err), err));
CL_CHECK((backend_ctx->kernel_tanh_f32_4 = clCreateKernel(prog, "kernel_tanh_f32_4", &err), err));
CL_CHECK((backend_ctx->kernel_tanh_f32_nc = clCreateKernel(prog, "kernel_tanh_f32_nc", &err), err));
const std::string kernel_src = read_file("neg.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_neg_f32 = clCreateKernel(prog, "kernel_neg_f32", &err), err));
CL_CHECK((backend_ctx->kernel_neg_f32_4 = clCreateKernel(prog, "kernel_neg_f32_4", &err), err));
CL_CHECK((backend_ctx->kernel_neg_f32_nc = clCreateKernel(prog, "kernel_neg_f32_nc", &err), err));
const std::string kernel_src = read_file("exp.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_exp_f32 = clCreateKernel(prog, "kernel_exp_f32", &err), err));
CL_CHECK((backend_ctx->kernel_exp_f32_4 = clCreateKernel(prog, "kernel_exp_f32_4", &err), err));
CL_CHECK((backend_ctx->kernel_exp_f32_nc = clCreateKernel(prog, "kernel_exp_f32_nc", &err), err));
const std::string kernel_src = read_file("expm1.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_expm1_f32 = clCreateKernel(prog, "kernel_expm1_f32", &err), err));
CL_CHECK((backend_ctx->kernel_expm1_f32_4 = clCreateKernel(prog, "kernel_expm1_f32_4", &err), err));
CL_CHECK((backend_ctx->kernel_expm1_f32_nc = clCreateKernel(prog, "kernel_expm1_f32_nc", &err), err));
GGML_LOG_CONT(".");
}
+ // abs
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "abs.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("abs.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_abs_f32 = clCreateKernel(prog, "kernel_abs_f32", &err), err));
+ CL_CHECK((backend_ctx->kernel_abs_f32_4 = clCreateKernel(prog, "kernel_abs_f32_4", &err), err));
+ CL_CHECK((backend_ctx->kernel_abs_f32_nc = clCreateKernel(prog, "kernel_abs_f32_nc", &err), err));
+ CL_CHECK((backend_ctx->kernel_abs_f16 = clCreateKernel(prog, "kernel_abs_f16", &err), err));
+ CL_CHECK((backend_ctx->kernel_abs_f16_4 = clCreateKernel(prog, "kernel_abs_f16_4", &err), err));
+ CL_CHECK((backend_ctx->kernel_abs_f16_nc = clCreateKernel(prog, "kernel_abs_f16_nc", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// softplus
{
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src = read_file("softplus.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_softplus_f32 = clCreateKernel(prog, "kernel_softplus_f32", &err), err));
CL_CHECK((backend_ctx->kernel_softplus_f32_4 = clCreateKernel(prog, "kernel_softplus_f32_4", &err), err));
CL_CHECK((backend_ctx->kernel_softplus_f32_nc = clCreateKernel(prog, "kernel_softplus_f32_nc", &err), err));
#endif
if (!kernel_src.empty()) {
backend_ctx->program_upscale =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_upscale = clCreateKernel(backend_ctx->program_upscale, "kernel_upscale", &err), err));
if (backend_ctx->program_upscale) {
cl_int err_bilinear;
const std::string kernel_src = read_file("concat.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_concat_f32 = clCreateKernel(prog, "kernel_concat_f32", &err), err));
CL_CHECK((backend_ctx->kernel_concat_f32_pack = clCreateKernel(prog, "kernel_concat_f32_pack", &err), err));
CL_CHECK(clReleaseProgram(prog));
#endif
if (!kernel_src.empty()) {
backend_ctx->program_tsembd =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_timestep_embedding = clCreateKernel(backend_ctx->program_tsembd, "kernel_timestep_embedding", &err), err));
GGML_LOG_CONT(".");
} else {
const std::string kernel_src = read_file("set_rows.cl");
#endif
backend_ctx->program_set_rows =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_set_rows_f32_i64 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_f32_i64", &err), err));
CL_CHECK((backend_ctx->kernel_set_rows_f32_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_f32_i32", &err), err));
#endif
if (!kernel_src.empty()) {
backend_ctx->program_conv_2d_f16 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), (std::string(compile_opts) + " -DUSE_FP16=1").c_str());
+ build_program_from_source(backend_ctx, kernel_src.c_str(), (std::string(compile_opts) + " -DUSE_FP16=1").c_str());
CL_CHECK((backend_ctx->kernel_conv_2d_f16 = clCreateKernel(backend_ctx->program_conv_2d_f16, "kernel_conv_2d", &err), err));
GGML_LOG_CONT(".");
backend_ctx->program_conv_2d_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_conv_2d_f32 = clCreateKernel(backend_ctx->program_conv_2d_f32, "kernel_conv_2d", &err), err));
GGML_LOG_CONT(".");
} else {
}
if (!kernel_src_f16_f32.empty()) {
backend_ctx->program_conv_2d_f16_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src_f16_f32.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src_f16_f32.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_conv_2d_f16_f32 = clCreateKernel(backend_ctx->program_conv_2d_f16_f32, "kernel_conv_2d", &err), err));
GGML_LOG_CONT(".");
} else {
const std::string kernel_src = read_file("ssm_conv.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_ssm_conv_f32_f32 = clCreateKernel(prog, "kernel_ssm_conv_f32_f32", &err), err));
CL_CHECK((backend_ctx->kernel_ssm_conv_f32_f32_4 = clCreateKernel(prog, "kernel_ssm_conv_f32_f32_4", &err), err));
// Please remember to implement code to handle it.
opts += " -DSUBGROUPS_PER_WG=" + std::to_string(spw);
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), opts);
CL_CHECK((backend_ctx->kernel_gated_delta_net_f32[si][kda][tgpp] =
clCreateKernel(prog, "kernel_gated_delta_net", &err), err));
GGML_LOG_CONT(".");
}
+ // moe_combine (fused router-weight mul + cross-expert sum)
+ {
+ #ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "moe_combine.cl.h"
+ };
+ #else
+ const std::string kernel_src = read_file("moe_combine.cl");
+ #endif
+ cl_program prog = build_program_from_source(
+ backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_moe_combine_f32 =
+ clCreateKernel(prog, "kernel_moe_combine_f32", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// mul_mv_id_q4_0_f32_8x_flat
{
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src = read_file("mul_mv_id_q4_0_f32_8x_flat.cl");
#endif
backend_ctx->program_mul_mv_id_q4_0_f32_8x_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_id_q4_0_f32_8x_flat = clCreateKernel(backend_ctx->program_mul_mv_id_q4_0_f32_8x_flat, "kernel_mul_mv_id_q4_0_f32_8x_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_id_q8_0_f32.cl");
#endif
backend_ctx->program_mul_mv_id_q8_0_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_id_q8_0_f32 = clCreateKernel(backend_ctx->program_mul_mv_id_q8_0_f32, "kernel_mul_mv_id_q8_0_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_id_q8_0_f32_flat.cl");
#endif
backend_ctx->program_mul_mv_id_q8_0_f32_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_id_q8_0_f32_flat = clCreateKernel(backend_ctx->program_mul_mv_id_q8_0_f32_flat, "kernel_mul_mv_id_q8_0_f32_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_id_mxfp4_f32.cl");
#endif
backend_ctx->program_mul_mv_id_mxfp4_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_id_mxfp4_f32 = clCreateKernel(backend_ctx->program_mul_mv_id_mxfp4_f32, "kernel_mul_mv_id_mxfp4_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("mul_mv_id_mxfp4_f32_flat.cl");
#endif
backend_ctx->program_mul_mv_id_mxfp4_f32_flat =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_mul_mv_id_mxfp4_f32_flat = clCreateKernel(backend_ctx->program_mul_mv_id_mxfp4_f32_flat, "kernel_mul_mv_id_mxfp4_f32_flat", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("transpose.cl");
#endif
backend_ctx->program_transpose =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_transpose_32_16 = clCreateKernel(backend_ctx->program_transpose, "kernel_transpose_32_16", &err), err));
CL_CHECK((backend_ctx->kernel_transpose_32 = clCreateKernel(backend_ctx->program_transpose, "kernel_transpose_32", &err), err));
#else
const std::string kernel_src = read_file("gemm_noshuffle_q1_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q1_0_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q1_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
const std::string kernel_src_CL_gemv_general = read_file("gemv_noshuffle_q1_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv_general.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv_general.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q1_0_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q1_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src_CL_gemv_general = read_file("gemv_noshuffle_q4_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv_general.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv_general.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src_CL_gemv = read_file("gemv_noshuffle_q4_0_f32_spec.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32_4096_1_4096 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
CL_gemv_compile_opts += " -DVECTOR_SUB_GROUP_BROADCAST ";
}
- prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
+ prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32_4096_1_11008 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
CL_gemv_compile_opts += " -DVECTOR_SUB_GROUP_BROADCAST ";
}
- prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
+ prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32_11008_1_4096 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
CL_gemv_compile_opts += " -DVECTOR_SUB_GROUP_BROADCAST ";
}
- prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
+ prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32000_1_4096 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
#else
const std::string kernel_src_CL_gemm = read_file("gemm_noshuffle_q4_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src_CL_gemm.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src_CL_gemm.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
#else
const std::string kernel_src = read_file("gemm_noshuffle_q4_1_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_1_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_1_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("gemv_noshuffle_q4_1_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_1_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_1_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
#else
const std::string kernel_src = read_file("gemm_noshuffle_q5_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_0_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
+ // gemm_noshuffle_q5_0_q8_1_dp4a (dp4a dense q5_0 prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_q5_0_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_q5_0_q8_1_dp4a.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_0_q8_1_dp4a", &err), err));
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_0_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_0_q8_1_dp4a_wimg", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// gemv_noshuffle_q5_0_f32
{
std::string CL_gemv_compile_opts = std::string("-cl-std=") + opencl_c_std +
#else
const std::string kernel_src = read_file("gemv_noshuffle_q5_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q5_0_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q5_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
#else
const std::string kernel_src = read_file("gemm_noshuffle_q5_1_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_1_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_1_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
#else
const std::string kernel_src = read_file("gemv_noshuffle_q5_1_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q5_1_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q5_1_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
#else
const std::string kernel_src = read_file("gemm_noshuffle_iq4_nl_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_iq4_nl_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_iq4_nl_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
+ // gemm_noshuffle_iq4_nl_q8_1_dp4a (dp4a dense IQ4_NL prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_iq4_nl_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_iq4_nl_q8_1_dp4a.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_noshuffle_q4_0_q8_1_dp4a (dp4a dense q4_0 prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_q4_0_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_q4_0_q8_1_dp4a.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_0_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// gemv_noshuffle_iq4_nl_f32
{
std::string CL_gemv_compile_opts = std::string("-cl-std=") + opencl_c_std +
const std::string kernel_src = read_file("gemv_noshuffle_iq4_nl_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_iq4_nl_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_iq4_nl_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
#else
const std::string kernel_src = read_file("gemm_noshuffle_q8_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
const std::string kernel_src_CL_gemv_general = read_file("gemv_noshuffle_q8_0_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src_CL_gemv_general.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src_CL_gemv_general.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q8_0_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q8_0_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
#else
const std::string kernel_src = read_file("gemm_noshuffle_q4_k_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
+ // gemm_noshuffle_q4_k_q8_1_dp4a (dp4a dense prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_q4_k_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_q4_k_q8_1_dp4a.cl");
+#endif
+ // Per-device dp4a dense tile. The X2-tuned TILESIZE_N=32 over-occupies LDS on
+ // X1 (1152 B/WG -> few resident WGs); TILESIZE_N=8 (288 B) lifts occupancy on
+ // X1, byte-identical. X2E keeps 32. Env override wins.
+ int q4k_dp4a_ts = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
+ if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) { q4k_dp4a_ts = atoi(e); }
+ std::string dp4a_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(q4k_dp4a_ts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), dp4a_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a", &err), err));
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_noshuffle_q8_0_q8_1_dp4a (dp4a dense q8_0 prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_q8_0_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_q8_0_q8_1_dp4a.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a", &err), err));
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_noshuffle_q5_k_q8_1_dp4a (dp4a dense prefill GEMM for q5_K)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_q5_k_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_q5_k_q8_1_dp4a.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_k_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_noshuffle_q6_k_q8_1_dp4a (dp4a dense prefill GEMM for q6_K ffn_down/output)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_noshuffle_q6_k_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_noshuffle_q6_k_q8_1_dp4a.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q6_k_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // quant_a_q8_1 (plain activation q8_1 pre-pass for the dense dp4a GEMM)
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "quant_a_q8_1.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("quant_a_q8_1.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_quant_a_q8_1 = clCreateKernel(prog, "kernel_quant_a_q8_1", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// gemv_noshuffle_q4_k_f32
{
std::string CL_gemv_compile_opts = std::string("-cl-std=") + opencl_c_std +
const std::string kernel_src = read_file("gemv_noshuffle_q4_k_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_k_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q4_k_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
#else
const std::string kernel_src = read_file("gemv_moe_q4_1_f32_ns.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q4_1_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q4_1_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
#else
const std::string kernel_src = read_file("gemm_moe_q4_1_f32_ns.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q4_1_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q4_1_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemv_moe_mxfp4_f32.cl");
#endif
backend_ctx->program_gemv_moe_mxfp4_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_mxfp4_f32 = clCreateKernel(backend_ctx->program_gemv_moe_mxfp4_f32, "kernel_gemv_moe_mxfp4_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("gemm_moe_mxfp4_f32.cl");
#endif
backend_ctx->program_gemm_moe_mxfp4_f32 =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_f32 = clCreateKernel(backend_ctx->program_gemm_moe_mxfp4_f32, "kernel_gemm_moe_mxfp4_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("gemv_moe_q4_0_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q4_0_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q4_0_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemm_moe_q4_0_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q4_0_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q4_0_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
}
}
+ // gemm_moe_q8_0_f32_ns
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_moe_q8_0_f32_ns.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_moe_q8_0_f32_ns.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
+
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q8_0_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q8_0_f32_ns", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// gemv_moe_q5_0_f32_ns
{
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src = read_file("gemv_moe_q5_0_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q5_0_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q5_0_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemm_moe_q5_0_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q5_0_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q5_0_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemv_moe_q5_1_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q5_1_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q5_1_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemm_moe_q5_1_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q5_1_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q5_1_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemv_moe_q4_k_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q4_k_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q4_k_f32_ns", &err), err));
+ CL_CHECK((backend_ctx->kernel_gemv_moe_q4_k_f32_ns_wimg = clCreateKernel(prog, "kernel_gemv_moe_q4_k_f32_ns_wimg", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
const std::string kernel_src = read_file("gemm_moe_q4_k_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q4_k_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q4_k_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
}
}
+ // gemm_moe_q4_k_q8_1_dp4a (dp4a prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_moe_q4_k_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_moe_q4_k_q8_1_dp4a.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
+
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q4_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_moe_q4_k_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_moe_mxfp4_q8_1_dp4a (dp4a prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_moe_mxfp4_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_moe_mxfp4_q8_1_dp4a.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
+
+ CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_moe_mxfp4_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_moe_q4_0_q8_1_dp4a (dp4a prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_moe_q4_0_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_moe_q4_0_q8_1_dp4a.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
+
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_moe_q4_0_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // gemm_moe_q8_1_dp4a (generic dp4a MoE GEMM; MOE_QT=80 -> q8_0 expert variant)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_moe_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_moe_q8_1_dp4a.cl");
+#endif
+ const std::string opts80 = CL_moe_compile_opts + " -DMOE_QT=80";
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), opts80.c_str());
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q8_1_dp4a_q80 = clCreateKernel(prog, "kernel_gemm_moe_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+
+ const std::string opts50 = CL_moe_compile_opts + " -DMOE_QT=50";
+ cl_program prog50 =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), opts50.c_str());
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q8_1_dp4a_q50 = clCreateKernel(prog50, "kernel_gemm_moe_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog50));
+
+ const std::string opts5 = CL_moe_compile_opts + " -DMOE_QT=5";
+ cl_program prog5 =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), opts5.c_str());
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q8_1_dp4a_q5k = clCreateKernel(prog5, "kernel_gemm_moe_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog5));
+ GGML_LOG_CONT(".");
+ }
+
+ // moe_reorder_quant_a_q8_1 (fused reorder + q8_1 quant)
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "moe_reorder_quant_a_q8_1.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("moe_reorder_quant_a_q8_1.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
+
+ CL_CHECK((backend_ctx->kernel_moe_reorder_quant_a_q8_1 = clCreateKernel(prog, "kernel_moe_reorder_quant_a_q8_1", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// gemv_moe_q5_k_f32_ns
{
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src = read_file("gemv_moe_q5_k_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q5_k_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q5_k_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemm_moe_q5_k_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q5_k_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q5_k_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemv_moe_q6_k_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_q6_k_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_q6_k_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("gemm_moe_q6_k_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_q6_k_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_q6_k_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
+ // gemm_moe_q6_k_f32_ns_bin
+ {
+ size_t bin_size = 0;
+ backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin = nullptr;
+
+ if (use_adreno_bin_kernels(backend_ctx)) {
+ const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_q6_k_f32_ns_ila", &bin_size);
+ if (kernel_bin && bin_size > 0) {
+ cl_program prog =
+ build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size);
+
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin = clCreateKernel(prog, "kernel_gemm_moe_q6_k_f32_ns_ila", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+ }
+ }
+
+ // gemm_moe_q6_k_q8_1_dp4a (dp4a q6_K MoE prefill GEMM)
+ if (backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "gemm_moe_q6_k_q8_1_dp4a.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("gemm_moe_q6_k_q8_1_dp4a.cl");
+#endif
+ cl_program prog =
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
+
+ CL_CHECK((backend_ctx->kernel_gemm_moe_q6_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_moe_q6_k_q8_1_dp4a", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// gemv_moe_mxfp4_f32_ns
{
#ifdef GGML_OPENCL_EMBED_KERNELS
const std::string kernel_src = read_file("gemv_moe_mxfp4_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_moe_mxfp4_f32_ns = clCreateKernel(prog, "kernel_gemv_moe_mxfp4_f32_ns", &err), err));
+ CL_CHECK((backend_ctx->kernel_gemv_moe_mxfp4_f32_ns_wimg = clCreateKernel(prog, "kernel_gemv_moe_mxfp4_f32_ns_wimg", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
}
const std::string kernel_src = read_file("gemm_moe_mxfp4_f32_ns.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_f32_ns = clCreateKernel(prog, "kernel_gemm_moe_mxfp4_f32_ns", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("moe_reorder_b.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_moe_reorder_b = clCreateKernel(prog, "kernel_moe_reorder_b", &err), err));
CL_CHECK(clReleaseProgram(prog));
const std::string kernel_src = read_file("moe_sort_by_expert.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_moe_histogram = clCreateKernel(prog, "kernel_moe_histogram", &err), err));
CL_CHECK((backend_ctx->kernel_moe_scan = clCreateKernel(prog, "kernel_moe_scan", &err), err));
}
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q6_K_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q6_K_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("gemm_noshuffle_q6_k_f32.cl");
#endif
cl_program prog =
- build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_moe_compile_opts);
+ build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q6_K_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q6_K_f32", &err), err));
GGML_LOG_CONT(".");
const std::string kernel_src = read_file("gemv_noshuffle_q5_k_f32.cl");
#endif
- cl_program prog = build_program_from_source(
- backend_ctx->context, backend_ctx->device, kernel_src.c_str(), CL_gemv_compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), CL_gemv_compile_opts);
CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q5_k_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q5_k_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
#else
const std::string kernel_src = read_file("gemm_noshuffle_q5_k_f32.cl");
#endif
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_k_f32 = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_k_f32", &err), err));
CL_CHECK(clReleaseProgram(prog));
GGML_LOG_CONT(".");
variant == FA_VARIANT_Q4_0_SPLIT;
if (is_split) {
opts += " -D N_SPLIT=" + std::to_string(cfg->n_split);
- if (backend_ctx->has_subgroup_shuffle) {
- opts += backend_ctx->has_qcom_subgroup_shuffle
- ? " -D cl_qcom_subgroup_shuffle=1"
- : " -D cl_khr_subgroup_shuffle=1";
- }
+ }
+ // Shuffle define for the split tile paths AND the cluster-parallel decode
+ // kernel (q1_vec_mq_split_c8) in the plain F32_F16 program. Without it the
+ // c8 kernel is compiled out (HAS_SUBGROUP_SHUFFLE guard) and dispatch
+ // falls back to the baseline mq_split.
+ if ((is_split || variant == FA_VARIANT_F32_F16) && backend_ctx->has_subgroup_shuffle) {
+ opts += backend_ctx->has_qcom_subgroup_shuffle
+ ? " -D cl_qcom_subgroup_shuffle=1"
+ : " -D cl_khr_subgroup_shuffle=1";
+ }
+ // X1E drops the explicit sub-group size pin on the c8 kernels, compiler
+ // routes the fp16-heavy kernel to a slow variant with explicit subgroup size
+ if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
+ opts += " -D FA_C8_NO_SG_PIN";
}
return opts;
}
opts += " -D FA_DECODE_ONLY -D FA_DECODE_MINIMAL";
}
+ // c8 cluster width (GGML_OPENCL_FA_CL_C overrides): value = GQA4 cluster
+ // width (kernel default 8); the g8 programs use 2x the value (default 16).
+ // Wider clusters halve per-lane o_acc at the cost of position streams per
+ // subgroup
+ static const int fa_cl_c_env = []{
+ const char * e = std::getenv("GGML_OPENCL_FA_CL_C");
+ const int x = (e && e[0]) ? atoi(e) : 0;
+ return (x == 8 || x == 16 || x == 32) ? x : 0; // 0 = per-gen default
+ }();
+ const int fa_cl_c_gqa4 = fa_cl_c_env ? fa_cl_c_env
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 16 : 0);
+ const std::string opts_cl_c_gqa4 = fa_cl_c_gqa4
+ ? " -D FA_CL_C=" + std::to_string(fa_cl_c_gqa4) : std::string();
+ const std::string fa_cl_c_g8_val = std::to_string(fa_cl_c_gqa4 ? fa_cl_c_gqa4 * 2 : 16);
+
const char * tag = nullptr;
switch (variant) {
case FA_VARIANT_F16: tag = "fa f16"; break;
default: break;
}
cl_program prog = build_program_from_source_ex(
- backend_ctx->context, backend_ctx->device, src.c_str(), opts,
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts + opts_cl_c_gqa4,
/*fatal=*/false, tag, backend_ctx->queue);
if (!prog) { return false; }
clReleaseKernel(k_q1_vec_mq_split_k_img);
}
}
+ // Cluster-parallel decode variant
+ cl_kernel k_q1_vec_mq_split_c8 = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_c8, 256,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_c8[{dk, dv}] = k_q1_vec_mq_split_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_c8);
+ }
+ }
cl_kernel k_merge = clCreateKernel(prog, "flash_attn_f32_merge", &err);
if (err == CL_SUCCESS) {
backend_ctx->fa.f32_merge[{dk, dv}] = k_merge;
// second compile of the same source with -DMQ_GQA=8.
// FA_MQ_ONLY keeps only the vec_mq kernels so that the program
// compiles within the Adreno compiler's memory budget at DK>=256.
- const std::string opts_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3 -D FA_MQ_ONLY";
+ // FA_CL_C for the g8 program: MQ_GQA=8 doubles the c8 kernel's
+ // per-lane o_acc, so widen the cluster to keep the register
+ // footprint inside the 192-thread WG cap (see fa_cl_c_gqa4 above
+ // for the per-gen default).
+ const std::string opts_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3 -D FA_MQ_ONLY -D FA_CL_C=" + fa_cl_c_g8_val;
cl_program prog_g8 = fa_decode_only ? nullptr : build_program_from_source_ex(
backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8,
/*fatal=*/false, "fa f32_f16 MQ_GQA=8", backend_ctx->queue);
clReleaseKernel(k_q1_vec_mq_split_g8_k_img);
}
}
+ // Cluster-parallel decode, MQ_GQA=8 / FA_CL_C=16 specialization
+ cl_kernel k_q1_vec_mq_split_g8_c8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_g8_c8, mq_g8_required_wg,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (g8)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8[{dk, dv}] = k_q1_vec_mq_split_g8_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_g8_c8", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_g8_c8);
+ }
+ }
// hybrid local-tile + MQ_GQA=8
if (dk == 128 && dv == 128) {
cl_kernel k_lmq_g8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_local_mq_split", &err);
}
clReleaseProgram(prog_g8);
}
+ // NSG_SPLIT=2 programs for the cluster-parallel kernel: its register
+ // footprint caps the per-kernel WG at 128 on X2 (< the stock 256/192
+ // requirement), so it can never register from the stock programs.
+ // With FA_CL_NCL position streams per subgroup, 2 subgroups still
+ // carry 16 in-flight rows per WG (baseline WG has 4). FA_MQ_ONLY
+ // keeps these compiles minimal; skipped when the stock program c8
+ // registered (some other device) or shuffles are absent.
+ if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count({dk, dv}) == 0) {
+ const std::string opts_c8_ns2 = opts + " -D FA_MQ_ONLY -D MQ_GQA=4 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2" + opts_cl_c_gqa4;
+ cl_program prog_c8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8_ns2,
+ /*fatal=*/false, "fa f32_f16 c8 NSG2", backend_ctx->queue);
+ if (prog_c8) {
+ cl_kernel k_c8 = clCreateKernel(prog_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ // WG = MQ_NSG(2) × Q1_WG_SIZE(=FA_SG): 128 Adreno (64), 64 Intel (32).
+ const size_t c8_ns2_wg = backend_ctx->gpu_family == INTEL ? 64 : 128;
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8, c8_ns2_wg,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (ns2)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2[{dk, dv}] = k_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8_ns2", dk, dv);
+ } else {
+ clReleaseKernel(k_c8);
+ }
+ }
+ clReleaseProgram(prog_c8);
+ }
+ }
+ // FA_CL_C=32 g8 program for the DK=DV=256 GQA=8
+ if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
+ dk == 256 && dv == 256) {
+ const std::string opts_g8_c32 = opts + " -D FA_MQ_ONLY -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2 -D FA_CL_C=32";
+ cl_program prog_g8_c32 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8_c32,
+ /*fatal=*/false, "fa f32_f16 c32 g8 d256 NSG2", backend_ctx->queue);
+ if (prog_g8_c32) {
+ cl_kernel k_g8_c32 = clCreateKernel(prog_g8_c32, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8_c32, 128,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (g8 c32 d256)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32[{dk, dv}] = k_g8_c32;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8_c32, "flash_attn_f32_f16_q1_vec_mq_split_g8_c32", dk, dv);
+ } else {
+ clReleaseKernel(k_g8_c32);
+ }
+ }
+ clReleaseProgram(prog_g8_c32);
+ }
+ }
+ if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count({dk, dv}) == 0) {
+ const std::string opts_g8_c8_ns2 = opts + " -D FA_MQ_ONLY -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2 -D FA_CL_C=" + fa_cl_c_g8_val;
+ cl_program prog_g8_c8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8_c8_ns2,
+ /*fatal=*/false, "fa f32_f16 c8 g8 NSG2", backend_ctx->queue);
+ if (prog_g8_c8) {
+ cl_kernel k_g8_c8 = clCreateKernel(prog_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8_c8, 128,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (g8 ns2)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2[{dk, dv}] = k_g8_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_g8_c8_ns2", dk, dv);
+ } else {
+ clReleaseKernel(k_g8_c8);
+ }
+ }
+ clReleaseProgram(prog_g8_c8);
+ }
+ }
break;
}
case FA_VARIANT_Q8_0:
}
clReleaseProgram(prog_mq_g8);
}
+ // GQA=4 cluster-parallel program (NSG_SPLIT=2 / WG=128)
+ if (backend_ctx->has_subgroup_shuffle) {
+ auto & m_c8_gqa4 = is_q8 ? backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8
+ : backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8;
+ const std::string name_c8_gqa4 = name_q1 + "_vec_mq_split_c8";
+ const std::string opts_c8_gqa4 = opts + " -D MQ_GQA=4 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2" + opts_cl_c_gqa4;
+ cl_program prog_c8_gqa4 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8_gqa4,
+ /*fatal=*/false, is_q8 ? "fa q8_0 c8 GQA4 NSG2" : "fa q4_0 c8 GQA4 NSG2",
+ backend_ctx->queue);
+ if (prog_c8_gqa4) {
+ cl_kernel k_c8_gqa4 = clCreateKernel(prog_c8_gqa4, name_c8_gqa4.c_str(), &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8_gqa4, 128,
+ name_c8_gqa4.c_str(), dk, dv)) {
+ m_c8_gqa4[{dk, dv}] = k_c8_gqa4;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8_gqa4, name_c8_gqa4.c_str(), dk, dv);
+ } else {
+ clReleaseKernel(k_c8_gqa4);
+ }
+ }
+ clReleaseProgram(prog_c8_gqa4);
+ }
+ }
+ // Cluster-parallel q4_0 decode kernel
+ if (!is_q8 && backend_ctx->has_subgroup_shuffle) {
+ const std::string opts_c8 = opts + " -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2";
+ cl_program prog_c8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8,
+ /*fatal=*/false, "fa q4_0 c8 NSG2", backend_ctx->queue);
+ if (prog_c8) {
+ cl_kernel k_c8 = clCreateKernel(prog_c8, "flash_attn_f32_q4_0_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8, 128,
+ "flash_attn_f32_q4_0_q1_vec_mq_split_c8 (g8 ns2)", dk, dv)) {
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8[{dk, dv}] = k_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8, "flash_attn_f32_q4_0_q1_vec_mq_split_g8_c8", dk, dv);
+ } else {
+ clReleaseKernel(k_c8);
+ }
+ }
+ clReleaseProgram(prog_c8);
+ }
+ }
break;
}
case FA_VARIANT_F32_F16_SPLIT: {
backend_ctx->has_subgroup_shuffle ? "true" : "false");
GGML_LOG_INFO("ggml_opencl: device FP16 support: %s\n",
backend_ctx->fp16_support ? "true" : "false");
+ GGML_LOG_INFO("ggml_opencl: khr dot product support: %s\n",
+ backend_ctx->has_integer_dot ? "true" : "false");
GGML_LOG_INFO("ggml_opencl: mem base addr align: %u\n",
backend_ctx->alignment);
GGML_LOG_INFO("ggml_opencl: global mem size: %zu MB\n",
// check Adreno large buffer support
backend_ctx->adreno_has_large_buffer = strstr(ext_buffer, "cl_qcom_large_buffer") != NULL;
+
// subgroup shuffle support (N_SPLIT>1 FA kernel)
backend_ctx->has_qcom_subgroup_shuffle = strstr(ext_buffer, "cl_qcom_subgroup_shuffle") != NULL;
backend_ctx->has_subgroup_shuffle =
strstr(ext_buffer, "cl_khr_subgroup_shuffle") != NULL ||
backend_ctx->has_qcom_subgroup_shuffle;
+ // check for cl_khr_integer_dot_product
+ // cl_qcom_dot_product8 uses signed * unsigned
+ // while cl_khr_integer_dot_product uses signed * signed -- we stick with khr for now
+ backend_ctx->has_integer_dot =
+ strstr(ext_buffer, "cl_khr_integer_dot_product") != NULL;
+
cl_uint base_align_in_bits;
CL_CHECK(clGetDeviceInfo(device, CL_DEVICE_MEM_BASE_ADDR_ALIGN, sizeof(cl_uint), &base_align_in_bits, NULL));
GGML_ASSERT(base_align_in_bits % 8u == 0);
static const char * ragged_gran_env = getenv("GGML_OPENCL_MOE_RAGGED_GRAN");
backend_ctx->adreno_moe_ragged_skip_gran = (ragged_gran_env != NULL) ? atoi(ragged_gran_env) : 8;
+ // whether fuse moe combine
+ static const char * fuse_moe_combine_env = getenv("GGML_OPENCL_FUSE_MOE_COMBINE");
+ backend_ctx->fuse_moe_combine = fuse_moe_combine_env == NULL ? 1 : (atoi(fuse_moe_combine_env) != 0);
+
+ // ragged moe dp4 variant
+ static const char * ragged_dp4_env = getenv("GGML_OPENCL_MOE_RAGGED");
+ backend_ctx->adreno_use_moe_ragged_dp4 = ragged_dp4_env == NULL ? 1 : (atoi(ragged_dp4_env) != 0);
+
#ifdef GGML_OPENCL_USE_ADRENO_BIN_KERNELS
// try loading adreno binary kernels if enabled
// if fails to load, builtin kernels will be used
cl_kernel kernel,
cl_mem src, cl_mem dst, size_t size,
cl_int stride, cl_int rows,
- bool blocking = true
+ bool blocking = true,
+ bool auto_local = false // let driver pick local size for non-uniform workgroups
) {
static ggml_cl_buffer buf;
size_t local_size[3] = {64, 1, 1};
size_t global_size[3] = {(size_t)stride, (size_t)rows, 1};;
CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL,
- global_size, local_size, 0, NULL, NULL));
+ global_size, auto_local ? NULL : local_size, 0, NULL, NULL));
if (blocking) {
CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, trans, dst, 0, 0, size, 0, NULL, &evt));
ggml_backend_opencl_context * backend_ctx,
cl_mem src, cl_mem dst, size_t size,
cl_int stride, cl_int rows,
- bool blocking = true
+ bool blocking = true,
+ bool auto_local = false
) {
transpose_2d(backend_ctx, backend_ctx->kernel_transpose_8_buf,
- src, dst, size, stride, rows, blocking);
+ src, dst, size, stride, rows, blocking, auto_local);
}
static void transpose_2d_as_16b(
cl_mem d = nullptr;
// Scales in image1d_buffer_t.
cl_mem d_img = nullptr;
+ // Uniform per-32-block scale (2/block) + min (1/block, = d*16 for the -16 centering)
+ // for the generic dp4a MoE GEMM. Built from d.
+ cl_mem scale = nullptr;
+ cl_mem min = nullptr;
// Size of quantized values.
size_t size_qs = 0;
// Size of 5-th bit values.
CL_CHECK(clReleaseMemObject(qs_img));
qs_img = nullptr;
}
+ if (scale != nullptr) {
+ CL_CHECK(clReleaseMemObject(scale));
+ scale = nullptr;
+ }
+ if (min != nullptr) {
+ CL_CHECK(clReleaseMemObject(min));
+ min = nullptr;
+ }
qh_img = nullptr;
d_img = nullptr;
cl_mem d = nullptr;
cl_mem d_img = nullptr;
+ // Uniform per-16-segment scale (16/superblock) for the generic dp4a MoE GEMM.
+ // Expanded from d at set_tensor; the int8 codes are reused from q.
+ // q8_0 is symmetric so no min buffer (has_min=0).
+ cl_mem scale = nullptr;
+
size_t size_q = 0;
size_t size_d = 0;
CL_CHECK(clReleaseMemObject(d));
d = nullptr;
}
+ if (scale != nullptr) {
+ CL_CHECK(clReleaseMemObject(scale));
+ scale = nullptr;
+ }
// Currently, q_img and d_img are not used. They can be image1d_buffer_t
// that wraps around q and d to utilize image access path.
q_img = nullptr;
cl_mem d = nullptr;
// Min for each super block.
cl_mem dm = nullptr;
+ // Uniform per-32-block scale (2/block) + min (1/block, = dm*mn) decoded from the
+ // 6-bit packed s[] for the generic dp4a MoE GEMM kernel_gemm_moe_q8_1_dp4a.
+ // Built from s/d/dm at set_tensor; q/qh are reused as-is.
+ cl_mem scale = nullptr;
+ cl_mem min = nullptr;
size_t size_q = 0;
size_t size_qh = 0;
CL_CHECK(clReleaseMemObject(q_img));
q_img = nullptr;
}
+ if (scale != nullptr) {
+ CL_CHECK(clReleaseMemObject(scale));
+ scale = nullptr;
+ }
+ if (min != nullptr) {
+ CL_CHECK(clReleaseMemObject(min));
+ min = nullptr;
+ }
size_q = 0;
size_qh = 0;
sync_with_other_backends(backend_ctx);
}
+// True if two tensors share a device buffer with overlapping byte ranges. The pool
+// allocator may place a fused op's output over a sequentially-dead input (safe for the
+// original separate kernels, but a read/write race inside one fused kernel).
+static bool ggml_cl_tensors_overlap(const ggml_tensor * x, const ggml_tensor * y) {
+ ggml_tensor_extra_cl * ex = (ggml_tensor_extra_cl *)x->extra;
+ ggml_tensor_extra_cl * ey = (ggml_tensor_extra_cl *)y->extra;
+ if (!ex || !ey || ex->data_device != ey->data_device) { return false; }
+ const cl_ulong xo = ex->offset + x->view_offs, xe = xo + ggml_nbytes(x);
+ const cl_ulong yo = ey->offset + y->view_offs, ye = yo + ggml_nbytes(y);
+ return xo < ye && yo < xe;
+}
+
+// Detect the MoE combine epilogue: router-weight MUL ([n_embd,k,nt] * [1,k,nt]) followed
+// by k VIEWs of it and a (k-1)-long ADD reduction chain producing [n_embd, nt]. When it
+// matches (and the output does not alias the inputs), the whole subgraph collapses to one
+// weighted-sum-across-experts kernel.
+static bool ggml_opencl_can_fuse_moe_combine(const struct ggml_cgraph * cgraph, int node_idx,
+ const ggml_tensor ** out_final_add) {
+ const ggml_tensor * mul = cgraph->nodes[node_idx];
+ if (mul->op != GGML_OP_MUL) { return false; }
+ const ggml_tensor * experts = mul->src[0];
+ const ggml_tensor * weights = mul->src[1];
+ if (!experts || !weights) { return false; }
+ if (experts->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { return false; }
+
+ const int64_t n_embd = experts->ne[0];
+ const int64_t k = experts->ne[1];
+ const int64_t nt = experts->ne[2];
+ if (k < 2 || k > 64 || experts->ne[3] != 1 || n_embd % 4 != 0) { return false; }
+ if (weights->ne[0] != 1 || weights->ne[1] != k || weights->ne[2] != nt || weights->ne[3] != 1) { return false; }
+ if (mul->ne[0] != n_embd || mul->ne[1] != k || mul->ne[2] != nt) { return false; }
+ // the fused kernel needs contiguous experts/weights and a contiguous 2D dst
+ if (!ggml_is_contiguous(experts) || !ggml_is_contiguous(weights)) { return false; }
+
+ const int n_nodes = 1 + (int)k + (int)(k - 1); // MUL + k*VIEW + (k-1)*ADD
+ if (n_nodes >= 32) { return false; }
+ if (node_idx + n_nodes > cgraph->n_nodes) { return false; }
+
+ enum ggml_op ops[1 + 64 + 63];
+ int n = 0;
+ ops[n++] = GGML_OP_MUL;
+ for (int j = 0; j < (int)k; ++j) { ops[n++] = GGML_OP_VIEW; }
+ for (int j = 0; j < (int)k - 1; ++j) { ops[n++] = GGML_OP_ADD; }
+ const int outs[] = { node_idx + n_nodes - 1 };
+ if (!ggml_can_fuse_subgraph(cgraph, node_idx, n_nodes, ops, outs, 1)) { return false; }
+
+ for (int j = 0; j < (int)k; ++j) {
+ const ggml_tensor * vw = cgraph->nodes[node_idx + 1 + j];
+ if (vw->op != GGML_OP_VIEW || vw->src[0] != mul || vw->ne[0] != n_embd || vw->ne[1] != nt) { return false; }
+ }
+ const ggml_tensor * final_add = cgraph->nodes[node_idx + n_nodes - 1];
+ if (final_add->op != GGML_OP_ADD || final_add->type != GGML_TYPE_F32 ||
+ final_add->ne[0] != n_embd || final_add->ne[1] != nt || final_add->ne[2] != 1) { return false; }
+ if (!ggml_is_contiguous(final_add)) { return false; }
+ // the fused kernel reads experts + writes final_add in one pass; bail if the
+ // pool allocator overlapped the output with the (large) experts input -- would race.
+ // The small weights input is copied to a private scratch in the dispatch, so its own
+ // aliasing with the output is handled there and does not block the fusion.
+ if (ggml_cl_tensors_overlap(experts, final_add)) { return false; }
+
+ *out_final_add = final_add;
+ return true;
+}
+
+static void ggml_cl_moe_combine_fused(ggml_backend_t backend, const ggml_tensor * mul, const ggml_tensor * dst) {
+ ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *)backend->context;
+ const ggml_tensor * experts = mul->src[0];
+ const ggml_tensor * weights = mul->src[1];
+
+ ggml_tensor_extra_cl * ee = (ggml_tensor_extra_cl *)experts->extra;
+ ggml_tensor_extra_cl * ew = (ggml_tensor_extra_cl *)weights->extra;
+ ggml_tensor_extra_cl * ed = (ggml_tensor_extra_cl *)dst->extra;
+ cl_ulong off_e = ee->offset + experts->view_offs;
+ cl_ulong off_w = ew->offset + weights->view_offs;
+ cl_ulong off_d = ed->offset + dst->view_offs;
+
+ const int n_embd4 = (int)(experts->ne[0] / 4);
+ const int k = (int)experts->ne[1];
+ const int nt = (int)experts->ne[2];
+ const cl_uint e1 = (cl_uint)(experts->nb[1] / sizeof(float));
+ const cl_uint e2 = (cl_uint)(experts->nb[2] / sizeof(float));
+ const cl_uint w1 = (cl_uint)(weights->nb[1] / sizeof(float));
+ const cl_uint w2 = (cl_uint)(weights->nb[2] / sizeof(float));
+ const cl_uint d1 = (cl_uint)(dst->nb[1] / sizeof(float));
+
+ // The router weights are tiny ([1,k,nt]) and may share a pool buffer with the output;
+ // copy them into a private scratch so the fused kernel never reads aliased memory.
+ const size_t w_bytes = ggml_nbytes(weights);
+ backend_ctx->prealloc_moe_combine_w.allocate(backend_ctx->context, w_bytes);
+ CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, ew->data_device, backend_ctx->prealloc_moe_combine_w.buffer,
+ off_w, 0, w_bytes, 0, NULL, NULL));
+ cl_mem w_dev = backend_ctx->prealloc_moe_combine_w.buffer;
+ cl_ulong w_off = 0;
+
+ cl_kernel kernel = backend_ctx->kernel_moe_combine_f32;
+ int a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &ee->data_device));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &off_e));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &w_dev));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &w_off));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &ed->data_device));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &off_d));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_embd4));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &k));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nt));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_uint), &e1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_uint), &e2));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_uint), &w1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_uint), &w2));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_uint), &d1));
+
+ size_t lws[2] = { 64, 1 };
+ size_t gws[2] = { (size_t)(((n_embd4 + 63) / 64) * 64), (size_t)nt };
+ backend_ctx->enqueue_ndrange_kernel(kernel, 2, gws, lws, dst);
+}
+
static bool ggml_opencl_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops) {
if (!ggml_can_fuse(cgraph, node_idx, ops)) {
return false;
i += 2;
continue;
}
+ // Fuse the MoE combine: router-weight mul + cross-expert add chain ->
+ // one weighted-sum-across-experts kernel.
+ if (backend_ctx->fuse_moe_combine && !backend_ctx->disable_fusion) {
+ const ggml_tensor * combine_out = nullptr;
+ if (ggml_opencl_can_fuse_moe_combine(cgraph, i, &combine_out)) {
+ ggml_cl_moe_combine_fused(backend, node, combine_out);
+ i += 2 * (int)node->ne[1] - 1; // skip the k VIEWs + (k-1) ADDs
+ continue;
+ }
+ }
+
if (!backend_ctx->disable_fusion && ggml_opencl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) {
ggml_opencl_op_rms_norm_fused(backend, node, cgraph->nodes[i+1]);
i++;
return threashold_ok;
}
+static bool adreno_e17_compiler_quirks(const ggml_backend_opencl_context *backend_ctx) {
+ if (!backend_ctx || backend_ctx->gpu_family != GPU_FAMILY::ADRENO ||
+ backend_ctx->adreno_cl_compiler_version.type != ADRENO_CL_COMPILER_TYPE::E17) {
+ return false;
+ }
+ const char * env = getenv("GGML_OPENCL_ADRENO_E17_QUIRKS");
+ return !(env && env[0] == '0');
+}
+
inline bool use_adreno_moe_kernels(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) {
- GGML_UNUSED(backend_ctx);
+ // The moe weight repack kernels *_trans4_ns alias a private ushort8 through a uchar*.
+ // Certain compilers (found with some A7x and A6x) miscompiles this, corrupting the weights.
+ // So, exclude A6x and A7x from using Adreno MoE kernels for now.
+ // The quants that have a general mul_mat_id kernel fallback to the general version; the
+ // rest fallback to CPU.
+ if (backend_ctx && (backend_ctx->adreno_gen == ADRENO_GPU_GEN::A6X ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::ADRENO_UNKNOWN)) {
+ return false;
+ }
+
+ if (adreno_e17_compiler_quirks(backend_ctx)) {
+ return false;
+ }
+
int ne01 = tensor->ne[1];
return (((strstr(tensor->name, "ffn") != NULL) && (strstr(tensor->name, "exps") != NULL)) || (strstr(tensor->name, "as") != NULL)) && (ne01 % 32 == 0);
}
size_t elem_num = tensor->ne[0] * tensor->ne[1] * tensor->ne[2] * tensor->ne[3];
- return ((elem_num < 128 * 1024 * 1024) && adreno_kernel); // max element num: 2**27
+ // The 2D weight transpose (transpose_2d_as_*) tiles rows by 4 over a 2D matrix,
+ // so it requires K(ne0)%32==0, M(ne1)%4==0 and ne2==ne3==1.
+ const bool shape_ok = (tensor->ne[0] % 32 == 0) && (tensor->ne[1] % 4 == 0) &&
+ (tensor->ne[2] == 1) && (tensor->ne[3] == 1);
+
+ return ((elem_num < 128 * 1024 * 1024) && adreno_kernel && shape_ok); // max element num: 2**27
}
static inline bool use_flat_gemv_for_large_m_q4_K(const ggml_tensor *tensor) {
return op->src[0]->type == GGML_TYPE_F32;
case GGML_UNARY_OP_EXPM1:
return op->src[0]->type == GGML_TYPE_F32;
+ case GGML_UNARY_OP_ABS:
+ return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
case GGML_UNARY_OP_SOFTPLUS:
return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
default:
case GGML_OP_MEAN:
return op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_FLASH_ATTN_EXT: {
+ // The E17 compilers segfault while building FA kernels, skip E17 for now
+ if (adreno_e17_compiler_quirks(backend_ctx)) {
+ return false;
+ }
const ggml_tensor * q = op->src[0];
const ggml_tensor * k = op->src[1];
const ggml_tensor * v = op->src[2];
return false;
}
+ // Some compilers for A7x (Adreno 740, compiler E031.41) crashes when
+ // building FA kernels with mixed or quant types (f32_f16, f32_q8_0, f32_q4_0)
+ // Here we skip all A7x for these kernels to avoid crash
+ if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X &&
+ (is_f32_f16 || is_f32_q8_0 || is_f32_q4_0)) {
+ return false;
+ }
+
if (dk == 512) {
if (backend_ctx->gpu_family == INTEL) {
return false;
extra->qs_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_format_qs, &img_desc_qs, NULL, &err);
tensor->extra = extra;
+ // Generic dp4a MoE path
+ {
+ static const char * q5dp4a_env = getenv("GGML_OPENCL_Q5_MOE_DP4A");
+ const bool q5dp4a = q5dp4a_env ? (atoi(q5dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ if (q5dp4a && ne02 > 1 && (ne00 % 32 == 0)) {
+ size_t nb32 = (size_t)ne00 / 32;
+ size_t sc_elems = (size_t)ne02 * ne01 * nb32 * 2;
+ size_t mn_elems = (size_t)ne02 * ne01 * nb32;
+ extra->scale = clCreateBuffer(context, CL_MEM_READ_WRITE, sc_elems * sizeof(cl_half), NULL, &err); CL_CHECK(err);
+ extra->min = clCreateBuffer(context, CL_MEM_READ_WRITE, mn_elems * sizeof(cl_half), NULL, &err); CL_CHECK(err);
+ cl_kernel ek = backend_ctx->kernel_moe_expand_scale_q5_0;
+ CL_CHECK(clSetKernelArg(ek, 0, sizeof(cl_mem), &extra->d));
+ CL_CHECK(clSetKernelArg(ek, 1, sizeof(cl_mem), &extra->scale));
+ CL_CHECK(clSetKernelArg(ek, 2, sizeof(cl_mem), &extra->min));
+ CL_CHECK(clSetKernelArg(ek, 3, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(ek, 4, sizeof(int), &ne01));
+ size_t eg[3] = { (size_t)(((ne01 + 63) / 64) * 64), nb32, (size_t)ne02 };
+ size_t el[3] = { 64, 1, 1 };
+ cl_event evt;
+ CL_CHECK(clEnqueueNDRangeKernel(queue, ek, 3, NULL, eg, el, 0, NULL, &evt));
+ CL_CHECK(clWaitForEvents(1, &evt));
+ }
+ }
+
return;
}
#endif // GGML_OPENCL_USE_ADRENO_KERNELS
tensor->extra = extra;
ctx->q8_0_soa_tensors.insert(tensor);
+ // Generic dp4a MoE path (opt-in GGML_OPENCL_Q8_MOE_DP4A)
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ {
+ static const char * q8dp4a_env = getenv("GGML_OPENCL_Q8_MOE_DP4A");
+ const bool q8dp4a = q8dp4a_env ? (atoi(q8dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ if (q8dp4a && tensor->ne[2] > 1 && (tensor->ne[0] % 32 == 0)) {
+ int ne00 = (int)tensor->ne[0];
+ int ne01 = (int)tensor->ne[1];
+ int ne02 = (int)tensor->ne[2];
+ size_t nb32 = (size_t)ne00 / 32;
+ size_t scale_elems = (size_t)ne02 * ne01 * nb32 * 2; // 2 per-16-seg scales / 32-block
+ extra->scale = clCreateBuffer(context, CL_MEM_READ_WRITE, scale_elems * sizeof(cl_half), NULL, &err);
+ CL_CHECK(err);
+ cl_kernel ek = backend_ctx->kernel_moe_expand_scale_q8_0;
+ CL_CHECK(clSetKernelArg(ek, 0, sizeof(cl_mem), &extra->d));
+ CL_CHECK(clSetKernelArg(ek, 1, sizeof(cl_mem), &extra->scale));
+ CL_CHECK(clSetKernelArg(ek, 2, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(ek, 3, sizeof(int), &ne01));
+ size_t eg[3] = { (size_t)(((ne01 + 63) / 64) * 64), nb32, (size_t)ne02 };
+ size_t el[3] = { 64, 1, 1 };
+ cl_event evt;
+ CL_CHECK(clEnqueueNDRangeKernel(queue, ek, 3, NULL, eg, el, 0, NULL, &evt));
+ CL_CHECK(clWaitForEvents(1, &evt));
+ }
+ }
+#endif
+
// Transpose the weights and scales
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
if (enable_adreno_trans_weight(backend_ctx, tensor)) {
transpose_2d_as_16b(backend_ctx, extra->q, extra->q, size_q, K/4, M);
transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/256, M);
transpose_2d_as_16b(backend_ctx, extra->dm, extra->dm, size_dm, K/256, M);
+
+ // Transpose s as uchar
+ transpose_2d_as_8b(backend_ctx, extra->s, extra->s, size_s, K/256*12, M, true, true);
}
#endif // GGML_OPENCL_USE_ADRENO_KERNELS
return;
CL_CHECK(err);
tensor->extra = extra;
+ // Generic dp4a MoE path
+ {
+ static const char * q5kdp4a_env = getenv("GGML_OPENCL_Q5K_MOE_DP4A");
+ const bool q5kdp4a = q5kdp4a_env ? (atoi(q5kdp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ if (q5kdp4a && ne02 > 1 && (ne00 % 256 == 0)) {
+ size_t nb32 = (size_t)ne00 / 32;
+ size_t sc_elems = (size_t)ne02 * ne01 * nb32 * 2;
+ size_t mn_elems = (size_t)ne02 * ne01 * nb32;
+ extra->scale = clCreateBuffer(context, CL_MEM_READ_WRITE, sc_elems * sizeof(cl_half), NULL, &err); CL_CHECK(err);
+ extra->min = clCreateBuffer(context, CL_MEM_READ_WRITE, mn_elems * sizeof(cl_half), NULL, &err); CL_CHECK(err);
+ cl_kernel ek = backend_ctx->kernel_moe_expand_scale_q5_K;
+ CL_CHECK(clSetKernelArg(ek, 0, sizeof(cl_mem), &extra->s));
+ CL_CHECK(clSetKernelArg(ek, 1, sizeof(cl_mem), &extra->d));
+ CL_CHECK(clSetKernelArg(ek, 2, sizeof(cl_mem), &extra->dm));
+ CL_CHECK(clSetKernelArg(ek, 3, sizeof(cl_mem), &extra->scale));
+ CL_CHECK(clSetKernelArg(ek, 4, sizeof(cl_mem), &extra->min));
+ CL_CHECK(clSetKernelArg(ek, 5, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(ek, 6, sizeof(int), &ne01));
+ size_t eg[3] = { (size_t)(((ne01 + 63) / 64) * 64), (size_t)(ne00 / 256), (size_t)ne02 };
+ size_t el[3] = { 64, 1, 1 };
+ cl_event evt;
+ CL_CHECK(clEnqueueNDRangeKernel(queue, ek, 3, NULL, eg, el, 0, NULL, &evt));
+ CL_CHECK(clWaitForEvents(1, &evt));
+ }
+ }
+
return;
}
#endif // GGML_OPENCL_USE_ADRENO_KERNELS
size_t size_q = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*ggml_blck_size(tensor->type)/2;
size_t size_d = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*sizeof(ggml_fp16_t);
size_t size_dm = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*sizeof(ggml_fp16_t);
+ size_t size_s = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*12;
static ggml_cl_buffer buf_trans_q;
static ggml_cl_buffer buf_trans_d;
static ggml_cl_buffer buf_trans_dm;
+ static ggml_cl_buffer buf_trans_s;
buf_trans_q.allocate(backend_ctx->context, size_q);
buf_trans_d.allocate(backend_ctx->context, size_d);
buf_trans_dm.allocate(backend_ctx->context, size_dm);
+ buf_trans_s.allocate(backend_ctx->context, size_s);
- // Transpose q, d, dm back
+ // Transpose q, d, dm, s back
transpose_2d_as_16b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K/4);
transpose_2d_as_16b(backend_ctx, extra->d, buf_trans_d.buffer, size_d, M, K/256);
transpose_2d_as_16b(backend_ctx, extra->dm, buf_trans_dm.buffer, size_dm, M, K/256);
+ transpose_2d_as_8b (backend_ctx, extra->s, buf_trans_s.buffer, size_s, M, K/256*12, true, true);
cl_kernel kernel = backend_ctx->kernel_restore_block_q4_K_noshuffle;
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &buf_trans_q.buffer));
- CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->s));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &buf_trans_s.buffer));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &buf_trans_d.buffer));
CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &buf_trans_dm.buffer));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &data_device));
static ggml_backend_buffer_t ggml_backend_opencl_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buffer_type, size_t size) {
ggml_backend_opencl_context *backend_ctx = ggml_cl_init(buffer_type->device);
+
+ if (!backend_ctx->program_cache_initialized) {
+ backend_ctx->program_cache = cl_program_cache_init(backend_ctx->device);
+ backend_ctx->program_cache_initialized = true;
+ }
load_cl_kernels(backend_ctx);
// clCreateBuffer returns -61 for size 0
cl_int err;
cl_mem mem = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, size, NULL, &err);
+#if GGML_OPENCL_TARGET_VERSION >= 300
+ // clCreateBufferWithProperties and cl_mem_properties are OpenCL 3.0. Drivers older than
+ // that do not export the symbol, so a build targeting them fails to link. The large
+ // buffer extension is only ever enabled on drivers that are well past 3.0, so this path
+ // is dead there anyway.
if (err != CL_SUCCESS && backend_ctx->adreno_use_large_buffer) {
cl_mem_properties props[] = { 0x41A6 /* CL_LARGE_BUFFER_QCOM */, 1, 0 };
mem = clCreateBufferWithProperties(backend_ctx->context, props, CL_MEM_READ_WRITE, size, NULL, &err);
}
+#endif
if (err != CL_SUCCESS) {
GGML_LOG_INFO("%s: failed to allocate %.2f MiB\n", __func__, size / 1024.0 / 1024.0);
ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
cl_ulong offset0 = extra0->offset + src0->view_offs;
- cl_ulong offset1 = extra1->offset + src0->view_offs;
+ cl_ulong offset1 = extra1->offset + src1->view_offs;
cl_ulong offsetd = extrad->offset + dst->view_offs;
ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
}
}
+static void ggml_cl_abs(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+ GGML_ASSERT(src0);
+ GGML_ASSERT(src0->extra);
+ GGML_ASSERT(dst);
+ GGML_ASSERT(dst->extra);
+
+ UNUSED(src1);
+
+ ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+ ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra;
+ ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
+
+ cl_ulong offset0 = extra0->offset + src0->view_offs;
+ cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+ const int ne00 = src0->ne[0];
+ const int ne01 = src0->ne[1];
+ const int ne02 = src0->ne[2];
+ const int ne03 = src0->ne[3];
+
+ const cl_ulong nb00 = src0->nb[0];
+ const cl_ulong nb01 = src0->nb[1];
+ const cl_ulong nb02 = src0->nb[2];
+ const cl_ulong nb03 = src0->nb[3];
+
+ const cl_ulong nb0 = dst->nb[0];
+ const cl_ulong nb1 = dst->nb[1];
+ const cl_ulong nb2 = dst->nb[2];
+ const cl_ulong nb3 = dst->nb[3];
+
+ cl_kernel kernel;
+
+ if (ggml_is_contiguous(src0)) {
+ // Handle contiguous input
+ int n = ggml_nelements(dst);
+ if (n % 4 == 0) {
+ if (src0->type == GGML_TYPE_F32) {
+ kernel = backend_ctx->kernel_abs_f32_4;
+ } else {
+ kernel = backend_ctx->kernel_abs_f16_4;
+ }
+ n /= 4;
+ } else {
+ if (src0->type == GGML_TYPE_F32) {
+ kernel = backend_ctx->kernel_abs_f32;
+ } else {
+ kernel = backend_ctx->kernel_abs_f16;
+ }
+ }
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra0->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
+
+ size_t global_work_size[] = {(size_t)n, 1, 1};
+ size_t local_work_size[] = {64, 1, 1};
+
+ size_t * local_work_size_ptr = local_work_size;
+ if (n % 64 != 0 && !backend_ctx->non_uniform_workgroups) {
+ local_work_size_ptr = nullptr;
+ }
+
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size_ptr, dst);
+ } else {
+ // Handle non-contiguous input
+ if (src0->type == GGML_TYPE_F32) {
+ kernel = backend_ctx->kernel_abs_f32_nc;
+ } else {
+ kernel = backend_ctx->kernel_abs_f16_nc;
+ }
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra0->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &nb00));
+ CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &nb0));
+ CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &nb1));
+ CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &nb2));
+ CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &nb3));
+
+ int nth = 64;
+
+ size_t global_work_size[] = {(size_t)ne01*nth, (size_t)ne02, (size_t)ne03};
+ size_t local_work_size[] = {(size_t)nth, 1, 1};
+
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ }
+}
+
static void ggml_cl_softplus(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_ASSERT(src0);
GGML_ASSERT(src0->extra);
// Flash-Decoding K-split decision. Resolved here, before the prefill
// prepass, because KV-pad and blk prepass are pure overhead when FD fires.
- const int is_causal = (mask == NULL && n_q > 1 && n_q == n_kv);
+ // Do not infer causality from tensor shapes: a NULL mask means full
+ // (bidirectional) attention, e.g. ViT encoders, where n_q == n_kv as well.
+ // Causal attention in llama.cpp always comes with an explicit KQ mask.
+ // Inferring is_causal here corrupted mmproj output on OpenCL (see #23800).
+ const int is_causal = 0;
const int fd_max_n_q = (d_head_q <= FD_MAX_DK_MULTI) ? FD_MAX_N_Q_MULTI : 1;
cl_kernel fd_k_split = NULL;
bool use_fd_mq = false;
const bool nq_in_vec_range = (n_q >= 1) && (n_q <= N_MAX_VEC_NQ);
const bool nq1_only = (n_q == 1);
+
+ // Cluster-parallel decode default on for Adreno X2E/X1E
+ static const int c8_env_state = []{
+ const char * e = getenv("GGML_OPENCL_FA_C8");
+ if (e == NULL || e[0] == '\0') { return -1; }
+ return (e[0] != '0') ? 1 : 0;
+ }();
+ const bool c8_default_on = backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E;
+ const bool c8_f16_on = (c8_env_state >= 0) ? (c8_env_state == 1) : c8_default_on;
+ // Quant-KV (q4_0/q8_0) GQA4 c8: default-on X2E + X1E
+ const bool c8_quant_on = (c8_env_state >= 0) ? (c8_env_state == 1) : c8_default_on;
if (mq_enabled && mq_kv_ok && nq_in_vec_range && !is_causal &&
backend_ctx->gpu_family != INTEL &&
!use_local_tile &&
getenv("GGML_OPENCL_FA_K_IMG") != NULL &&
getenv("GGML_OPENCL_FA_K_IMG")[0] != '0' &&
backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img.count(dk_dv) > 0;
- if (k_img_on) {
+ // Cluster-parallel decode
+ const bool c8_env = d_head_q == 128 && d_head_v == 128 && c8_f16_on;
+ if (c8_env && backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.at(dk_dv);
+ use_fd_mq = true;
+ } else if (c8_env && backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ } else if (k_img_on) {
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img.at(dk_dv);
use_fd_mq = true;
use_fa_k_img = true;
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split.at(dk_dv);
use_fd_mq = true;
}
+ // Cluster-parallel decode, DK=DV=256 GQA=8
+ } else if (nq1_only && is_mixed && gqa_ratio_dispatch == 8 &&
+ d_head_q == 256 && d_head_v == 256 &&
+ c8_env_state == 1 &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ // Cluster-parallel decode for the g8
+ } else if (is_mixed && gqa_ratio_dispatch == 8 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ c8_f16_on &&
+ (backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count(dk_dv) > 0 ||
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2.count(dk_dv) > 0)) {
+ if (backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.at(dk_dv);
+ fd_mq_wg = 192;
+ } else {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2.at(dk_dv);
+ fd_mq_wg = 128;
+ }
+ use_fd_mq = true;
} else if (is_mixed && gqa_ratio_dispatch == 8 &&
d_head_q == 128 && d_head_v == 128 &&
getenv("GGML_OPENCL_FA_K_IMG") != NULL &&
} else if (nq1_only && is_q8_0 && gqa_ratio_dispatch == 4 &&
d_head_q == 128 && d_head_v == 128 &&
backend_ctx->fa.f32_q8_0_q1_vec_mq_split.count(dk_dv) > 0) {
- fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split.at(dk_dv);
+ // Cluster-parallel q8_0 GQA4
+ if (c8_quant_on &&
+ backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8.at(dk_dv);
+ fd_mq_wg = 128;
+ } else {
+ fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split.at(dk_dv);
+ }
use_fd_mq = true;
} else if (nq1_only && is_q4_0) {
const char * q4_mq_env = getenv("GGML_OPENCL_FA_Q4_MQ");
const bool q4_mq_on = (q4_mq_env != NULL) && (q4_mq_env[0] != '0');
- if (q4_mq_on && gqa_ratio_dispatch == 8 &&
+ // Cluster-parallel q4_0
+ const bool q4_c8_on = c8_env_state == 1 &&
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8.count(dk_dv) > 0;
+ if (q4_c8_on && gqa_ratio_dispatch == 8 &&
+ d_head_q == 64 && d_head_v == 64) {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ } else if (q4_mq_on && gqa_ratio_dispatch == 8 &&
d_head_q == 128 && d_head_v == 128 &&
backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8.count(dk_dv) > 0) {
fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8.at(dk_dv);
} else if (gqa_ratio_dispatch == 4 &&
d_head_q == 128 && d_head_v == 128 &&
backend_ctx->fa.f32_q4_0_q1_vec_mq_split.count(dk_dv) > 0) {
- fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split.at(dk_dv);
+ // Cluster-parallel q4_0 GQA4
+ if (c8_quant_on &&
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8.at(dk_dv);
+ fd_mq_wg = 128;
+ } else {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split.at(dk_dv);
+ }
use_fd_mq = true;
}
}
}
}
+ // Intel cluster-parallel decode FA
+ if (fd_k_split == NULL && backend_ctx->gpu_family == INTEL && n_q == 1 && !is_causal &&
+ is_mixed && gqa_ratio_dispatch == 4 && d_head_q == 128 && d_head_v == 128 &&
+ n_kv >= FD_MIN_N_KV &&
+ getenv("GGML_OPENCL_FA_C8") != NULL && getenv("GGML_OPENCL_FA_C8")[0] != '0' &&
+ backend_ctx->fa.f32_merge.count(dk_dv) > 0) {
+ if (backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ } else if (backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 64;
+ }
+ }
if (fd_k_split == NULL &&
n_q >= 1 && n_q <= fd_max_n_q && n_kv >= FD_MIN_N_KV && !is_causal &&
d_head_q <= FD_MAX_DK &&
// <--------------------------------------------> //
extra0 = src0->view_src ? (ggml_tensor_extra_cl *)src0->view_src->extra : (ggml_tensor_extra_cl *)src0->extra;
- region.origin = (extra0->offset);
+ region.origin = (extra0->offset + src0->view_offs);
if (nb01 > nb02) {
// KQ
region.size = nb01 * ne01;
// create sub-buffer for B
// <--------------------------------------------> //
- region.origin = (extra1->offset);
+ region.origin = (extra1->offset + src1->view_offs);
region.size = nb10 * ne10 * ne11 * ne12;
B_sub_buffer = clCreateSubBuffer((extra1->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
// create sub-buffer for output C
// <--------------------------------------------> //
- region.origin = (extrad->offset);
+ region.origin = (extrad->offset + dst->view_offs);
region.size = ne0 * ne1 * dst->ne[2] * dst->nb[0]; // size of C in bytes
D_sub_buffer = clCreateSubBuffer((extrad->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
CL_CHECK(clReleaseMemObject(b_sub_buf));
CL_CHECK(clReleaseMemObject(b_img));
} else {
+ // dp4a (int8) dense prefill GEMM, default off
+ static const char * q4_0_dense_dp4a_env = getenv("GGML_OPENCL_Q4_0_DENSE_DP4A");
+ bool q4_0_dense_dp4a_on = q4_0_dense_dp4a_env
+ ? (atoi(q4_0_dense_dp4a_env) != 0)
+ : false;
+ // dot prod has to be available
+ q4_0_dense_dp4a_on = backend_ctx->has_integer_dot && q4_0_dense_dp4a_on;
+
+ if (q4_0_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a
+ && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
+ cl_mem a_sub = nullptr;
+ region.origin = offset1;
+ region.size = (size_t)K * N * sizeof(float);
+ CL_CHECK((a_sub = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err));
+
+ const size_t n_blocks = (size_t)N * (K / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)N * K * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &a_sub));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a;
+ int ai = 0;
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q4_0->q));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q4_0->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &M));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &N));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &K));
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, 32) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ CL_CHECK(clReleaseMemObject(a_sub));
+ return;
+ }
+
cl_mem b_sub_buf = nullptr;
cl_mem b_sub_buf_trans = nullptr;
cl_mem b_img = nullptr;
CL_CHECK(clReleaseMemObject(b_sub_buf));
CL_CHECK(clReleaseMemObject(b_img));
} else {
+ // dp4a (int8) dense q5_0 prefill GEMM, default off
+ static const char * q5_dense_dp4a_env = getenv("GGML_OPENCL_Q5_DENSE_DP4A");
+ static const char * q5_dense_wimg_env = getenv("GGML_OPENCL_Q5_DENSE_DP4A_WIMG");
+ const bool q5_dense_wimg_on = q5_dense_wimg_env && (atoi(q5_dense_wimg_env) != 0);
+ bool q5_dense_dp4a_on = q5_dense_wimg_on
+ ? true
+ : (q5_dense_dp4a_env && (atoi(q5_dense_dp4a_env) != 0));
+ // dot prod has to be available
+ q5_dense_dp4a_on = backend_ctx->has_integer_dot && q5_dense_dp4a_on;
+
+ if (q5_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q5_0_q8_1_dp4a
+ && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
+ cl_mem a_sub = nullptr;
+ region.origin = offset1;
+ region.size = (size_t)K * N * sizeof(float);
+ CL_CHECK((a_sub = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err));
+
+ const size_t n_blocks = (size_t)N * (K / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)N * K * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &a_sub));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ // optional qs texture (image1d_buffer over the nibble plane; the same
+ // CL_R/UINT32 view, width M*K/8, the GEMV path builds).
+ cl_mem q5_qs_img = nullptr;
+ bool use_wimg = q5_dense_wimg_on;
+ if (use_wimg) {
+ const size_t tex = (size_t)M * (size_t)K / 8; // uint32 texels (2 ushorts/texel)
+ if (tex == 0 || tex > backend_ctx->image_max_buffer_size) {
+ use_wimg = false;
+ } else {
+ img_fmt = { CL_R, CL_UNSIGNED_INT32 };
+ memset(&img_desc, 0, sizeof(img_desc));
+ img_desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+ img_desc.image_width = tex;
+ img_desc.buffer = extra0_q5_0->qs;
+ q5_qs_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err);
+ if (err != CL_SUCCESS || q5_qs_img == nullptr) { use_wimg = false; q5_qs_img = nullptr; }
+ }
+ }
+
+ cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q5_0_q8_1_dp4a_wimg
+ : backend_ctx->kernel_gemm_noshuffle_q5_0_q8_1_dp4a;
+ int ai = 0;
+ if (use_wimg) {
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q5_qs_img));
+ } else {
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_0->qs));
+ }
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_0->qh));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_0->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &M));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &N));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &K));
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, 32) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ if (q5_qs_img != nullptr) {
+ CL_CHECK(clReleaseMemObject(q5_qs_img));
+ }
+ CL_CHECK(clReleaseMemObject(a_sub));
+ return;
+ }
+
cl_mem b_sub_buf = nullptr;
cl_mem b_sub_buf_trans = nullptr;
cl_mem b_img = nullptr;
CL_CHECK(clReleaseMemObject(b_sub_buf));
CL_CHECK(clReleaseMemObject(b_img));
} else {
+ // dp4a (int8) dense IQ4_NL prefill GEMM. Quantizes the [N,K] activations to
+ // q8_1 and runs the int8 dot instead of the f16 half-dot. Large-batch
+ // (ne1>8) only
+ static const char * iq4nl_dense_dp4a_env = getenv("GGML_OPENCL_IQ4NL_DENSE_DP4A");
+ bool iq4nl_dense_dp4a_on = iq4nl_dense_dp4a_env
+ ? (atoi(iq4nl_dense_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ // dot prod has to be available
+ iq4nl_dense_dp4a_on = backend_ctx->has_integer_dot && iq4nl_dense_dp4a_on;
+
+ if (iq4nl_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a
+ && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
+ cl_mem a_sub = nullptr;
+ region.origin = offset1;
+ region.size = (size_t)K * N * sizeof(float);
+ CL_CHECK((a_sub = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err));
+
+ const size_t n_blocks = (size_t)N * (K / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)N * K * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &a_sub));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a;
+ int ai = 0;
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_iq4_nl->q));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_iq4_nl->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &M));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &N));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &K));
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, 32) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ CL_CHECK(clReleaseMemObject(a_sub));
+ return;
+ }
+
cl_mem b_sub_buf = nullptr;
cl_mem b_sub_buf_trans = nullptr;
cl_mem b_img = nullptr;
cl_ulong offset1 = extra1->offset + src1->view_offs;
cl_ulong offsetd = extrad->offset + dst->view_offs;
- GGML_ASSERT(src1->view_offs == 0);
- GGML_ASSERT(dst->view_offs == 0);
-
const int ne00 = src0->ne[0];
const int ne01 = src0->ne[1];
const int ne02 = src0->ne[2];
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &q_img));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra0_q8_0->d));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &b_img));
- CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &extra1->offset));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset1));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &extrad->data_device));
- CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &extrad->offset));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offsetd));
CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne00));
CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &ne01));
CL_CHECK(clSetKernelArg(kernel, 8, sizeof(int), &ne02));
CL_CHECK(clReleaseMemObject(b_img));
CL_CHECK(clReleaseMemObject(b_sub_buf));
} else {
+ // dp4a dense q8_0 prefill GEMM. Quantizes the [N,K] activations to
+ // q8_1 and runs the int8 dot instead of the f16 half-dot. Large-batch
+ // (ne1>8) only; q8_0 weights are already int8 (no requant) and symmetric
+ // (no min term)
+ static const char * q8_dense_dp4a_env = getenv("GGML_OPENCL_Q8_DENSE_DP4A");
+ static const char * q8_dense_wimg_env = getenv("GGML_OPENCL_Q8_DENSE_DP4A_WIMG");
+ const bool q8_dense_wimg_on = q8_dense_wimg_env && (atoi(q8_dense_wimg_env) != 0);
+
+ const bool q8_bin_loaded = (backend_ctx->kernel_gemm_noshuffle_q8_0_f32_bin != nullptr);
+ // bin kernel takes precedence
+ bool q8_dense_dp4a_on = q8_dense_wimg_on
+ ? true
+ : q8_dense_dp4a_env
+ ? (atoi(q8_dense_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E && !q8_bin_loaded);
+ // dot prod has to be available
+ q8_dense_dp4a_on = backend_ctx->has_integer_dot && q8_dense_dp4a_on;
+
+ if (q8_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a
+ && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
+ cl_mem a_sub = nullptr;
+ region.origin = offset1;
+ region.size = (size_t)K * N * sizeof(float);
+ CL_CHECK((a_sub = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err));
+
+ const size_t n_blocks = (size_t)N * (K / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)N * K * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &a_sub));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ // optional weight texture, the same CL_R/UINT32 view, width M*K/4
+ cl_mem q8_q_img = nullptr;
+ bool use_wimg = q8_dense_wimg_on;
+ if (use_wimg) {
+ const size_t tex = (size_t)M * (size_t)K / 4; // uint32 texels
+ if (tex == 0 || tex > backend_ctx->image_max_buffer_size) {
+ use_wimg = false;
+ } else {
+ img_fmt = { CL_R, CL_UNSIGNED_INT32 };
+ memset(&img_desc, 0, sizeof(img_desc));
+ img_desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+ img_desc.image_width = tex;
+ img_desc.buffer = extra0_q8_0->q;
+ q8_q_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err);
+ if (err != CL_SUCCESS || q8_q_img == nullptr) { use_wimg = false; q8_q_img = nullptr; }
+ }
+ }
+
+ cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg
+ : backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a;
+ int ai = 0;
+ if (use_wimg) {
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q8_q_img));
+ } else {
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q8_0->q));
+ }
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q8_0->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &M));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &N));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &K));
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, 32) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ if (q8_q_img != nullptr) {
+ CL_CHECK(clReleaseMemObject(q8_q_img));
+ }
+ CL_CHECK(clReleaseMemObject(a_sub));
+ return;
+ }
+
// use bin kernel if available
if (backend_ctx->kernel_gemm_noshuffle_q8_0_f32_bin) {
int K_pad = K;
size_t global_work_size_t[2] = { (size_t)width_B, (size_t)padded_height_B };
backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size_t, local_work_size_t, dst);
+ // dp4a (int8) dense prefill GEMM and weight via texture
+ static const char * q4k_dense_dp4a_env = getenv("GGML_OPENCL_Q4K_DENSE_DP4A");
+ static const char * q4k_dense_wimg_env = getenv("GGML_OPENCL_Q4K_DENSE_DP4A_WIMG");
+
+ const bool q4k_dense_wimg_on = q4k_dense_wimg_env && (atoi(q4k_dense_wimg_env) != 0);
+ bool q4k_dense_dp4a_on = q4k_dense_wimg_on
+ ? true
+ : q4k_dense_dp4a_env
+ ? (atoi(q4k_dense_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+
+ // dp4 has to be available
+ q4k_dense_dp4a_on = backend_ctx->has_integer_dot && q4k_dense_dp4a_on;
+
+ // Min N for the dp4a prefill GEMM, default 9, i.e., ne1 > 8
+ static const char * q4k_dp4a_minn_env = getenv("GGML_OPENCL_Q4K_DP4A_MINN");
+ const int q4k_dp4a_minn = q4k_dp4a_minn_env ? atoi(q4k_dp4a_minn_env) : 9;
+
+ if (q4k_dense_dp4a_on && N >= q4k_dp4a_minn && (K % 32 == 0) && (M % 64 == 0)) {
+ const size_t n_blocks = (size_t)N * (K / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)N * K * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &b_sub_buf));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ // check if weights go through texture
+ cl_mem q4k_q_img = nullptr;
+ bool use_wimg = q4k_dense_wimg_on;
+ if (use_wimg) {
+ const size_t tex = (size_t)M * (size_t)K / 8; // uint32 texels = bytes/4
+ if (tex == 0 || tex > backend_ctx->image_max_buffer_size) {
+ use_wimg = false;
+ } else {
+ img_fmt = { CL_R, CL_UNSIGNED_INT32 };
+ memset(&img_desc, 0, sizeof(img_desc));
+ img_desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+ img_desc.image_width = tex;
+ img_desc.buffer = extra0_q4_k->q;
+ q4k_q_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err);
+ if (err != CL_SUCCESS || q4k_q_img == nullptr) {
+ use_wimg = false;
+ q4k_q_img = nullptr;
+ }
+ }
+ }
+
+ cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
+ : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a;
+ int ai = 0;
+ if (use_wimg) {
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q4k_q_img));
+ } else {
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q4_k->q));
+ }
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q4_k->s));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q4_k->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q4_k->dm));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &M));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &N));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &K));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d6));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d4));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_hi2));
+ // Must match the compile-time TILESIZE_N chosen at program build (per-device,
+ // X1E=8 else 32; env override). Same inputs -> same value.
+ int q4k_dp4a_ts = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
+ if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) q4k_dp4a_ts = atoi(e);
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, q4k_dp4a_ts) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ if (q4k_q_img != nullptr) {
+ CL_CHECK(clReleaseMemObject(q4k_q_img));
+ }
+ CL_CHECK(clReleaseMemObject(b_sub_buf));
+ CL_CHECK(clReleaseMemObject(b_sub_buf_trans));
+ CL_CHECK(clReleaseMemObject(b_img));
+ CL_CHECK(clReleaseMemObject(b_img_trans));
+ return;
+ }
+
// gemm
kernel = backend_ctx->kernel_gemm_noshuffle_q4_k_f32;
int padded_N = N + padding;
region.size = ne00 * ne1 * sizeof(float);
CL_CHECK((b_sub_buf = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err));
+ // dp4a (int8) dense q6_K prefill GEMM
+ static const char * q6k_dense_dp4a_env = getenv("GGML_OPENCL_Q6K_DENSE_DP4A");
+ bool q6k_dense_dp4a_on = (q6k_dense_dp4a_env != nullptr)
+ ? (atoi(q6k_dense_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen != ADRENO_GPU_GEN::X1E);
+ // dot prod has to be available
+ q6k_dense_dp4a_on = backend_ctx->has_integer_dot && q6k_dense_dp4a_on;
+
+ const bool is_output_w_dp4a = strncmp(src0->name, "output", 6) == 0 ||
+ strncmp(src0->name, "token_embd", 10) == 0;
+
+ if (q6k_dense_dp4a_on && !is_output_w_dp4a && ne1 > 8 && (ne00 % 32 == 0) && (ne01 % 64 == 0)) {
+ const int M = ne01, N = ne1, K = ne00;
+ const size_t n_blocks = (size_t)N * (K / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)N * K * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &b_sub_buf));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a;
+ int ai = 0;
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q6_K->ql));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q6_K->qh));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q6_K->s));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q6_K->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &M));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &N));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &K));
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, 32) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ CL_CHECK(clReleaseMemObject(b_sub_buf));
+ return;
+ }
+
// image for activation
img_fmt.image_channel_order = CL_RGBA;
img_fmt.image_channel_data_type = CL_FLOAT;
size_t global_work_size_t[2] = {(size_t)width_B, (size_t)padded_height_B};
backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size_t, local_work_size_t, dst);
+ // dp4a (int8) dense q5_K prefill GEMM
+ static const char * q5k_dense_dp4a_env = getenv("GGML_OPENCL_Q5K_DENSE_DP4A");
+ bool q5k_dense_dp4a_on = q5k_dense_dp4a_env
+ ? (atoi(q5k_dense_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ // dot prod has to be available
+ q5k_dense_dp4a_on = backend_ctx->has_integer_dot && q5k_dense_dp4a_on;
+
+ if (q5k_dense_dp4a_on && ne1 > 8 && (ne00 % 32 == 0) && (ne01 % 64 == 0)) {
+ const int Mm = ne01, Nn = ne1, Kk = ne00;
+ const size_t n_blocks = (size_t)Nn * (Kk / 32);
+ backend_ctx->prealloc_moe_qa.allocate(context, (size_t)Nn * Kk * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half));
+
+ cl_int tb = (cl_int)n_blocks;
+ cl_kernel qk = backend_ctx->kernel_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &b_sub_buf));
+ CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb));
+ size_t q_local[1] = { 64 };
+ size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
+ backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
+
+ cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a;
+ int ai = 0;
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_k->q));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_k->qh));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_k->s));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_k->d));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extra0_q5_k->dm));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &Mm));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &Nn));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_int), &Kk));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d6));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d4));
+ CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_hi2));
+ size_t d_local[3] = { 64, 1, 1 };
+ size_t d_global[3] = { 64, (size_t)(Mm / 64), (size_t)CEIL_DIV(Nn, 32) };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
+
+ CL_CHECK(clReleaseMemObject(b_sub_buf));
+ CL_CHECK(clReleaseMemObject(b_sub_buf_trans));
+ CL_CHECK(clReleaseMemObject(b_img));
+ CL_CHECK(clReleaseMemObject(b_img_trans));
+ return;
+ }
+
// gemm
kernel = backend_ctx->kernel_gemm_noshuffle_q5_k_f32;
int padded_N = N + padding;
GGML_ASSERT(ne00 == ne10);
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ // adreno GEMM/GEMV kernels do not support broadcast, assuming ne2 and ne3 are 1 for src1
+ // so we handle broadcast here
+ if ((ne12 > 1 || ne13 > 1) && ne02 == 1 && ne03 == 1 &&
+ src0t != GGML_TYPE_F16 && src0t != GGML_TYPE_F32) {
+ for (int i13 = 0; i13 < ne13; ++i13) {
+ for (int i12 = 0; i12 < ne12; ++i12) {
+ ggml_tensor s1 = *src1;
+ s1.ne[2] = 1; s1.ne[3] = 1;
+ s1.view_offs = src1->view_offs + (size_t)i12*nb12 + (size_t)i13*nb13;
+ ggml_tensor d = *dst;
+ d.ne[2] = 1; d.ne[3] = 1;
+ d.view_offs = dst->view_offs + (size_t)i12*nb2 + (size_t)i13*nb3;
+ ggml_cl_mul_mat(backend, src0, &s1, &d);
+ }
+ }
+ return;
+ }
+#endif
+
int nth0 = 32;
int nth1 = 1;
int nrows = 1;
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
if(src0t == GGML_TYPE_F16 && src1t == GGML_TYPE_F32){
if (ne01 >= 64 && ne1 >= 32 && ne00 >= 16 && (ne12 % ne02) == 0 &&
+ // the KQ/KQV image kernels do not handle dim 3 (multi-stream batches)
+ ne03 == 1 && ne13 == 1 &&
// dst is wrapped with image1d_buffer, the size limit applies, also src0
(ne0 * ne1 * dst->ne[2] * dst->nb[0] / 4 <= backend_ctx->image_max_buffer_size)) {
// For KQ
backend_ctx->enqueue_ndrange_kernel(kernel, 3, histogram_global_size, histogram_local_size, src);
+ // [MOE_TILES] env-gated padding probe: read back total_tiles (= Sum_e
+ // ceil(k_e/n_tile_size)) and compare to the ideal tile count for the real
+ // routing count. Quantifies the per-expert tile-padding waste. Blocking
+ // readback perturbs timing -> diagnostic only.
+ if (getenv("GGML_OPENCL_MOE_TILES_DEBUG")) {
+ int h_total = 0;
+ clFinish(backend_ctx->queue);
+ CL_CHECK(clEnqueueReadBuffer(backend_ctx->queue, total_tiles_buf, CL_TRUE, 0, sizeof(int), &h_total, 0, NULL, NULL));
+ const int routings = ne20 * ne21;
+ const int ideal = (routings + n_tile_size - 1) / n_tile_size;
+ const int slots = h_total * n_tile_size;
+ fprintf(stderr, "[MOE_TILES] routings=%d (ne20=%d ne21=%d nexp=%d) total_tiles=%d ideal=%d slots=%d pad=%.1f%%\n",
+ routings, ne20, ne21, ne02, h_total, ideal, slots,
+ routings > 0 ? 100.0 * (slots - routings) / routings : 0.0);
+ fflush(stderr);
+ }
+
CL_CHECK(clReleaseMemObject(original_router_buf));
CL_CHECK(clReleaseMemObject(hist_buf));
CL_CHECK(clReleaseMemObject(tile_offset_buf));
backend_ctx->toggle_reorder = false;
}
- cl_mem sub_buf_src1_pre, buf_src1_reordered, image_src1_reordered, sub_buf_dst, buf_dst_image;
+ cl_mem sub_buf_src1_pre, sub_buf_dst, buf_dst_image;
+ cl_mem buf_src1_reordered = nullptr, image_src1_reordered = nullptr;
cl_mem buf_src2, buf_src2_emap;
+ // dp4a (int8) prefill GEMM variant
+ static const char * q4_0_moe_dp4a_env = getenv("GGML_OPENCL_Q4_0_MOE_DP4A");
+ bool use_moe_dp4a = q4_0_moe_dp4a_env
+ ? (atoi(q4_0_moe_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ // dot prod has to be available
+ use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
+ // bin kernel takes precedence
+ use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr;
+
cl_buffer_region region;
region.origin = 0;
region.size = sizeof(int) * max_post_router_tile * n_tile_size;
sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
- // Create image for reordered src1
- // Use pre-allocated placeholder
- region.origin = 0;
- region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
- backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
- buf_src1_reordered = clCreateSubBuffer(
- backend_ctx->prealloc_act_trans.buffer,
- 0,
- CL_BUFFER_CREATE_TYPE_REGION,
- ®ion,
- &status);
- CL_CHECK(status);
- cl_image_format image_format_buf_src1;
- cl_image_desc image_desc_buf_src1;
- image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
- image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
- if (backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin) {
- // bin kernel uses slightly different image format
- image_format_buf_src1 = {CL_R, CL_FLOAT};
- image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
- }
- image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
- CL_CHECK(status);
-
unsigned short map_ratio = ne20 / ne11;
GGML_ASSERT(((map_ratio == 1) || (map_ratio == ne20)) && "Map ratio not supported\n");
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
- size_t reorder_b_local_size[3] = {256, 1, 1};
- size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
-
- // Dispatch reorder kernel
- backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ if (!use_moe_dp4a) {
+ // Create image for reordered src1
+ // Use pre-allocated placeholder
+ region.origin = 0;
+ region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
+ backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
+ buf_src1_reordered = clCreateSubBuffer(
+ backend_ctx->prealloc_act_trans.buffer,
+ 0,
+ CL_BUFFER_CREATE_TYPE_REGION,
+ ®ion,
+ &status);
+ CL_CHECK(status);
+ cl_image_format image_format_buf_src1;
+ cl_image_desc image_desc_buf_src1;
+ image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
+ image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
+ if (backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin) {
+ // bin kernel uses slightly different image format
+ image_format_buf_src1 = {CL_R, CL_FLOAT};
+ image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
+ }
+ image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
+ CL_CHECK(status);
+
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
+
+ size_t reorder_b_local_size[3] = {256, 1, 1};
+ size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
+
+ // Dispatch reorder kernel
+ backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ }
// MoE kernel prepare
// Create sub buffer for dst
buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &image_format_buf_dst, &image_desc_buf_dst, NULL, &status);
CL_CHECK(status);
+ if (use_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ // fused reorder + q8_1 quant straight from the original activations
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ // dp4a GEMM
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a;
+ int aidx = 0;
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->q_img));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->d));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(sub_buf_dst);
+ clReleaseMemObject(buf_dst_image);
+ return;
+ }
+
// Set kernel args
int arg_idx = 0;
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q4_0->q_img));
sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
+ // Generic dp4a MoE GEMM
+ {
+ static const char * q5mdp4a_env = getenv("GGML_OPENCL_Q5_MOE_DP4A");
+ const bool q5mdp4a_on = q5mdp4a_env ? (atoi(q5mdp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ const bool use_q5_moe_dp4a = q5mdp4a_on
+ && backend_ctx->kernel_gemm_moe_q8_1_dp4a_q50 != nullptr
+ && extra0_q5_0->scale != nullptr;
+
+ if (use_q5_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ unsigned short map_ratio_q5 = ne20 / ne11;
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio_q5));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ region.origin = offsetd;
+ region.size = ne0 * ne1 * ne2 * sizeof(float);
+ cl_mem dp_sub_buf_dst = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+ cl_image_format dp_ifd = {CL_R, CL_FLOAT};
+ cl_image_desc dp_idd = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne0 * ne1 * ne2), 0,0,0,0,0,0,0, {dp_sub_buf_dst}};
+ cl_mem dp_buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &dp_ifd, &dp_idd, NULL, &status);
+ CL_CHECK(status);
+
+ int ne00i = (int)ne00, ne01i = (int)ne01;
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_q8_1_dp4a_q50;
+ int has_min_q5 = 1;
+ int aidx = 0;
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_0->qs_img));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_0->qh));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_0->scale));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_0->min));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &dp_buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne00i));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne01i));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &has_min_q5));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(dp_sub_buf_dst);
+ clReleaseMemObject(dp_buf_dst_image);
+ return;
+ }
+ }
+
// Create image for reordered src1
// Use pre-allocated placeholder
region.origin = 0;
#endif //GGML_OPENCL_USE_ADRENO_KERNELS
}
case GGML_TYPE_Q8_0: {
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ // MoE GEMM for q8_0 at prefill (ne12>1)
+ // There is no corresponding gemv_moe, so the code path is different here
+ static const char * moe_gemm_q8_env = getenv("GGML_OPENCL_MOE_GEMM_Q8");
+ const bool moe_gemm_q8 = moe_gemm_q8_env
+ ? (atoi(moe_gemm_q8_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ if (moe_gemm_q8 && use_adreno_moe_kernels(backend_ctx, src0) && ne12 > 1) {
+ cl_int status;
+
+ size_t local_size[3] = {64, 2, 1};
+ size_t global_size[3] = {64, 2, 1};
+
+ kernel = backend_ctx->kernel_gemm_moe_q8_0_f32_ns;
+
+ if ((strstr(src0->name, "as") != NULL) || backend_ctx->toggle_reorder) {
+ moe_router_reoerder(backend, src2, ne20);
+ backend_ctx->toggle_reorder = false;
+ }
+
+ cl_mem sub_buf_src1_pre, buf_src1_reordered, image_src1_reordered, sub_buf_dst, buf_dst_image;
+ cl_mem buf_src2, buf_src2_emap;
+
+ cl_buffer_region region;
+ region.origin = 0;
+ region.size = sizeof(int) * max_post_router_tile * n_tile_size;
+ buf_src2 = clCreateSubBuffer(backend_ctx->prealloc_post_router.buffer, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+
+ region.origin = 0;
+ region.size = sizeof(short) * max_post_router_tile;
+ buf_src2_emap = clCreateSubBuffer(backend_ctx->prealloc_emap.buffer, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+
+ // Reorder activations (group tokens by expert into tiles of 32)
+ region.origin = offset1;
+ region.size = ne10 * ne11 * ne12 * sizeof(float);
+ sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+
+ // Generic dp4a MoE GEMM
+ {
+ static const char * q8mdp4a_env = getenv("GGML_OPENCL_Q8_MOE_DP4A");
+ const bool q8mdp4a_on = q8mdp4a_env ? (atoi(q8mdp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ const bool use_q8_moe_dp4a = q8mdp4a_on
+ && backend_ctx->kernel_gemm_moe_q8_1_dp4a_q80 != nullptr
+ && extra0_q8_0->scale != nullptr;
+ if (use_q8_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ unsigned short map_ratio_q8 = ne20 / ne11;
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio_q8));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ // dst image
+ region.origin = offsetd;
+ region.size = ne0 * ne1 * ne2 * sizeof(float);
+ cl_mem dp_sub_buf_dst = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+ cl_image_format dp_ifd = {CL_R, CL_FLOAT};
+ cl_image_desc dp_idd = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne0 * ne1 * ne2), 0,0,0,0,0,0,0, {dp_sub_buf_dst}};
+ cl_mem dp_buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &dp_ifd, &dp_idd, NULL, &status);
+ CL_CHECK(status);
+
+ int ne00i = (int)ne00, ne01i = (int)ne01;
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_q8_1_dp4a_q80;
+ int has_min_q8 = 0;
+ int aidx = 0;
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q8_0->q)); // flat int8 codes [expert][row][K]
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q8_0->scale)); // uniform scale[16]
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q8_0->scale)); // dummy min (has_min=0, unread)
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &dp_buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne00i));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne01i));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &has_min_q8));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(dp_sub_buf_dst);
+ clReleaseMemObject(dp_buf_dst_image);
+ return;
+ }
+ }
+
+ region.origin = 0;
+ region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
+ backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
+ buf_src1_reordered = clCreateSubBuffer(
+ backend_ctx->prealloc_act_trans.buffer, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+ cl_image_format image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
+ cl_image_desc image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
+ image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
+ CL_CHECK(status);
+
+ unsigned short map_ratio = ne20 / ne11;
+ GGML_ASSERT(((map_ratio == 1) || (map_ratio == ne20)) && "Map ratio not supported\n");
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
+
+ size_t reorder_b_local_size[3] = {256, 1, 1};
+ size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
+ backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+
+ // dst image
+ region.origin = offsetd;
+ region.size = ne0 * ne1 * ne2 * sizeof(float);
+ sub_buf_dst = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+ cl_image_format image_format_buf_dst = {CL_R, CL_FLOAT};
+ cl_image_desc image_desc_buf_dst = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne0 * ne1 * ne2), 0,0,0,0,0,0,0, {sub_buf_dst}};
+ buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &image_format_buf_dst, &image_desc_buf_dst, NULL, &status);
+ CL_CHECK(status);
+
+ int arg_idx = 0;
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q8_0->q)); // flat q8_0 quants
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q8_0->d)); // flat q8_0 scales
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &image_src1_reordered));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &buf_dst_image));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(int), &ne01));
+
+ global_size[1] = static_cast<size_t>((ne01 + 63) / 64);
+ global_size[2] = static_cast<size_t>(max_post_router_tile);
+ local_size[1] = 1;
+ local_size[2] = 1;
+
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_size, local_size, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src1_reordered);
+ clReleaseMemObject(image_src1_reordered);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(sub_buf_dst);
+ clReleaseMemObject(buf_dst_image);
+ return;
+ }
+#endif // GGML_OPENCL_USE_ADRENO_KERNELS
#ifdef GGML_OPENCL_SOA_Q
kernel = backend_ctx->kernel_mul_mv_id_q8_0_f32_flat;
if (ne12 == 1) { // for gemv
kernel = backend_ctx->kernel_gemv_moe_q4_k_f32_ns;
+ // Weight-as-texture MoE decode GEMV
+ static const char * moe_decode_wimg_env = getenv("GGML_OPENCL_MOE_DECODE_WIMG");
+ const bool moe_decode_wimg_on = moe_decode_wimg_env
+ ? (atoi(moe_decode_wimg_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ const bool use_moe_decode_wimg = moe_decode_wimg_on
+ && backend_ctx->kernel_gemv_moe_q4_k_f32_ns_wimg != nullptr
+ && extra0_q4_K->q_img != nullptr;
+ if (use_moe_decode_wimg) {
+ kernel = backend_ctx->kernel_gemv_moe_q4_k_f32_ns_wimg;
+ }
+
cl_mem src1_sub_buffer, buf_src1_image, buf_src2;
// create a sub_buffer for src2
// Set kernel args
int arg_idx = 0;
- CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q4_K->q));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), use_moe_decode_wimg ? &extra0_q4_K->q_img : &extra0_q4_K->q));
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q4_K->d));
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q4_K->dm));
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q4_K->s));
backend_ctx->toggle_reorder = false;
}
- cl_mem sub_buf_src1_pre, buf_src1_reordered, image_src1_reordered, sub_buf_dst, buf_dst_image;
+ cl_mem sub_buf_src1_pre, sub_buf_dst, buf_dst_image;
+ cl_mem buf_src1_reordered = nullptr, image_src1_reordered = nullptr;
cl_mem buf_src2, buf_src2_emap;
+ // dp4a (int8) prefill GEMM variant
+ static const char * q4k_moe_dp4a_env = getenv("GGML_OPENCL_Q4K_MOE_DP4A");
+ bool use_moe_dp4a = (q4k_moe_dp4a_env != nullptr)
+ ? (atoi(q4k_moe_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E || backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E);
+ // dot prod has to be available
+ use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
+ // bin kernel takes precedence
+ use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_k_f32_ns_bin == nullptr;
+
cl_buffer_region region;
region.origin = 0;
region.size = sizeof(int) * max_post_router_tile * n_tile_size;
sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
- // Create image for reordered src1
- region.origin = 0;
- region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
- backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
- buf_src1_reordered = clCreateSubBuffer(
- backend_ctx->prealloc_act_trans.buffer,
- 0,
- CL_BUFFER_CREATE_TYPE_REGION,
- ®ion,
- &status);
- CL_CHECK(status);
- cl_image_format image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
- cl_image_desc image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
- if (backend_ctx->kernel_gemm_moe_q4_k_f32_ns_bin) {
- // bin kernel uses slightly different image format
- image_format_buf_src1 = {CL_R, CL_FLOAT};
- image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
- }
- image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
- CL_CHECK(status);
-
unsigned short map_ratio = ne20 / ne11;
GGML_ASSERT(((map_ratio == 1) || (map_ratio == ne20)) && "Map ratio not supported\n");
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
-
- size_t reorder_b_local_size[3] = {256, 1, 1};
- size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
- // Dispatch reorder kernel
- backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ if (!use_moe_dp4a) {
+ // Create image for reordered src1
+ region.origin = 0;
+ region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
+ backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
+ buf_src1_reordered = clCreateSubBuffer(
+ backend_ctx->prealloc_act_trans.buffer,
+ 0,
+ CL_BUFFER_CREATE_TYPE_REGION,
+ ®ion,
+ &status);
+ CL_CHECK(status);
+ cl_image_format image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
+ cl_image_desc image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
+ if (backend_ctx->kernel_gemm_moe_q4_k_f32_ns_bin) {
+ // bin kernel uses slightly different image format
+ image_format_buf_src1 = {CL_R, CL_FLOAT};
+ image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
+ }
+ image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
+ CL_CHECK(status);
+
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
+
+ size_t reorder_b_local_size[3] = {256, 1, 1};
+ size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
+
+ // Dispatch reorder kernel
+ backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ }
// MoE kernel prepare
region.origin = offsetd;
buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &image_format_buf_dst, &image_desc_buf_dst, NULL, &status);
CL_CHECK(status);
+ if (use_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ // fused reorder + q8_1 quant straight from the original
+ // activations (no intermediate f32 reorder buffer)
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ // dp4a GEMM
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_q4_k_q8_1_dp4a;
+ int aidx = 0;
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_K->q_img));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_K->d));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_K->dm));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_K->s));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(sub_buf_dst);
+ clReleaseMemObject(buf_dst_image);
+ return;
+ }
+
// Set kernel args
int arg_idx = 0;
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q4_K->q_img));
sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
+ // Generic dp4a MoE GEMM
+ {
+ static const char * q5kmdp4a_env = getenv("GGML_OPENCL_Q5K_MOE_DP4A");
+ const bool q5kmdp4a_on = q5kmdp4a_env ? (atoi(q5kmdp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ bool use_moe_dp4a = q5kmdp4a_on
+ && backend_ctx->kernel_gemm_moe_q8_1_dp4a_q5k != nullptr
+ && extra0_q5_K->scale != nullptr;
+ // dot prod has to be available
+ use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
+
+ if (use_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ unsigned short map_ratio_q5k = ne20 / ne11;
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio_q5k));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ region.origin = offsetd;
+ region.size = ne0 * ne1 * ne2 * sizeof(float);
+ cl_mem dp_sub_buf_dst = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ CL_CHECK(status);
+ cl_image_format dp_ifd = {CL_R, CL_FLOAT};
+ cl_image_desc dp_idd = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne0 * ne1 * ne2), 0,0,0,0,0,0,0, {dp_sub_buf_dst}};
+ cl_mem dp_buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &dp_ifd, &dp_idd, NULL, &status);
+ CL_CHECK(status);
+
+ int ne00i = (int)ne00, ne01i = (int)ne01;
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_q8_1_dp4a_q5k;
+ int has_min_q5k = 1;
+ int aidx = 0;
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_K->q_img));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_K->qh));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_K->scale));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q5_K->min));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &dp_buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne00i));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne01i));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &has_min_q5k));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(dp_sub_buf_dst);
+ clReleaseMemObject(dp_buf_dst_image);
+ return;
+ }
+ }
+
// Create image for reordered src1
// Use pre-allocated placeholder
region.origin = 0;
} else { // for gemm
kernel = backend_ctx->kernel_gemm_moe_q6_k_f32_ns;
+ if (backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin) {
+ kernel = backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin;
+ }
// Reorder router if called from test-backend-ops or when new router is generated.
// Otherwise reuse the reordered result from previous mul_mat_id call.
backend_ctx->toggle_reorder = false;
}
- cl_mem sub_buf_src1_pre, buf_src1_reordered, image_src1_reordered, sub_buf_dst, buf_dst_image;
+ cl_mem sub_buf_src1_pre, sub_buf_dst, buf_dst_image;
+ cl_mem buf_src1_reordered = nullptr, image_src1_reordered = nullptr;
cl_mem buf_src2, buf_src2_emap;
+ // dp4a (int8) q6_K MoE prefill GEMM
+ static const char * q6k_moe_dp4a_env = getenv("GGML_OPENCL_Q6K_MOE_DP4A");
+ bool use_moe_dp4a = (q6k_moe_dp4a_env != nullptr)
+ ? (atoi(q6k_moe_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E
+ || backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E);
+ // dot prod has to be available
+ use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
+ // bin kernel takes precedence
+ use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin == nullptr;
+
cl_buffer_region region;
region.origin = 0;
region.size = sizeof(int) * max_post_router_tile * n_tile_size;
sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
- // Create image for reordered src1
- region.origin = 0;
- region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
- backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
- buf_src1_reordered = clCreateSubBuffer(
- backend_ctx->prealloc_act_trans.buffer,
- 0,
- CL_BUFFER_CREATE_TYPE_REGION,
- ®ion,
- &status);
- CL_CHECK(status);
- cl_image_format image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
- cl_image_desc image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
- image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
- CL_CHECK(status);
-
unsigned short map_ratio = ne20 / ne11;
GGML_ASSERT(((map_ratio == 1) || (map_ratio == ne20)) && "Map ratio not supported\n");
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
- size_t reorder_b_local_size[3] = {256, 1, 1};
- size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
-
- // Dispatch reorder kernel
- backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ if (!use_moe_dp4a) {
+ // Create image for reordered src1
+ region.origin = 0;
+ region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
+ backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
+ buf_src1_reordered = clCreateSubBuffer(
+ backend_ctx->prealloc_act_trans.buffer,
+ 0,
+ CL_BUFFER_CREATE_TYPE_REGION,
+ ®ion,
+ &status);
+ CL_CHECK(status);
+ cl_image_format image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
+ cl_image_desc image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
+ if (backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin) {
+ // bin kernel uses slightly different image format
+ image_format_buf_src1 = {CL_R, CL_FLOAT};
+ image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
+ }
+ image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
+ CL_CHECK(status);
+
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
+
+ size_t reorder_b_local_size[3] = {256, 1, 1};
+ size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
+
+ // Dispatch reorder kernel
+ backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ }
// MoE kernel prepare
// Create sub buffer for dst
buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &image_format_buf_dst, &image_desc_buf_dst, NULL, &status);
CL_CHECK(status);
+ if (use_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ // fused reorder + q8_1 quant from the original activations
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_q6_k_q8_1_dp4a;
+ int qi = 0;
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &extra0_q6_K->ql_img));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &extra0_q6_K->qh));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &extra0_q6_K->s));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &extra0_q6_K->d));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(dk, qi++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(sub_buf_dst);
+ clReleaseMemObject(buf_dst_image);
+ return;
+ }
+
// Set kernel args
int arg_idx = 0;
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_q6_K->ql_img));
if (ne12 == 1) { // for gemv
kernel = backend_ctx->kernel_gemv_moe_mxfp4_f32_ns;
+ // Weight-as-texture MoE decode GEMV (see q4_K _wimg)
+ static const char * moe_decode_wimg_env = getenv("GGML_OPENCL_MOE_DECODE_WIMG");
+ const bool use_moe_decode_wimg = (moe_decode_wimg_env && (atoi(moe_decode_wimg_env) != 0))
+ && backend_ctx->kernel_gemv_moe_mxfp4_f32_ns_wimg != nullptr
+ && extra0_mxfp4->q_img != nullptr;
+ if (use_moe_decode_wimg) {
+ kernel = backend_ctx->kernel_gemv_moe_mxfp4_f32_ns_wimg;
+ }
+
cl_mem src1_sub_buffer, buf_src1_image, buf_src2;
// create a sub_buffer for src2
// Set kernel args
int arg_idx = 0;
- CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_mxfp4->q));
+ CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), use_moe_decode_wimg ? &extra0_mxfp4->q_img : &extra0_mxfp4->q));
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_mxfp4->e));
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &buf_src1_image));
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &buf_src2));
backend_ctx->toggle_reorder = false;
}
- cl_mem sub_buf_src1_pre, buf_src1_reordered, image_src1_reordered, sub_buf_dst, buf_dst_image;
+ cl_mem sub_buf_src1_pre, sub_buf_dst, buf_dst_image;
+ cl_mem buf_src1_reordered = nullptr, image_src1_reordered = nullptr;
cl_mem buf_src2, buf_src2_emap;
+ // dp4a (int8) prefill GEMM variant
+ static const char * mxfp4_moe_dp4a_env = getenv("GGML_OPENCL_MXFP4_MOE_DP4A");
+ bool use_moe_dp4a = mxfp4_moe_dp4a_env
+ ? (atoi(mxfp4_moe_dp4a_env) != 0)
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+ // dot prod has to be available
+ use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
+ // bin kernel takes precedence
+ use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr;
+
cl_buffer_region region;
region.origin = 0;
region.size = sizeof(int) * max_post_router_tile * n_tile_size;
sub_buf_src1_pre = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
CL_CHECK(status);
- // Create image for reordered src1
- // Use pre-allocated placeholder
- region.origin = 0;
- region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
- backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
- buf_src1_reordered = clCreateSubBuffer(
- backend_ctx->prealloc_act_trans.buffer,
- 0,
- CL_BUFFER_CREATE_TYPE_REGION,
- ®ion,
- &status);
- CL_CHECK(status);
- cl_image_format image_format_buf_src1;
- cl_image_desc image_desc_buf_src1;
- image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
- image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
- if (backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin) {
- // bin kernel uses slightly different image format
- image_format_buf_src1 = {CL_R, CL_FLOAT};
- image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
- }
- image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
- CL_CHECK(status);
-
unsigned short map_ratio = ne20 / ne11;
GGML_ASSERT(((map_ratio == 1) || (map_ratio == ne20)) && "Map ratio not supported\n");
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
- CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
-
- size_t reorder_b_local_size[3] = {256, 1, 1};
- size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
- // Dispatch reorder kernel
- backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ if (!use_moe_dp4a) {
+ // Create image for reordered src1
+ // Use pre-allocated placeholder
+ region.origin = 0;
+ region.size = ne00 * max_post_router_tile * n_tile_size * sizeof(float);
+ backend_ctx->prealloc_act_trans.allocate(backend_ctx->context, region.size);
+ buf_src1_reordered = clCreateSubBuffer(
+ backend_ctx->prealloc_act_trans.buffer,
+ 0,
+ CL_BUFFER_CREATE_TYPE_REGION,
+ ®ion,
+ &status);
+ CL_CHECK(status);
+ cl_image_format image_format_buf_src1;
+ cl_image_desc image_desc_buf_src1;
+ image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
+ image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
+ if (backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin) {
+ // bin kernel uses slightly different image format
+ image_format_buf_src1 = {CL_R, CL_FLOAT};
+ image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
+ }
+ image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
+ CL_CHECK(status);
+
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 2, sizeof(cl_mem), &buf_src1_reordered));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 3, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 4, sizeof(unsigned int), &ne00));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 5, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(backend_ctx->kernel_moe_reorder_b, 6, sizeof(unsigned int), &n_tile_size));
+
+ size_t reorder_b_local_size[3] = {256, 1, 1};
+ size_t reorder_b_global_size[3] = {static_cast<size_t>(((ne00 / 4) + 255) / 256 * 256), static_cast<size_t>(max_post_router_tile * n_tile_size), 1};
+
+ // Dispatch reorder kernel
+ backend_ctx->enqueue_ndrange_kernel(backend_ctx->kernel_moe_reorder_b, 3, reorder_b_global_size, reorder_b_local_size, dst);
+ }
// MoE kernel prepare
// Create sub buffer for dst
buf_dst_image = clCreateImage(backend_ctx->context, CL_MEM_WRITE_ONLY, &image_format_buf_dst, &image_desc_buf_dst, NULL, &status);
CL_CHECK(status);
+ if (use_moe_dp4a) {
+ const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
+ const size_t n_blocks = tok_slots * (ne00 / 32);
+ backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char));
+ backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+ backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half));
+
+ // fused reorder + q8_1 quant straight from the original
+ // activations (no intermediate f32 reorder buffer). mxfp4 has no
+ // min term so the GEMM ignores sa, but reorder_quant still writes it.
+ const cl_uint n_kblocks = (cl_uint)(ne00 / 32);
+ cl_kernel rq = backend_ctx->kernel_moe_reorder_quant_a_q8_1;
+ CL_CHECK(clSetKernelArg(rq, 0, sizeof(cl_mem), &sub_buf_src1_pre));
+ CL_CHECK(clSetKernelArg(rq, 1, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(rq, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(rq, 4, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer));
+ CL_CHECK(clSetKernelArg(rq, 5, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(rq, 6, sizeof(cl_uint), &ne00));
+ CL_CHECK(clSetKernelArg(rq, 7, sizeof(unsigned short), &map_ratio));
+ CL_CHECK(clSetKernelArg(rq, 8, sizeof(cl_uint), &n_tile_size));
+ CL_CHECK(clSetKernelArg(rq, 9, sizeof(cl_uint), &n_kblocks));
+ size_t rq_local[2] = { 32, 1 };
+ size_t rq_global[2] = { (size_t)(((n_kblocks + 31) / 32) * 32), tok_slots };
+ backend_ctx->enqueue_ndrange_kernel(rq, 2, rq_global, rq_local, dst);
+
+ // dp4a GEMM
+ cl_kernel dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a;
+ int aidx = 0;
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->q_img));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->e));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_dst_image));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer)));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(int), &backend_ctx->adreno_use_moe_ragged_dp4));
+
+ size_t dp_global[3] = { 64, (size_t)((ne01 + 63) / 64), (size_t)max_post_router_tile };
+ size_t dp_local[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(dk, 3, dp_global, dp_local, dst);
+
+ clReleaseMemObject(sub_buf_src1_pre);
+ clReleaseMemObject(buf_src2);
+ clReleaseMemObject(buf_src2_emap);
+ clReleaseMemObject(sub_buf_dst);
+ clReleaseMemObject(buf_dst_image);
+ return;
+ }
+
// Set kernel args
int arg_idx = 0;
CL_CHECK(clSetKernelArg(kernel, arg_idx++, sizeof(cl_mem), &extra0_mxfp4->q_img));
}
func = ggml_cl_expm1;
break;
+ case GGML_UNARY_OP_ABS:
+ if (!any_on_device) {
+ return false;
+ }
+ func = ggml_cl_abs;
+ break;
case GGML_UNARY_OP_SOFTPLUS:
if (!any_on_device) {
return false;
--- /dev/null
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+//------------------------------------------------------------------------------
+// abs
+//------------------------------------------------------------------------------
+
+kernel void kernel_abs_f32(
+ global const float * src0,
+ ulong offset0,
+ global float * dst,
+ ulong offsetd
+) {
+ src0 = (global float*)((global char*)src0 + offset0);
+ dst = (global float*)((global char*)dst + offsetd);
+
+ dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f32_4(
+ global const float4 * src0,
+ ulong offset0,
+ global float4 * dst,
+ ulong offsetd
+) {
+ src0 = (global float4*)((global char*)src0 + offset0);
+ dst = (global float4*)((global char*)dst + offsetd);
+
+ dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f16(
+ global const half * src0,
+ ulong offset0,
+ global half * dst,
+ ulong offsetd
+) {
+ src0 = (global half*)((global char*)src0 + offset0);
+ dst = (global half*)((global char*)dst + offsetd);
+
+ dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f16_4(
+ global const half4 * src0,
+ ulong offset0,
+ global half4 * dst,
+ ulong offsetd
+) {
+ src0 = (global half4*)((global char*)src0 + offset0);
+ dst = (global half4*)((global char*)dst + offsetd);
+
+ dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f32_nc(
+ global const char * src0,
+ ulong offset0,
+ global char * dst,
+ ulong offsetd,
+ int ne00,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ ulong nb0,
+ ulong nb1,
+ ulong nb2,
+ ulong nb3
+) {
+ src0 = src0 + offset0;
+ dst = dst + offsetd;
+
+ const int i3 = get_group_id(2);
+ const int i2 = get_group_id(1);
+ const int i1 = get_group_id(0);
+
+ for (int i0 = get_local_id(0); i0 < ne00; i0 += get_local_size(0)) {
+ global const float * x = (global const float *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
+ global float * y = (global float *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
+
+ *y = fabs(*x);
+ }
+}
+
+kernel void kernel_abs_f16_nc(
+ global const char * src0,
+ ulong offset0,
+ global char * dst,
+ ulong offsetd,
+ int ne00,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ ulong nb0,
+ ulong nb1,
+ ulong nb2,
+ ulong nb3
+) {
+ src0 = src0 + offset0;
+ dst = dst + offsetd;
+
+ const int i3 = get_group_id(2);
+ const int i2 = get_group_id(1);
+ const int i1 = get_group_id(0);
+
+ for (int i0 = get_local_id(0); i0 < ne00; i0 += get_local_size(0)) {
+ global const half * x = (global const half *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
+ global half * y = (global half *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
+
+ *y = fabs(*x);
+ }
+}
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);
+ }
+}
#elif defined(cl_qcom_subgroup_shuffle)
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
#define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
#endif
#define ACC_TYPE float
}
}
+// Cluster-parallel variant of _q1_vec_mq_split
+//
+// Tthe baseline keeps one 256B K row in flight per subgroup (32 lanes cooperate
+// on one position, serialized by the reduce+exp chain). This kernel
+// takes q1_split's memory-level parallelism at MQ's read-once traffic:
+// - the 64-lane subgroup is split into FA_CL_NCL clusters of FA_CL_C lanes;
+// - each cluster owns its own KV position stream (positions strided by
+// FA_CL_NCL) with private per-cluster online-softmax state, hence FA_CL_NCL
+// independent K rows in flight per subgroup, no cross-cluster serial chain;
+// - within a cluster, lanes split DK for the dot (cluster-reduce via
+// sub_group_shuffle_xor, steps < FA_CL_C stay inside the cluster) and
+// split DV for o_acc (each lane owns dv indices {lic + FA_CL_C*i} — the
+// same slice for every position, so accumulation is lane-local);
+// - merge stage 1 folds the FA_CL_NCL cluster partials with cross-cluster
+// shuffles (distances >= FA_CL_C); stage 2 is the baseline cross-subgroup
+// LDS merge (o published by cluster 0's lanes, layout-identical to the
+// baseline's sg_o).
+// The KV sweep runs a UNIFORM trip count (max over clusters) with a clamped
+// row address + FA_M_INIT score on the tail — keeps every shuffle convergent
+// (p = exp(FA_M_INIT - m) underflows to 0, so clamped-row reads are inert).
+// Register cost vs baseline: o_acc grows from DV_VEC/64 to DV_VEC/FA_CL_C
+// float4 per lane per head — FA_CL_C=8 / MQ_GQA=4 => 16 float4 (256B).
+
+#ifdef HAS_SUBGROUP_SHUFFLE // cluster reduce/merge needs shuffles; absent -> kernel dropped, dispatch falls back
+
+#ifndef FA_CL_C
+#define FA_CL_C 8
+#endif
+
+// The lane striping requires DK/DV to divide evenly across the cluster;
+// otherwise (e.g. DK=40 with FA_CL_C=16 -> zero-size arrays) compile the
+// kernel out — host soft-create falls back silently.
+#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
+#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
+#define FA_CL_DK (DK_VEC / FA_CL_C) // half4s of K per lane per row
+#define FA_CL_DV (DV_VEC / FA_CL_C) // float4s of o_acc per lane per head
+
+// explicit "half" sub-group attribute routes this fp16-heavy kernel to a slow
+// codegen path on the X1 compiler. X2 keeps the pin: its driver miscompile
+// without it.
+#ifdef FA_C8_NO_SG_PIN
+#define FA_C8_SG_ATTR
+#else
+// REQD_FA_SG pins the HW subgroup on Intel (intel_reqd_sub_group_size(FA_SG),
+// host passes -D FA_SG=32); empty on Adreno. REQD_SUBGROUP_SIZE_64 pins 64 on
+// Adreno; empty on Intel.
+#define FA_C8_SG_ATTR REQD_FA_SG REQD_SUBGROUP_SIZE_64
+#endif
+
+FA_C8_SG_ATTR
+__kernel void flash_attn_f32_f16_q1_vec_mq_split_c8(
+ const global void * q_void, ulong q_offset,
+ const global void * k_void, ulong k_offset,
+ const global void * v_void, ulong v_offset,
+ const float scale,
+ const int n_q,
+ const int n_kv,
+ const int n_head,
+ const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+ const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+ const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+ const float max_bias,
+ const float m0,
+ const float m1,
+ const int n_head_log2,
+ const float logit_softcap,
+ const int n_head_kv,
+ const global void * mask_void,
+ const ulong mask_offset,
+ const ulong mask_nb1,
+ const ulong mask_nb2,
+ const ulong mask_nb3,
+ const int mask_ne2,
+ const int mask_ne3,
+ global float * partial_void,
+ const int n_splits,
+ const int kv_per_split
+) {
+ const int tid = get_local_id(0);
+ const int sgid = tid / Q1_WG_SIZE;
+ const int tid_sg = tid % Q1_WG_SIZE;
+ const int cl = tid_sg / FA_CL_C; // cluster id
+ const int lic = tid_sg % FA_CL_C; // lane in cluster
+ const int kvhead_batch_idx = get_global_id(1);
+ const int split_q_idx = get_global_id(2);
+ const int split_idx = split_q_idx % n_splits;
+ const int q_idx = split_q_idx / n_splits;
+
+ const int batch_idx = kvhead_batch_idx / n_head_kv;
+ const int head_kv_idx = kvhead_batch_idx % n_head_kv;
+
+ const int kv_start = split_idx * kv_per_split;
+ const int kv_end = min(kv_start + kv_per_split, n_kv);
+
+ const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
+
+ if (kv_start >= kv_end) {
+ if (tid == 0) {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + rec_idx * record_stride;
+ rec[0] = FA_M_INIT;
+ rec[1] = 0.0f;
+ }
+ }
+ return;
+ }
+
+ const global char * q_base = (const global char *) q_void + q_offset;
+ const global char * k_base = (const global char *) k_void + k_offset;
+ const global char * v_base = (const global char *) v_void + v_offset;
+
+ // Stage MQ_GQA Q rows in __local once (uniform across WG).
+ __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+ for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE) {
+ const int h = i / DK_VEC;
+ const int k = i % DK_VEC;
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+ const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+ q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ float slope[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
+ }
+
+ const global char * mask_base[MQ_GQA];
+ if (mask_void != NULL) {
+ const int mask_batch_idx = batch_idx % mask_ne3;
+ const global char * mask_base_b = (const global char *) mask_void + mask_offset +
+ mask_batch_idx * mask_nb3 +
+ (ulong) q_idx * mask_nb1;
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const int mask_head_idx = head_idx % mask_ne2;
+ mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
+ }
+ } else {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
+ }
+
+ // Per-CLUSTER online-softmax state (uniform across the cluster's lanes);
+ // o_acc holds this lane's DV slice {lic + FA_CL_C*i}.
+ ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DV];
+ ACC_TYPE m_i[MQ_GQA];
+ ACC_TYPE l_i[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ m_i[h] = FA_M_INIT;
+ l_i[h] = 0.0f;
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
+ }
+
+ const int kv_len = kv_end - kv_start;
+ const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
+ const int kv_lo = kv_start + sgid * kv_per_sg;
+ const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
+
+ // Uniform trip count across the subgroup: every cluster runs n_iter
+ // iterations; tail positions clamp the row address and drop the score to
+ // FA_M_INIT so shuffles stay convergent and the contribution is exactly 0.
+ const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
+ const ulong kv_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
+ const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
+
+ for (int it = 0; it < n_iter; ++it) {
+ const int k_idx = kv_lo + cl + it * FA_CL_NCL;
+ const int valid = k_idx < kv_hi;
+ const int k_safe = valid ? k_idx : (kv_hi - 1);
+
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + kv_row_base + (ulong) k_safe * k_nb1);
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_base + (ulong) k_safe * v_nb1);
+
+ // Dot: this lane covers DK elements {lic + FA_CL_C*i} of the cluster's row.
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DK; ++i) {
+ const int kk = lic + FA_CL_C * i;
+ const ACC_TYPE4 k_vec = CONVERT_KV_ACC4(k_ptr[kk]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + kk], k_vec, dot4[h]);
+ }
+ }
+
+ // Cluster-reduce (xor steps < FA_CL_C stay inside the cluster) + score.
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
+ #pragma unroll
+ for (int step = 1; step < FA_CL_C; step <<= 1) {
+ s += sub_group_shuffle_xor(s, step);
+ }
+ s *= scale;
+ if (mask_base[h] != NULL) {
+ const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+ s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = valid ? s : FA_M_INIT;
+ }
+
+ // Per-cluster online update — identical math to the baseline, but the
+ // serial chain is per cluster (depth n_iter, not kv_per_sg).
+ ACC_TYPE p_h[MQ_GQA];
+ ACC_TYPE sp_h[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const ACC_TYPE m_new = max(m_i[h], score[h]);
+ sp_h[h] = native_exp(m_i[h] - m_new);
+ p_h[h] = native_exp(score[h] - m_new);
+ l_i[h] = l_i[h] * sp_h[h] + p_h[h];
+ m_i[h] = m_new;
+ }
+
+ // V accumulate on this lane's DV slice (p = 0 on tail -> inert).
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) {
+ const ACC_TYPE4 v_vec = CONVERT_KV_ACC4(v_ptr[lic + FA_CL_C * i]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][i] = mad(p_h[h], v_vec, o_acc[h][i] * sp_h[h]);
+ }
+ }
+ }
+
+ // Merge stage 1: fold the FA_CL_NCL cluster partials inside the subgroup.
+ // Lanes with equal lic across clusters hold the SAME dv slice, so a
+ // cross-cluster xor-reduce (distances FA_CL_C..Q1_WG_SIZE/2) sums o
+ // slice-wise; m/l fold the same way. All shuffles are subgroup-convergent.
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE m_c = m_i[h];
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
+ }
+ const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
+ ACC_TYPE l_c = l_i[h] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ l_c += sub_group_shuffle_xor(l_c, step);
+ }
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) {
+ ACC_TYPE4 o = o_acc[h][i] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ o.s0 += sub_group_shuffle_xor(o.s0, step);
+ o.s1 += sub_group_shuffle_xor(o.s1, step);
+ o.s2 += sub_group_shuffle_xor(o.s2, step);
+ o.s3 += sub_group_shuffle_xor(o.s3, step);
+ }
+ o_acc[h][i] = o;
+ }
+ m_i[h] = m_c;
+ l_i[h] = l_c;
+ }
+
+ // Merge stage 2: baseline cross-subgroup LDS merge. Cluster 0's lanes hold
+ // the subgroup's merged o (dv indices {lic + FA_CL_C*i}) — same sg_o layout
+ // and fold loop as q1_vec_mq_split.
+ __local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
+ __local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
+ __local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
+
+ if (tid_sg == 0) {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ sg_m[h][sgid] = m_i[h];
+ sg_l[h][sgid] = l_i[h];
+ }
+ }
+
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ if (cl == 0) {
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) {
+ sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ if (sgid == 0) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+
+ ACC_TYPE m_c = sg_m[h][0];
+ #pragma unroll
+ for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
+ m_c = max(m_c, sg_m[h][s]);
+ }
+ ACC_TYPE l_c = 0.0f;
+ #pragma unroll
+ for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+ l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
+ }
+
+ const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + rec_idx * record_stride;
+ global float4 * rec_o = (global float4 *) (rec + 2);
+
+ if (tid_sg == 0) {
+ rec[0] = (float) m_c;
+ rec[1] = (float) l_c;
+ }
+ for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
+ ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+ const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
+ o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
+ }
+ rec_o[dv_idx] = o_merged;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+}
+
+#endif // DK_VEC/DV_VEC divisible by FA_CL_C
+#endif // HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
+
REQD_SUBGROUP_SIZE_64
__kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img(
const global void * q_void, ulong q_offset,
#elif defined(cl_qcom_subgroup_shuffle)
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
#define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
#endif
// Flash attention: Q=f32, K=q4_0, V=q4_0.
}
}
+// flash_attn_f32_q4_0_q1_vec_mq_split_c8 — cluster-parallel variant of the MQ
+// split, port of flash_attn_f32_f16_q1_vec_mq_split_c8
+// Requires dp4a + subgroup shuffles
+
+#if defined(FA_HAVE_INT_DOT) && defined(HAS_SUBGROUP_SHUFFLE)
+
+#ifndef FA_CL_C
+#define FA_CL_C 8
+#endif
+
+// Lane striping requires DK/DV to divide across the cluster (see f16 c8).
+#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
+#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
+#define FA_CL_DKQ (DK_VEC / FA_CL_C) // K quartets per lane per row
+#define FA_CL_DVQ (DV_VEC / FA_CL_C) // V quartets (o_acc float4s) per lane per head
+
+#ifdef FA_C8_NO_SG_PIN
+#define FA_C8_SG_ATTR_Q4
+#else
+#define FA_C8_SG_ATTR_Q4 REQD_SUBGROUP_SIZE_64
+#endif
+
+FA_C8_SG_ATTR_Q4
+__kernel void flash_attn_f32_q4_0_q1_vec_mq_split_c8(
+ const global void * q_void, ulong q_offset,
+ const global void * k_void, ulong k_offset,
+ const global void * v_void, ulong v_offset,
+ const float scale,
+ const int n_q,
+ const int n_kv,
+ const int n_head,
+ const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+ const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+ const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+ const float max_bias,
+ const float m0,
+ const float m1,
+ const int n_head_log2,
+ const float logit_softcap,
+ const int n_head_kv,
+ const global void * mask_void,
+ const ulong mask_offset,
+ const ulong mask_nb1,
+ const ulong mask_nb2,
+ const ulong mask_nb3,
+ const int mask_ne2,
+ const int mask_ne3,
+ global float * partial_void,
+ const int n_splits,
+ const int kv_per_split
+) {
+ const int tid = get_local_id(0);
+ const int sgid = tid / Q1_WG_SIZE;
+ const int tid_sg = tid % Q1_WG_SIZE;
+ const int cl = tid_sg / FA_CL_C; // cluster id
+ const int lic = tid_sg % FA_CL_C; // lane in cluster
+ const int kvhead_batch_idx = get_global_id(1);
+ const int split_q_idx = get_global_id(2);
+ const int split_idx = split_q_idx % n_splits;
+ const int q_idx = split_q_idx / n_splits;
+
+ const int batch_idx = kvhead_batch_idx / n_head_kv;
+ const int head_kv_idx = kvhead_batch_idx % n_head_kv;
+
+ const int kv_start = split_idx * kv_per_split;
+ const int kv_end = min(kv_start + kv_per_split, n_kv);
+
+ const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
+
+ if (kv_start >= kv_end) {
+ if (tid == 0) {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + rec_idx * record_stride;
+ rec[0] = FA_M_INIT;
+ rec[1] = 0.0f;
+ }
+ }
+ return;
+ }
+
+ const global char * q_base = (const global char *) q_void + q_offset;
+ const global char * k_base = (const global char *) k_void + k_offset;
+ const global char * v_base = (const global char *) v_void + v_offset;
+
+ // Stage MQ_GQA Q rows in __local as float4 (source for the quantize pass).
+ __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+ for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q4) {
+ const int h = i / DK_VEC;
+ const int k = i % DK_VEC;
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+ const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+ q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ // Per-(h, block) int8-packed Q + (qd, q_sum), quantized once per WG.
+ __local uint q_packed_shared[MQ_GQA * DK_Q4_BLOCKS * 8];
+ __local float q_d_shared[MQ_GQA * DK_Q4_BLOCKS];
+ __local int q_sum_shared[MQ_GQA * DK_Q4_BLOCKS];
+ {
+ const int active = MQ_GQA * DK_Q4_BLOCKS;
+ if (tid < active) {
+ const int h = tid / DK_Q4_BLOCKS;
+ const int block_id = tid % DK_Q4_BLOCKS;
+ ACC_TYPE4 q_block[8];
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) q_block[i] = q_shared[h * DK_VEC + block_id * 8 + i];
+ uint packed[8];
+ q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) q_packed_shared[(h * DK_Q4_BLOCKS + block_id) * 8 + i] = packed[i];
+ q_d_shared[h * DK_Q4_BLOCKS + block_id] = info.qd;
+ q_sum_shared[h * DK_Q4_BLOCKS + block_id] = info.q_sum;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ float slope[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
+ }
+
+ const global char * mask_base[MQ_GQA];
+ if (mask_void != NULL) {
+ const int mask_batch_idx = batch_idx % mask_ne3;
+ const global char * mask_base_b = (const global char *) mask_void + mask_offset +
+ mask_batch_idx * mask_nb3 +
+ (ulong) q_idx * mask_nb1;
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const int mask_head_idx = head_idx % mask_ne2;
+ mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
+ }
+ } else {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
+ }
+
+ // Per-CLUSTER online state; o_acc holds this lane's V quartets {lic + FA_CL_C*i}.
+ ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DVQ];
+ ACC_TYPE m_i[MQ_GQA];
+ ACC_TYPE l_i[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ m_i[h] = FA_M_INIT;
+ l_i[h] = 0.0f;
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
+ }
+
+ const int kv_len = kv_end - kv_start;
+ const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
+ const int kv_lo = kv_start + sgid * kv_per_sg;
+ const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
+
+ // Uniform trip count; tail clamps the row address and drops the score to
+ // FA_M_INIT (p underflows to 0) so shuffles stay convergent.
+ const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
+ const ulong k_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
+ const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
+
+ for (int it = 0; it < n_iter; ++it) {
+ const int k_idx = kv_lo + cl + it * FA_CL_NCL;
+ const int valid = k_idx < kv_hi;
+ const int k_safe = valid ? k_idx : (kv_hi - 1);
+
+ const global char * k_row = k_base + k_row_base + (ulong) k_safe * k_nb1;
+ const global char * v_row = v_base + v_row_base + (ulong) k_safe * v_nb1;
+
+ // dp4a K dot over this lane's quartets of the cluster's row.
+ ACC_TYPE lane_contrib[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
+
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DKQ; ++i) {
+ const int qk = lic + FA_CL_C * i;
+ const int block_idx = qk / 8;
+ const int lane_in_block = qk % 8;
+ const int g = lane_in_block & 3;
+ const int shift = (lane_in_block < 4) ? 0 : 4;
+ const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
+ const float kd = vload_half(0, (const global half *)k_block);
+ const global uchar * k_qs = (const global uchar *)(k_block + 2);
+ const uchar b0 = k_qs[g*4 + 0];
+ const uchar b1 = k_qs[g*4 + 1];
+ const uchar b2 = k_qs[g*4 + 2];
+ const uchar b3 = k_qs[g*4 + 3];
+ const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
+ ((uint)((b1 >> shift) & 0x0F)) << 8 |
+ ((uint)((b2 >> shift) & 0x0F)) << 16 |
+ ((uint)((b3 >> shift) & 0x0F)) << 24;
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const uint q_packed_lane = q_packed_shared[(h * DK_Q4_BLOCKS + block_idx) * 8 + lane_in_block];
+ const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
+ const float qd = q_d_shared[h * DK_Q4_BLOCKS + block_idx];
+ const float block_scale = qd * kd;
+ float contrib = (float) raw_dot * block_scale;
+ if (lane_in_block == 0) {
+ const int q_sum_b = q_sum_shared[h * DK_Q4_BLOCKS + block_idx];
+ contrib -= 8.0f * block_scale * (float) q_sum_b;
+ }
+ lane_contrib[h] += contrib;
+ }
+ }
+
+ // Cluster-reduce + score.
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = lane_contrib[h];
+ #pragma unroll
+ for (int step = 1; step < FA_CL_C; step <<= 1) {
+ s += sub_group_shuffle_xor(s, step);
+ }
+ s *= scale;
+ if (mask_base[h] != NULL) {
+ const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+ s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = valid ? s : FA_M_INIT;
+ }
+
+ // Per-cluster online update (serial chain depth n_iter, not kv_per_sg).
+ ACC_TYPE p_h[MQ_GQA];
+ ACC_TYPE sp_h[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const ACC_TYPE m_new = max(m_i[h], score[h]);
+ sp_h[h] = native_exp(m_i[h] - m_new);
+ p_h[h] = native_exp(score[h] - m_new);
+ l_i[h] = l_i[h] * sp_h[h] + p_h[h];
+ m_i[h] = m_new;
+ }
+
+ // V accumulate on this lane's quartets (p = 0 on tail -> inert).
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ const int dv = lic + FA_CL_C * i;
+ const float4 v_v = dequant_q4_0_lane(v_row + (dv / 8) * Q4_0_BLOCK_SIZE, dv % 8);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][i] = mad(p_h[h], v_v, o_acc[h][i] * sp_h[h]);
+ }
+ }
+ }
+
+ // Merge stage 1: fold cluster partials inside the subgroup via shuffles.
+ // Lanes with equal lic across clusters hold the SAME dv slice.
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE m_c = m_i[h];
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
+ }
+ const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
+ ACC_TYPE l_c = l_i[h] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ l_c += sub_group_shuffle_xor(l_c, step);
+ }
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ ACC_TYPE4 o = o_acc[h][i] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ o.s0 += sub_group_shuffle_xor(o.s0, step);
+ o.s1 += sub_group_shuffle_xor(o.s1, step);
+ o.s2 += sub_group_shuffle_xor(o.s2, step);
+ o.s3 += sub_group_shuffle_xor(o.s3, step);
+ }
+ o_acc[h][i] = o;
+ }
+ m_i[h] = m_c;
+ l_i[h] = l_c;
+ }
+
+ // Merge stage 2: baseline cross-subgroup LDS merge (o published by
+ // cluster 0's lanes; layout identical to the baseline sg_o).
+ __local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
+ __local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
+ __local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
+
+ if (tid_sg == 0) {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ sg_m[h][sgid] = m_i[h];
+ sg_l[h][sgid] = l_i[h];
+ }
+ }
+
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ if (cl == 0) {
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ if (sgid == 0) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+
+ ACC_TYPE m_c = sg_m[h][0];
+ #pragma unroll
+ for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
+ m_c = max(m_c, sg_m[h][s]);
+ }
+ ACC_TYPE l_c = 0.0f;
+ #pragma unroll
+ for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+ l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
+ }
+
+ const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + rec_idx * record_stride;
+ global float4 * rec_o = (global float4 *) (rec + 2);
+
+ if (tid_sg == 0) {
+ rec[0] = (float) m_c;
+ rec[1] = (float) l_c;
+ }
+ for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
+ ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+ const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
+ o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
+ }
+ rec_o[dv_idx] = o_merged;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+}
+
+#endif // DK_VEC/DV_VEC divisible by FA_CL_C
+#endif // FA_HAVE_INT_DOT && HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
+
__kernel void flash_attn_f32_q4_0(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
#elif defined(cl_qcom_subgroup_shuffle)
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
#define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
#endif
// Flash attention: Q=f32, K=q8_0, V=q8_0.
}
}
+// flash_attn_f32_q8_0_q1_vec_mq_split_c8 — cluster-parallel variant of the MQ
+// split above, port of the f16/q4_0 c8 kernels
+
+#ifdef HAS_SUBGROUP_SHUFFLE
+
+#ifndef FA_CL_C
+#define FA_CL_C 8
+#endif
+
+// Lane striping requires DK/DV to divide across the cluster (see f16 c8).
+#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
+#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
+#define FA_CL_DKQ (DK_VEC / FA_CL_C) // K quartets per lane per row
+#define FA_CL_DVQ (DV_VEC / FA_CL_C) // V quartets (o_acc float4s) per lane per head
+
+#ifdef FA_C8_NO_SG_PIN
+#define FA_C8_SG_ATTR_Q8
+#else
+#define FA_C8_SG_ATTR_Q8 REQD_SUBGROUP_SIZE_64
+#endif
+
+FA_C8_SG_ATTR_Q8
+__kernel void flash_attn_f32_q8_0_q1_vec_mq_split_c8(
+ const global void * q_void, ulong q_offset,
+ const global void * k_void, ulong k_offset,
+ const global void * v_void, ulong v_offset,
+ const float scale,
+ const int n_q,
+ const int n_kv,
+ const int n_head,
+ const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+ const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+ const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+ const float max_bias,
+ const float m0,
+ const float m1,
+ const int n_head_log2,
+ const float logit_softcap,
+ const int n_head_kv,
+ const global void * mask_void,
+ const ulong mask_offset,
+ const ulong mask_nb1,
+ const ulong mask_nb2,
+ const ulong mask_nb3,
+ const int mask_ne2,
+ const int mask_ne3,
+ global float * partial_void,
+ const int n_splits,
+ const int kv_per_split
+) {
+ const int tid = get_local_id(0);
+ const int sgid = tid / Q1_WG_SIZE;
+ const int tid_sg = tid % Q1_WG_SIZE;
+ const int cl = tid_sg / FA_CL_C; // cluster id
+ const int lic = tid_sg % FA_CL_C; // lane in cluster
+ const int kvhead_batch_idx = get_global_id(1);
+ const int split_q_idx = get_global_id(2);
+ const int split_idx = split_q_idx % n_splits;
+ const int q_idx = split_q_idx / n_splits;
+
+ const int batch_idx = kvhead_batch_idx / n_head_kv;
+ const int head_kv_idx = kvhead_batch_idx % n_head_kv;
+
+ const int kv_start = split_idx * kv_per_split;
+ const int kv_end = min(kv_start + kv_per_split, n_kv);
+
+ const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
+
+ if (kv_start >= kv_end) {
+ if (tid == 0) {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + rec_idx * record_stride;
+ rec[0] = FA_M_INIT;
+ rec[1] = 0.0f;
+ }
+ }
+ return;
+ }
+
+ const global char * q_base = (const global char *) q_void + q_offset;
+ const global char * k_base = (const global char *) k_void + k_offset;
+ const global char * v_base = (const global char *) v_void + v_offset;
+
+ // Stage MQ_GQA Q rows in __local once (uniform across WG).
+ __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+ for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q8) {
+ const int h = i / DK_VEC;
+ const int k = i % DK_VEC;
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+ const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+ q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ float slope[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
+ }
+
+ const global char * mask_base[MQ_GQA];
+ if (mask_void != NULL) {
+ const int mask_batch_idx = batch_idx % mask_ne3;
+ const global char * mask_base_b = (const global char *) mask_void + mask_offset +
+ mask_batch_idx * mask_nb3 +
+ (ulong) q_idx * mask_nb1;
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const int mask_head_idx = head_idx % mask_ne2;
+ mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
+ }
+ } else {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
+ }
+
+ // Per-CLUSTER online state; o_acc holds this lane's V quartets {lic + FA_CL_C*i}.
+ ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DVQ];
+ ACC_TYPE m_i[MQ_GQA];
+ ACC_TYPE l_i[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ m_i[h] = FA_M_INIT;
+ l_i[h] = 0.0f;
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
+ }
+
+ const int kv_len = kv_end - kv_start;
+ const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
+ const int kv_lo = kv_start + sgid * kv_per_sg;
+ const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
+
+ // Uniform trip count; tail clamps the row address and drops the score to
+ // FA_M_INIT (p underflows to 0) so shuffles stay convergent.
+ const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
+ const ulong k_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
+ const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
+
+ for (int it = 0; it < n_iter; ++it) {
+ const int k_idx = kv_lo + cl + it * FA_CL_NCL;
+ const int valid = k_idx < kv_hi;
+ const int k_safe = valid ? k_idx : (kv_hi - 1);
+
+ const global char * k_row = k_base + k_row_base + (ulong) k_safe * k_nb1;
+ const global char * v_row = v_base + v_row_base + (ulong) k_safe * v_nb1;
+
+ // Float-dequant K dot over this lane's quartets of the cluster's row.
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DKQ; ++i) {
+ const int qk = lic + FA_CL_C * i;
+ const float4 k_v = dequant_q8_0_lane(k_row + (qk / 8) * Q8_0_BLOCK_SIZE, qk % 8);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
+ }
+ }
+
+ // Cluster-reduce (xor steps < FA_CL_C stay inside the cluster) + score.
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
+ #pragma unroll
+ for (int step = 1; step < FA_CL_C; step <<= 1) {
+ s += sub_group_shuffle_xor(s, step);
+ }
+ s *= scale;
+ if (mask_base[h] != NULL) {
+ const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+ s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = valid ? s : FA_M_INIT;
+ }
+
+ // Per-cluster online update (serial chain depth n_iter, not kv_per_sg).
+ ACC_TYPE p_h[MQ_GQA];
+ ACC_TYPE sp_h[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const ACC_TYPE m_new = max(m_i[h], score[h]);
+ sp_h[h] = native_exp(m_i[h] - m_new);
+ p_h[h] = native_exp(score[h] - m_new);
+ l_i[h] = l_i[h] * sp_h[h] + p_h[h];
+ m_i[h] = m_new;
+ }
+
+ // V accumulate on this lane's quartets (p = 0 on tail -> inert).
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ const int dv = lic + FA_CL_C * i;
+ const float4 v_v = dequant_q8_0_lane(v_row + (dv / 8) * Q8_0_BLOCK_SIZE, dv % 8);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][i] = mad(p_h[h], v_v, o_acc[h][i] * sp_h[h]);
+ }
+ }
+ }
+
+ // Merge stage 1: fold cluster partials inside the subgroup via shuffles.
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE m_c = m_i[h];
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
+ }
+ const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
+ ACC_TYPE l_c = l_i[h] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ l_c += sub_group_shuffle_xor(l_c, step);
+ }
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ ACC_TYPE4 o = o_acc[h][i] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ o.s0 += sub_group_shuffle_xor(o.s0, step);
+ o.s1 += sub_group_shuffle_xor(o.s1, step);
+ o.s2 += sub_group_shuffle_xor(o.s2, step);
+ o.s3 += sub_group_shuffle_xor(o.s3, step);
+ }
+ o_acc[h][i] = o;
+ }
+ m_i[h] = m_c;
+ l_i[h] = l_c;
+ }
+
+ // Merge stage 2: baseline cross-subgroup LDS merge (o published by
+ // cluster 0's lanes; layout identical to the baseline sg_o).
+ __local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
+ __local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
+ __local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
+
+ if (tid_sg == 0) {
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ sg_m[h][sgid] = m_i[h];
+ sg_l[h][sgid] = l_i[h];
+ }
+ }
+
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ if (cl == 0) {
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ if (sgid == 0) {
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+
+ ACC_TYPE m_c = sg_m[h][0];
+ #pragma unroll
+ for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
+ m_c = max(m_c, sg_m[h][s]);
+ }
+ ACC_TYPE l_c = 0.0f;
+ #pragma unroll
+ for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+ l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
+ }
+
+ const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + rec_idx * record_stride;
+ global float4 * rec_o = (global float4 *) (rec + 2);
+
+ if (tid_sg == 0) {
+ rec[0] = (float) m_c;
+ rec[1] = (float) l_c;
+ }
+ for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
+ ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+ const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
+ o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
+ }
+ rec_o[dv_idx] = o_merged;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+}
+
+#endif // DK_VEC/DV_VEC divisible by FA_CL_C
+#endif // HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
+
__kernel void flash_attn_f32_q8_0(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
shared_b[b_local_offset.y] = bx8_f16.hi;
// Dequantization
- reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * s;
- reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * s;
+ // Cast the e8m0 scale to half to satisfy E17 compilers
+ reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * (half)s;
+ reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * (half)s;
sub_group_barrier(CLK_LOCAL_MEM_FENCE);
shared_b[b_local_offset.y] = bx8_f16.hi;
// Dequantization
- reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * s;
- reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * s;
+ // Cast the e8m0 scale to half to satisfy E17 compilers
+ reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * (half)s;
+ reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * (half)s;
sub_group_barrier(CLK_LOCAL_MEM_FENCE);
--- /dev/null
+#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 { \
+ uint4 a0 = vload4(0, &sh_qa[t][0]); \
+ uint4 a1 = vload4(0, &sh_qa[t][4]); \
+ int raw = 0; \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, 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
+ // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+ const uint vlim = (uint)n_real * 2;
+ for (uint idx = lid; idx < vlim; idx += 64) {
+ const uint t = idx >> 1;
+ const uint h = (idx & 1) << 2; // 0 or 4
+ uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+ vstore4(v, 0, &sh_qa[t][h]);
+ }
+ 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]);
+ }
+ }
+}
--- /dev/null
+#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 { \
+ uint4 a0 = vload4(0, &sh_qa[t][0]); \
+ uint4 a1 = vload4(0, &sh_qa[t][4]); \
+ int raw = 0; \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, 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
+ // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+ const uint vlim = (uint)n_real * 2;
+ for (uint idx = lid; idx < vlim; idx += 64) {
+ const uint t = idx >> 1;
+ const uint h = (idx & 1) << 2; // 0 or 4
+ uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+ vstore4(v, 0, &sh_qa[t][h]);
+ }
+ 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]);
+ }
+ }
+}
--- /dev/null
+#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].
+// The 8 activation uints are read as two 128-bit uint4 loads staged to private (Adreno
+// wants 128-bit local reads, and a __local operand fed straight to the dp4a builtin is
+// slower and can miscompile).
+#define MOE_Q4K_DP4A_T(t) do { \
+ uint4 a0 = vload4(0, &sh_qa[t][0]); \
+ uint4 a1 = vload4(0, &sh_qa[t][4]); \
+ int raw = 0; \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw); \
+ raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, 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 lm
+ // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+ const uint vlim = (uint)n_real * 2;
+ for (uint idx = lid; idx < vlim; idx += 64) {
+ const uint t = idx >> 1;
+ const uint h = (idx & 1) << 2; // 0 or 4
+ uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+ vstore4(v, 0, &sh_qa[t][h]);
+ }
+ 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]);
+ }
+ }
+}
--- /dev/null
+#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 { \
+ uint4 a0 = vload4(0, &sh_qa[t][0]); \
+ uint4 a1 = vload4(0, &sh_qa[t][4]); \
+ const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], a0.s0, a0.s1, a0.s2, a0.s3); \
+ const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], a1.s0, a1.s1, a1.s2, a1.s3); \
+ 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));
+
+ // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+ const uint vlim = (uint)n_real * 2;
+ for (uint idx = lid; idx < vlim; idx += 64) {
+ const uint t = idx >> 1;
+ const uint h = (idx & 1) << 2; // 0 or 4
+ uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+ vstore4(v, 0, &sh_qa[t][h]);
+ }
+ 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]);
+ }
+ }
+}
--- /dev/null
+#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));
+}
--- /dev/null
+#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 { \
+ uint4 a0 = vload4(0, &sh_qa[t][0]); \
+ uint4 a1 = vload4(0, &sh_qa[t][4]); \
+ const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], a0.s0,a0.s1,a0.s2,a0.s3); \
+ const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], a1.s0,a1.s1,a1.s2,a1.s3); \
+ 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)
+
+ // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+ const uint vlim = (uint)n_real * 2;
+ for (uint idx = lid; idx < vlim; idx += 64) {
+ const uint t = idx >> 1;
+ const uint h = (idx & 1) << 2; // 0 or 4
+ uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+ vstore4(v, 0, &sh_qa[t][h]);
+ }
+ 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]);
+ }
+}
--- /dev/null
+#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
+}
--- /dev/null
+#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
+}
#define QK_K 256
#define K_SCALE_SIZE 12
+// scales are transposed: consecutive codes of a row are `stride` apart
inline void get_scale_min_k4(
int j,
global const uchar * q,
+ int stride,
uchar * d,
uchar * m,
uchar mask_d6,
uchar mask_hi2
) {
if (j < 4) {
- *d = q[j] & mask_d6;
- *m = q[j+4] & mask_d6;
+ *d = q[j*stride] & mask_d6;
+ *m = q[(j+4)*stride] & 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);
+ *d = (q[(j+4)*stride] & mask_d4) | ((q[(j-4)*stride] & mask_hi2) >> 2);
+ *m = ((q[(j+4)*stride] >> 4) & mask_d4) | ((q[j*stride] & mask_hi2) >> 2);
}
}
half8 B;
half4 dequantized_weights;
- int num_blocks_K = k / QK_K;
global const ushort * weight_ptr = src0_q + gx_2;
global const half * d_ptr = src0_d + gx_2;
half4 d = vload4(0, d_ptr + sb_idx * m);
half4 dm = vload4(0, dm_ptr + sb_idx * m);
- global const uchar * sc0 = src0_s + (gx_2+0) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
- global const uchar * sc1 = src0_s + (gx_2+1) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
- global const uchar * sc2 = src0_s + (gx_2+2) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
- global const uchar * sc3 = src0_s + (gx_2+3) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
+ global const uchar * sc0 = src0_s + sb_idx * K_SCALE_SIZE * m + (gx_2+0);
+ global const uchar * sc1 = sc0 + 1;
+ global const uchar * sc2 = sc0 + 2;
+ global const uchar * sc3 = sc0 + 3;
uchar sv0, mn0, sv1, mn1, sv2, mn2, sv3, mn3;
- get_scale_min_k4(sub_idx, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
- get_scale_min_k4(sub_idx, sc1, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
- get_scale_min_k4(sub_idx, sc2, &sv2, &mn2, mask_d6, mask_d4, mask_hi2);
- get_scale_min_k4(sub_idx, sc3, &sv3, &mn3, mask_d6, mask_d4, mask_hi2);
+ get_scale_min_k4(sub_idx, sc0, m, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
+ get_scale_min_k4(sub_idx, sc1, m, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
+ get_scale_min_k4(sub_idx, sc2, m, &sv2, &mn2, mask_d6, mask_d4, mask_hi2);
+ get_scale_min_k4(sub_idx, sc3, m, &sv3, &mn3, mask_d6, mask_d4, mask_hi2);
half4 scale = convert_half4(convert_float4(d) * convert_float4((uchar4)(sv0, sv1, sv2, sv3)));
half4 mval = convert_half4(convert_float4(dm) * convert_float4((uchar4)(mn0, mn1, mn2, mn3)));
--- /dev/null
+#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
+
+// scales are transposed: consecutive codes of a row are `stride` apart
+inline void get_scale_min_k4(
+ int j,
+ global const uchar * q,
+ uint stride,
+ uchar * d,
+ uchar * m,
+ uchar mask_d6,
+ uchar mask_d4,
+ uchar mask_hi2
+) {
+ if (j < 4) {
+ *d = q[j*stride] & mask_d6;
+ *m = q[(j+4)*stride] & mask_d6;
+ } else {
+ *d = (q[(j+4)*stride] & mask_d4) | ((q[(j-4)*stride] & mask_hi2) >> 2);
+ *m = ((q[(j+4)*stride] >> 4) & mask_d4) | ((q[j*stride] & 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 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 + sb_idx * K_SCALE_SIZE * (uint)m + rrow;
+ uchar sv, mn;
+ get_scale_min_k4(sub_idx, sc, (uint)m, &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
+
+ __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 + sb_idx * K_SCALE_SIZE * (uint)m + rrow;
+ uchar sv, mn;
+ get_scale_min_k4(sub_idx, sc, (uint)m, &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
+}
--- /dev/null
+#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
+}
--- /dev/null
+#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
+}
--- /dev/null
+#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
+}
--- /dev/null
+#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
+}
}
}
+
+__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;
+ }
+}
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;
+ }
+}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- dst[gid] = totalSum;
+ // Guard the output row. The x-grid is padded to CEIL_DIV(M,wavesize)*wavesize,
+ // so when ne01 is not a multiple of the wave size the tail work-items run past
+ // row ne01 and would overrun dst into the adjacent tensor. No-op / byte-identical
+ // when ne01 is wave-aligned (no padding).
+ if (gid < M) dst[gid] = totalSum;
}
}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows against the padded x-grid tail overrunning dst.
+ // The current shape specializations are all ne01 % 128 == 0 (no padding), so
+ // this is a no-op / byte-identical today; keep it in lockstep with the base kernel.
+ if (gid * 2 + 0 < ne01) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < ne01) dst[gid * 2 + 1] = totalSum.s1;
}
}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
#define NSUBGROUPS 4
#define SUBGROUP_SIZE 64
+// scales are transposed: consecutive codes of a row are `stride` apart
inline void get_scale_min_k4(
int j,
global const uchar * q,
+ uint stride,
uchar * d,
uchar * m,
uchar mask_d6,
uchar mask_hi2
) {
if (j < 4) {
- *d = q[j] & mask_d6;
- *m = q[j+4] & mask_d6;
+ *d = q[j*stride] & mask_d6;
+ *m = q[(j+4)*stride] & 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);
+ *d = (q[(j+4)*stride] & mask_d4) | ((q[(j-4)*stride] & mask_hi2) >> 2);
+ *m = ((q[(j+4)*stride] >> 4) & mask_d4) | ((q[j*stride] & mask_hi2) >> 2);
}
}
uint LINE_STRIDE_A = M / 2;
uint BLOCK_STRIDE_A = NSUBGROUPS * M;
- uint scales_per_row = (K / QK_K) * 12;
private uint4 regA;
private half2 regS;
half2 d = src0_d[gid + sb * LINE_STRIDE_A];
half2 dm = src0_m[gid + sb * LINE_STRIDE_A];
- global const uchar * sc0 = src0_s + 2 * gid * scales_per_row + sb * 12;
- global const uchar * sc1 = src0_s + (2 * gid + 1) * scales_per_row + sb * 12;
+ global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid;
+ global const uchar * sc1 = sc0 + 1;
uchar sv0, mn0, sv1, mn1;
- get_scale_min_k4(j, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
- get_scale_min_k4(j, sc1, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
+ get_scale_min_k4(j, sc0, M, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
+ get_scale_min_k4(j, sc1, M, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
regS = convert_half2(convert_float2(d) * convert_float2((uchar2)(sv0, sv1)));
regM = convert_half2(convert_float2(dm) * convert_float2((uchar2)(mn0, mn1)));
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
// 2 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(totalSum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+ // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+ if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+ if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
}
}
if (grp == 0) {
dst = (global float*)((global char*)dst + offsetd);
- vstore2(total_sum, 0, &(dst[gid * 2]));
+ // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+ // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+ // and would overrun dst into the adjacent tensor (garbage downstream).
+ // No-op / byte-identical when ne01 % 128 == 0 (no padding).
+ if (gid * 2 + 0 < ne01) dst[gid * 2 + 0] = total_sum.s0;
+ if (gid * 2 + 1 < ne01) dst[gid * 2 + 1] = total_sum.s1;
}
}
// 1 outputs per fiber in wave 0
if (groupId == 0) {
dst = (global float*)((global char*)dst + offsetd);
- dst[gid] = totalSum;
+ // Guard the output row. The x-grid is padded to CEIL_DIV(M,wavesize)*wavesize,
+ // so when ne01 is not a multiple of the wave size the tail work-items run past
+ // row ne01 and would overrun dst into the adjacent tensor. No-op / byte-identical
+ // when ne01 is wave-aligned (no padding).
+ if (gid < M) dst[gid] = totalSum;
}
}
--- /dev/null
+// 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);
+}
--- /dev/null
+#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);
+}
global half * x = (global half *) (src0 + offset_src0);
- if (ne00 < 128) {
+ // The vector path below casts the row pointers to half4, which must be 8-byte aligned.
+ // A row address is r0*nb01 + ..., and a permuted or strided src leaves nb01/nb11
+ // unconstrained -- an odd ne00, say, gives a row that is only 2-byte aligned. Every
+ // src1 row this work-item walks is src1_base + r1*nb11, so require both.
+ const ulong src1_base = (ulong) (src1 + (i12)*nb12 + (i13)*nb13);
+ const bool row_aligned = (((ulong) x) & 7) == 0 && (src1_base & 7) == 0 && (nb11 & 7) == 0;
+
+ if (ne00 < 128 || !row_aligned) {
for (int row = 0; row < N_F16_F16; ++row) {
int r1 = rb + row;
if (r1 >= ne11) {
global half * x = (global half *) (src0 + offset_src0);
- if (ne00 < 128) {
+ // The vector path below casts the row pointers to half4/float4, which must be 8- and
+ // 16-byte aligned. A row address is r0*nb01 + ..., and a permuted or strided src leaves
+ // nb01/nb11 unconstrained -- an odd ne00, say, gives a row that is only 2-byte aligned.
+ // Every src1 row this work-item walks is src1_base + r1*nb11, so require both.
+ const ulong src1_base = (ulong) (src1 + (i12)*nb12 + (i13)*nb13);
+ const bool row_aligned = (((ulong) x) & 7) == 0 && (src1_base & 15) == 0 && (nb11 & 15) == 0;
+
+ if (ne00 < 128 || !row_aligned) {
for (int row = 0; row < N_F16_F32; ++row) {
int r1 = rb + row;
if (r1 >= ne11) {
global half * x = (global half *) (src0 + offset_src0);
global float * y = (global float *) (src1 + offset_src1);
+ // The vector path below casts the row pointers to half4/float4, which must be 8- and
+ // 16-byte aligned. A row address is r0*nb01 + ..., and a permuted or strided src leaves
+ // nb01/nb11 unconstrained -- an odd ne00, say, gives a row that is only 2-byte aligned.
+ // Take the vector path only when the rows this work-item touches are actually aligned;
+ // the scalar loop has no such requirement.
+ const bool row_aligned = (((ulong) x) & 7) == 0 && (((ulong) y) & 15) == 0;
+
float sumf = 0;
- if (ne00 < 128) {
+ if (ne00 < 128 || !row_aligned) {
for (int i = get_sub_group_local_id(); i < ne00; i += get_max_sub_group_size()) {
sumf += (float) x[i] * (float) y[i];
}
#elif defined(cl_qcom_subgroup_shuffle)
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
#define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
#endif
// Assumes row size (ne00) is a multiple of 4
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
#ifdef cl_intel_required_subgroup_size
#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
#define INTEL_GPU 1
global ushort * q2 = q1 + 32;
- float4 acc1 = {0.f, 0.f, 0.f, 0.f};
- float4 acc2 = {0.f, 0.f, 0.f, 0.f};
- for (int i = 0; i < 8; i += 2) {
- acc1.s0 += yl[i+0] * (q1[i/2] & 0x000F);
- acc1.s1 += yl[i+1] * (q1[i/2] & 0x0F00);
- acc1.s2 += yl[i+8] * (q1[i/2] & 0x00F0);
- acc1.s3 += yl[i+9] * (q1[i/2] & 0xF000);
- acc2.s0 += yh[i+0] * (q2[i/2] & 0x000F);
- acc2.s1 += yh[i+1] * (q2[i/2] & 0x0F00);
- acc2.s2 += yh[i+8] * (q2[i/2] & 0x00F0);
- acc2.s3 += yh[i+9] * (q2[i/2] & 0xF000);
- }
+ // Load the 4 q1 / 4 q2 quant ushorts as 2 uints each. 16-bit integer ops are
+ // disproportionately slow on the A7X (E031.41) compiler; keeping the dequant
+ // operands in 32-bit registers avoids the ushort path. q1/q2 are 4-byte aligned
+ // (ib*128 + (32*iq+8*ir) bytes; q1 += blk*128 bytes/row). Math is unchanged:
+ // w & 0x0F00 on the low/high halves equals the original ushort mask value.
+ global uint * q1u = (global uint *)q1;
+ global uint * q2u = (global uint *)q2;
+ uint a0 = q1u[0], a1 = q1u[1];
+ uint b0 = q2u[0], b1 = q2u[1];
+ uint w0 = a0 & 0xFFFF, w1 = a0 >> 16, w2 = a1 & 0xFFFF, w3 = a1 >> 16;
+ uint v0 = b0 & 0xFFFF, v1 = b0 >> 16, v2 = b1 & 0xFFFF, v3 = b1 >> 16;
+
+ float4 acc1, acc2;
+ acc1.s0 = yl[0]*(w0&0x000F) + yl[ 2]*(w1&0x000F) + yl[ 4]*(w2&0x000F) + yl[ 6]*(w3&0x000F);
+ acc1.s1 = yl[1]*(w0&0x0F00) + yl[ 3]*(w1&0x0F00) + yl[ 5]*(w2&0x0F00) + yl[ 7]*(w3&0x0F00);
+ acc1.s2 = yl[8]*(w0&0x00F0) + yl[10]*(w1&0x00F0) + yl[12]*(w2&0x00F0) + yl[14]*(w3&0x00F0);
+ acc1.s3 = yl[9]*(w0&0xF000) + yl[11]*(w1&0xF000) + yl[13]*(w2&0xF000) + yl[15]*(w3&0xF000);
+ acc2.s0 = yh[0]*(v0&0x000F) + yh[ 2]*(v1&0x000F) + yh[ 4]*(v2&0x000F) + yh[ 6]*(v3&0x000F);
+ acc2.s1 = yh[1]*(v0&0x0F00) + yh[ 3]*(v1&0x0F00) + yh[ 5]*(v2&0x0F00) + yh[ 7]*(v3&0x0F00);
+ acc2.s2 = yh[8]*(v0&0x00F0) + yh[10]*(v1&0x00F0) + yh[12]*(v2&0x00F0) + yh[14]*(v3&0x00F0);
+ acc2.s3 = yh[9]*(v0&0xF000) + yh[11]*(v1&0xF000) + yh[13]*(v2&0xF000) + yh[15]*(v3&0xF000);
float dall = *d;
float dmin = *dm;
global ushort * q2 = q1 + 32;
- float4 acc1 = {0.f, 0.f, 0.f, 0.f};
- float4 acc2 = {0.f, 0.f, 0.f, 0.f};
- for (int i = 0; i < 8; i += 2) {
- acc1.s0 += yl[i+0] * ((q1[i/2] & 0x000F) + (qh[i+0] & u1_lo ? 16.f : 0.f));
- acc1.s1 += yl[i+1] * ((q1[i/2] & 0x0F00) + (qh[i+1] & u1_lo ? 16.f*256.f : 0.f));
- acc1.s2 += yl[i+8] * ((q1[i/2] & 0x00F0) + (qh[i+0] & u2_lo ? 16.f*16.f : 0.f));
- acc1.s3 += yl[i+9] * ((q1[i/2] & 0xF000) + (qh[i+1] & u2_lo ? 16.f*4096.f: 0.f));
- acc2.s0 += yh[i+0] * ((q2[i/2] & 0x000F) + (qh[i+0] & u1_hi ? 16.f : 0.f));
- acc2.s1 += yh[i+1] * ((q2[i/2] & 0x0F00) + (qh[i+1] & u1_hi ? 16.f*256.f : 0.f));
- acc2.s2 += yh[i+8] * ((q2[i/2] & 0x00F0) + (qh[i+0] & u2_hi ? 16.f*16.f : 0.f));
- acc2.s3 += yh[i+9] * ((q2[i/2] & 0xF000) + (qh[i+1] & u2_hi ? 16.f*4096.f: 0.f));
- }
+ // Load the 4 q1 / 4 q2 quant ushorts as 2 uints each. 16-bit integer ops are
+ // disproportionately slow on the A7X (E031.41) compiler; keeping the dequant
+ // operands in 32-bit registers avoids the ushort path (same fix as q4_K flat).
+ // q1/q2 are 4-byte aligned; w & 0x0F00 on the low/high half of a uint equals the
+ // original ushort mask value, so the math is unchanged. The qh high-bit term is
+ // byte-indexed (qh[0..7]) and left as-is.
+ global uint * q1u = (global uint *)q1;
+ global uint * q2u = (global uint *)q2;
+ uint a0 = q1u[0], a1 = q1u[1], b0 = q2u[0], b1 = q2u[1];
+ uint w0 = a0 & 0xFFFF, w1 = a0 >> 16, w2 = a1 & 0xFFFF, w3 = a1 >> 16;
+ uint v0 = b0 & 0xFFFF, v1 = b0 >> 16, v2 = b1 & 0xFFFF, v3 = b1 >> 16;
+
+ float4 acc1, acc2;
+ acc1.s0 =
+ yl[0]*((w0&0x000F)+(qh[0]&u1_lo?16.f:0.f)) +
+ yl[2]*((w1&0x000F)+(qh[2]&u1_lo?16.f:0.f)) +
+ yl[4]*((w2&0x000F)+(qh[4]&u1_lo?16.f:0.f)) +
+ yl[6]*((w3&0x000F)+(qh[6]&u1_lo?16.f:0.f));
+ acc1.s1 =
+ yl[1]*((w0&0x0F00)+(qh[1]&u1_lo?16.f*256.f:0.f)) +
+ yl[3]*((w1&0x0F00)+(qh[3]&u1_lo?16.f*256.f:0.f)) +
+ yl[5]*((w2&0x0F00)+(qh[5]&u1_lo?16.f*256.f:0.f)) +
+ yl[7]*((w3&0x0F00)+(qh[7]&u1_lo?16.f*256.f:0.f));
+ acc1.s2 =
+ yl[ 8]*((w0&0x00F0)+(qh[0]&u2_lo?16.f*16.f:0.f)) +
+ yl[10]*((w1&0x00F0)+(qh[2]&u2_lo?16.f*16.f:0.f)) +
+ yl[12]*((w2&0x00F0)+(qh[4]&u2_lo?16.f*16.f:0.f)) +
+ yl[14]*((w3&0x00F0)+(qh[6]&u2_lo?16.f*16.f:0.f));
+ acc1.s3 =
+ yl[ 9]*((w0&0xF000)+(qh[1]&u2_lo?16.f*4096.f:0.f)) +
+ yl[11]*((w1&0xF000)+(qh[3]&u2_lo?16.f*4096.f:0.f)) +
+ yl[13]*((w2&0xF000)+(qh[5]&u2_lo?16.f*4096.f:0.f)) +
+ yl[15]*((w3&0xF000)+(qh[7]&u2_lo?16.f*4096.f:0.f));
+ acc2.s0 =
+ yh[0]*((v0&0x000F)+(qh[0]&u1_hi?16.f:0.f)) +
+ yh[2]*((v1&0x000F)+(qh[2]&u1_hi?16.f:0.f)) +
+ yh[4]*((v2&0x000F)+(qh[4]&u1_hi?16.f:0.f)) +
+ yh[6]*((v3&0x000F)+(qh[6]&u1_hi?16.f:0.f));
+ acc2.s1 =
+ yh[1]*((v0&0x0F00)+(qh[1]&u1_hi?16.f*256.f:0.f)) +
+ yh[3]*((v1&0x0F00)+(qh[3]&u1_hi?16.f*256.f:0.f)) +
+ yh[5]*((v2&0x0F00)+(qh[5]&u1_hi?16.f*256.f:0.f)) +
+ yh[7]*((v3&0x0F00)+(qh[7]&u1_hi?16.f*256.f:0.f));
+ acc2.s2 =
+ yh[ 8]*((v0&0x00F0)+(qh[0]&u2_hi?16.f*16.f:0.f)) +
+ yh[10]*((v1&0x00F0)+(qh[2]&u2_hi?16.f*16.f:0.f)) +
+ yh[12]*((v2&0x00F0)+(qh[4]&u2_hi?16.f*16.f:0.f)) +
+ yh[14]*((v3&0x00F0)+(qh[6]&u2_hi?16.f*16.f:0.f));
+ acc2.s3 =
+ yh[ 9]*((v0&0xF000)+(qh[1]&u2_hi?16.f*4096.f:0.f)) +
+ yh[11]*((v1&0xF000)+(qh[3]&u2_hi?16.f*4096.f:0.f)) +
+ yh[13]*((v2&0xF000)+(qh[5]&u2_hi?16.f*4096.f:0.f)) +
+ yh[15]*((v3&0xF000)+(qh[7]&u2_hi?16.f*4096.f:0.f));
float dall = *d;
float dmin = *dm;
--- /dev/null
+#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);
+}
return ®
}
+
+GGML_BACKEND_DL_IMPL(ggml_backend_openvino_reg)
RPC_CMD_HELLO,
RPC_CMD_DEVICE_COUNT,
RPC_CMD_GRAPH_RECOMPUTE,
+ RPC_CMD_MEMSET_TENSOR,
RPC_CMD_COUNT,
};
uint8_t value;
};
+struct rpc_msg_memset_tensor_req {
+ rpc_tensor tensor;
+ uint64_t offset;
+ uint64_t size;
+ uint8_t value;
+};
+
struct rpc_msg_set_tensor_hash_req {
rpc_tensor tensor;
uint64_t offset;
return GGML_STATUS_SUCCESS;
}
+static void ggml_backend_rpc_buffer_memset_tensor(
+ ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
+ ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
+ rpc_msg_memset_tensor_req request = {
+ /* .tensor = */ serialize_tensor(tensor),
+ /* .offset = */ offset,
+ /* .size = */ size,
+ /* .value = */ value,
+ };
+ bool status = send_rpc_cmd(ctx->sock, RPC_CMD_MEMSET_TENSOR, &request, sizeof(request), nullptr, 0);
+ RPC_STATUS_ASSERT(status);
+}
+
static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
rpc_tensor rpc_tensor = serialize_tensor(tensor);
/* .free_buffer = */ ggml_backend_rpc_buffer_free_buffer,
/* .get_base = */ ggml_backend_rpc_buffer_get_base,
/* .init_tensor = */ ggml_backend_rpc_buffer_init_tensor,
- /* .memset_tensor = */ NULL,
+ /* .memset_tensor = */ ggml_backend_rpc_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_rpc_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_rpc_buffer_get_tensor,
/* .set_tensor_2d = */ NULL,
bool buffer_get_base(const rpc_msg_buffer_get_base_req & request, rpc_msg_buffer_get_base_rsp & response);
bool free_buffer(const rpc_msg_free_buffer_req & request);
bool buffer_clear(const rpc_msg_buffer_clear_req & request);
+ bool memset_tensor(const rpc_msg_memset_tensor_req & request);
bool set_tensor(const std::vector<uint8_t> & input);
bool set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response);
bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector<uint8_t> & response);
return true;
}
+bool rpc_server::memset_tensor(const rpc_msg_memset_tensor_req & request) {
+ struct ggml_init_params params {
+ /*.mem_size =*/ ggml_tensor_overhead(),
+ /*.mem_buffer =*/ NULL,
+ /*.no_alloc =*/ true,
+ };
+ ggml_context_ptr ctx_ptr { ggml_init(params) };
+ GGML_ASSERT(ctx_ptr != nullptr);
+ ggml_context * ctx = ctx_ptr.get();
+ ggml_tensor * tensor = deserialize_tensor(ctx, &request.tensor);
+ if (tensor == nullptr || tensor->buffer == nullptr) {
+ GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__);
+ return false;
+ }
+
+ const uint64_t tensor_size = ggml_nbytes(tensor);
+ if (request.offset > tensor_size || request.size > tensor_size - request.offset) {
+ GGML_LOG_ERROR("[%s] tensor region (offset=%" PRIu64 ", size=%" PRIu64 ") out of tensor bounds [0, %" PRIu64 ")\n",
+ __func__, request.offset, request.size, tensor_size);
+ return false;
+ }
+
+ const uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(tensor->buffer);
+ const uint64_t buffer_size = ggml_backend_buffer_get_size(tensor->buffer);
+ if (request.tensor.data < buffer_start) {
+ GGML_LOG_ERROR("[%s] tensor data before buffer start\n", __func__);
+ return false;
+ }
+ const uint64_t data_offset = request.tensor.data - buffer_start;
+ if (data_offset > buffer_size ||
+ request.offset > buffer_size - data_offset ||
+ request.size > buffer_size - data_offset - request.offset) {
+ GGML_LOG_ERROR("[%s] tensor region out of buffer bounds\n", __func__);
+ return false;
+ }
+ if (tensor->buffer->iface.memset_tensor == nullptr) {
+ GGML_LOG_ERROR("[%s] memset not implemented by backend buffer\n", __func__);
+ return false;
+ }
+
+ LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %" PRIu64 ", value: %u\n",
+ __func__, (void *) tensor->buffer, tensor->data, request.offset, request.size, request.value);
+ ggml_backend_tensor_memset(tensor, request.value, request.offset, request.size);
+ return true;
+}
+
ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor) {
// Validate tensor type before using it
if (tensor->type >= GGML_TYPE_COUNT) {
}
break;
}
+ case RPC_CMD_MEMSET_TENSOR: {
+ rpc_msg_memset_tensor_req request;
+ if (!recv_msg(sock, &request, sizeof(request))) {
+ return;
+ }
+ if (!server.memset_tensor(request)) {
+ return;
+ }
+ if (!send_msg(sock, nullptr, 0)) {
+ return;
+ }
+ break;
+ }
case RPC_CMD_SET_TENSOR: {
std::vector<uint8_t> input;
if (!recv_msg(sock, input)) {
-fsycl-targets=spir64_gen
"SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\""
)
+
+ # Pass through parallel job (process) count for parallelising the
+ # `llvm-foreach -- ocloc` invocation for compiling AOT device images.
+ include(ProcessorCount)
+ ProcessorCount(_ggml_sycl_nproc)
+ if (_ggml_sycl_nproc LESS 1)
+ set(_ggml_sycl_nproc 1)
+ endif()
+ set(GGML_SYCL_MAX_PARALLEL_LINK_JOBS ${_ggml_sycl_nproc} CACHE STRING
+ "Parallel ocloc jobs for spir64_gen AOT device-image lowering")
target_link_options(
ggml-sycl PRIVATE
-fsycl-targets=spir64_gen
"SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\""
+ -fsycl-max-parallel-link-jobs=${GGML_SYCL_MAX_PARALLEL_LINK_JOBS}
)
endif()
#include "set_rows.hpp"
#include "ssm_conv.hpp"
#include "softmax.hpp"
+#include "topk-moe.hpp"
#include "tsembd.hpp"
#include "upscale.hpp"
#include "wkv.hpp"
extern int g_ggml_sycl_debug;
extern int g_ggml_sycl_enable_optimize;
+extern int g_ggml_sycl_enable_fusion;
extern int g_ggml_sycl_prioritize_dmmv;
extern int g_ggml_sycl_enable_flash_attention;
extern int g_ggml_sycl_dev2dev_memcpy;
+extern int g_ggml_sycl_fa_onednn;
+extern int g_ggml_sycl_fa_onednn_max_kv;
#if defined(__clang__) && __has_builtin(__builtin_expect)
}
};
-template <typename Layout>
-static void conv2d_dw_kernel(const float * input, const float * kernel, float * output,
+template <typename KernelT, typename Layout>
+static void conv2d_dw_kernel(const float * input, const KernelT * kernel, float * output,
const conv2d_dw_params p, const sycl::nd_item<3> & item_ct1) {
const int global_idx = item_ct1.get_local_id(2) +
item_ct1.get_group(2) * item_ct1.get_local_range(2);
for (int kx = bounds.x_min; kx < bounds.x_max; ++kx) {
const int in_x = dw_calculate_input_coord(out_x, kx, p.stride_x, p.dilation_x, p.padding_x);
acc += input[Layout::input_index(n, c, in_y, in_x, p)] *
- kernel[Layout::kernel_index(c, ky, kx, p)];
+ static_cast<float>(kernel[Layout::kernel_index(c, ky, kx, p)]);
}
}
output[Layout::output_index(n, c, out_y, out_x, p)] = acc;
}
-template <typename Layout>
-static void conv2d_dw_sycl(const float * x_d, const float * w_d, float * y_d,
+template <typename KernelT, typename Layout>
+static void conv2d_dw_sycl(const float * x_d, const KernelT * w_d, float * y_d,
const conv2d_dw_params p, const queue_ptr & stream) {
const int total = p.batches * p.channels * p.out_h * p.out_w;
const int num_blocks = (total + SYCL_CONV2D_DW_BLOCK_SIZE - 1) / SYCL_CONV2D_DW_BLOCK_SIZE;
const sycl::range<3> block_nums(1, 1, num_blocks);
stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
[=](sycl::nd_item<3> item_ct1) {
- conv2d_dw_kernel<Layout>(x_d, w_d, y_d, p, item_ct1);
+ conv2d_dw_kernel<KernelT, Layout>(x_d, w_d, y_d, p, item_ct1);
});
}
const ggml_tensor * kernel = dst->src[0];
const ggml_tensor * input = dst->src[1];
- GGML_ASSERT(kernel->type == GGML_TYPE_F32 && input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32);
+ GGML_ASSERT((kernel->type == GGML_TYPE_F32 || kernel->type == GGML_TYPE_F16) &&
+ input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32);
- const float * w_d = (const float *) kernel->data;
const float * x_d = (const float *) input->data;
float * y_d = (float *) dst->data;
const queue_ptr stream = ctx.stream();
- if (ggml_is_contiguous(input)) {
- conv2d_dw_sycl<dw_whcn_layout>(x_d, w_d, y_d, params, stream);
- } else if (ggml_is_contiguous_channels(input)) {
- conv2d_dw_sycl<dw_cwhn_layout>(x_d, w_d, y_d, params, stream);
+ if (kernel->type == GGML_TYPE_F16) {
+ const sycl::half * w_d = (const sycl::half *) kernel->data;
+ if (ggml_is_contiguous(input)) {
+ conv2d_dw_sycl<sycl::half, dw_whcn_layout>(x_d, w_d, y_d, params, stream);
+ } else if (ggml_is_contiguous_channels(input)) {
+ conv2d_dw_sycl<sycl::half, dw_cwhn_layout>(x_d, w_d, y_d, params, stream);
+ } else {
+ GGML_ABORT("Unsupported memory layout for conv2d_dw");
+ }
} else {
- GGML_ABORT("Unsupported memory layout for conv2d_dw");
+ const float * w_d = (const float *) kernel->data;
+ if (ggml_is_contiguous(input)) {
+ conv2d_dw_sycl<float, dw_whcn_layout>(x_d, w_d, y_d, params, stream);
+ } else if (ggml_is_contiguous_channels(input)) {
+ conv2d_dw_sycl<float, dw_cwhn_layout>(x_d, w_d, y_d, params, stream);
+ } else {
+ GGML_ABORT("Unsupported memory layout for conv2d_dw");
+ }
}
}
typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, dfloat2 & v);
typedef void (*dequantize_kernel_t_reorder)(const void *d, const int64_t ib, const void *qs,
const int iqs, dfloat2 &v);
+typedef void (*dequantize_kernel_f32_t)(const void * vx, const int64_t ib, const int iqs, float & v0, float & v1);
#if QK_K == 256
static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m);
v.y() = (2 * bit_1 - 1) * d;
}
+static __dpct_inline__ void dequantize_q1_0(const void *vx, const int64_t ib,
+ const int iqs, dfloat2 &v) {
+ const block_q1_0 * x = (const block_q1_0 *) vx;
+ const dfloat d = x[ib].d;
+
+ const int bit_index_0 = iqs + 0;
+ const int bit_index_1 = iqs + 1;
+
+ const int bit_0 = (x[ib].qs[bit_index_0 / 8] >> (bit_index_0 % 8)) & 1;
+ const int bit_1 = (x[ib].qs[bit_index_1 / 8] >> (bit_index_1 % 8)) & 1;
+
+ v.x() = (2 * bit_0 - 1) * d;
+ v.y() = (2 * bit_1 - 1) * d;
+}
+
static __dpct_inline__ void dequantize_q4_1(const void *vx, const int64_t ib,
const int iqs, dfloat2 &v) {
const block_q4_1 * x = (const block_q4_1 *) vx;
#endif
}
+static __dpct_inline__ void dequantize_q4_K_f32(const void *vx, const int64_t ib,
+ const int iqs, float &v0, float &v1) {
+#if QK_K == 256
+ const block_q4_K * x = (const block_q4_K *) vx;
+ const sycl::half2 dm = x[ib].dm;
+ const float dall = dm[0];
+ const float dmin = dm[1];
+
+ auto dequantize_one = [&](const int idx) -> float {
+ const int il = idx / 64;
+ const int in = idx % 64;
+ const int is = 2 * il + (in >= 32 ? 1 : 0);
+ const int qsi = 32 * il + (in & 31);
+
+ uint8_t sc;
+ uint8_t m;
+ get_scale_min_k4(is, x[ib].scales, sc, m);
+
+ const float d = dall * sc;
+ const float mn = dmin * m;
+ const uint8_t q = x[ib].qs[qsi];
+ const uint8_t qv = (in >= 32) ? (q >> 4) : (q & 0xF);
+
+ return d * qv - mn;
+ };
+
+ v0 = dequantize_one(iqs + 0);
+ v1 = dequantize_one(iqs + 1);
+#else
+ GGML_ABORT("Q4_K dequantize not supported for QK_K != 256");
+#endif
+}
+
static __dpct_inline__ void dequantize_q2_K(const void *vx, const int64_t ib,
const int iqs, dfloat2 &v) {
#if QK_K == 256
const float d = dall * (sc & 0xF);
const float m = dmin * (sc >> 4);
- return sycl::fma((dfloat) ((q >> (2 * g)) & 3), (dfloat) d, (dfloat) (-m));
+ return (dfloat) d * (dfloat) ((q >> (2 * g)) & 3) - (dfloat) m;
};
v.x() = dequantize_one(iqs + 0);
#endif
}
+static __dpct_inline__ void dequantize_q2_K_f32(const void *vx, const int64_t ib,
+ const int iqs, float &v0, float &v1) {
+#if QK_K == 256
+ const block_q2_K * x = (const block_q2_K *) vx;
+ const float dall = x[ib].dm[0];
+ const float dmin = x[ib].dm[1];
+
+ auto dequantize_one = [&](const int idx) -> float {
+ const int n = idx / 128;
+ const int r = idx % 128;
+ const int g = r / 32;
+ const int l = r % 32;
+ const int is = 8 * n + l / 16;
+
+ const uint8_t q = x[ib].qs[32 * n + l];
+ const uint8_t sc = x[ib].scales[is + 2 * g];
+ const float d = dall * (sc & 0xF);
+ const float m = dmin * (sc >> 4);
+
+ return d * ((q >> (2 * g)) & 3) - m;
+ };
+
+ v0 = dequantize_one(iqs + 0);
+ v1 = dequantize_one(iqs + 1);
+#else
+ GGML_ABORT("Q2_K dequantize not supported for QK_K != 256");
+#endif
+}
+
static __dpct_inline__ void dequantize_q3_K(const void *vx, const int64_t ib,
const int iqs, dfloat2 &v) {
#if QK_K == 256
#endif
}
+static __dpct_inline__ void dequantize_q5_K_f32(const void *vx, const int64_t ib,
+ const int iqs, float &v0, float &v1) {
+#if QK_K == 256
+ const block_q5_K * x = (const block_q5_K *) vx;
+ const float dall = x[ib].dm[0];
+ const float dmin = x[ib].dm[1];
+
+ auto dequantize_one = [&](const int idx) -> float {
+ const int il = idx / 64;
+ const int in = idx % 64;
+ const int is = 2 * il + (in >= 32 ? 1 : 0);
+ const int ir = (in & 31) / 2;
+ const int iq = in & 1;
+
+ const uint8_t q = x[ib].qs[32 * il + 2 * ir + iq];
+ const uint8_t h = x[ib].qh[2 * ir + iq];
+ const uint8_t qv = (in >= 32) ? (q >> 4) : (q & 0xF);
+
+ uint8_t sc;
+ uint8_t m;
+ get_scale_min_k4(is, x[ib].scales, sc, m);
+
+ const float d = dall * sc;
+ const float mn = dmin * m;
+ const uint8_t hm = 1 << (2 * il + (in >= 32 ? 1 : 0));
+
+ return (qv + ((h & hm) ? 16 : 0)) * d - mn;
+ };
+
+ v0 = dequantize_one(iqs + 0);
+ v1 = dequantize_one(iqs + 1);
+#else
+ GGML_ABORT("Q5_K dequantize not supported for QK_K != 256");
+#endif
+}
+
static __dpct_inline__ void dequantize_q6_K(const void *vx, const int64_t ib,
const int iqs, dfloat2 &v) {
#if QK_K == 256
v.y() = d * kvalues_mxfp4[q >> 4] * 0.5f;
}
-static __dpct_inline__ void dequantize_q1_0(const void *vx, const int64_t ib,
- const int iqs, dfloat2 &v) {
- const block_q1_0 * x = (const block_q1_0 *) vx;
- const dfloat d = x[ib].d;
-
- const int bit_index_0 = iqs + 0;
- const int bit_index_1 = iqs + 1;
-
- const int bit_0 = (x[ib].qs[bit_index_0 / 8] >> (bit_index_0 % 8)) & 1;
- const int bit_1 = (x[ib].qs[bit_index_1 / 8] >> (bit_index_1 % 8)) & 1;
-
- v.x() = (2 * bit_0 - 1) * d;
- v.y() = (2 * bit_1 - 1) * d;
-}
-
static __dpct_inline__ void dequantize_nvfp4(const void *vx, const int64_t ib,
const int iqs, dfloat2 &v) {
const block_nvfp4 & xb = ((const block_nvfp4 *) vx)[ib];
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
}
}
+static void dequantize_mul_mat_vec_q2_k_reorder(const void *__restrict__ vx,
+ const float *__restrict__ yy,
+ float *__restrict__ dst,
+ const int ncols, int nrows,
+ const sycl::nd_item<3> &item_ct1) {
+
+ static_assert(16%K_QUANTS_PER_ITERATION == 0, "16 must be divisible by K_QUANTS_PER_ITERATION");
+
+ const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
+ item_ct1.get_local_id(1);
+ if (row >= nrows) return;
+
+ const int num_blocks_per_row = ncols / QK_K;
+ const int ib0 = row*num_blocks_per_row;
+
+ // SOA base pointers for the reordered layout:
+ // [qs: nb * (QK_K/4)] [scales: nb * (QK_K/16)] [dm: nb * sizeof(half2)]
+ const int nb = nrows * num_blocks_per_row;
+ const uint8_t * qs_base = (const uint8_t *)vx;
+ const uint8_t * scales_base = qs_base + (size_t)nb * (QK_K / 4);
+ const sycl::half2 * dm_base = (const sycl::half2 *)(scales_base + (size_t)nb * (QK_K / 16));
+
+ float tmp = 0; // partial sum for thread in warp
+
+#if QK_K == 256
+ const int tid =
+ item_ct1.get_local_id(2) / K_QUANTS_PER_ITERATION; // 0...7 or 0...15
+ const int ix =
+ item_ct1.get_local_id(2) % K_QUANTS_PER_ITERATION; // 0 or 0,1
+
+ const int step = 16/K_QUANTS_PER_ITERATION;
+
+ const int in = tid % step; // 0...15 or 0...7
+
+ const int l0 = K_QUANTS_PER_ITERATION*in; // 0...15 or 0...14 in steps of 2
+
+ uint32_t aux[4];
+ const uint8_t * d = (const uint8_t *)aux;
+ const uint8_t * m = (const uint8_t *)(aux + 2);
+
+ for (int i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) {
+ const int bi = ib0 + i;
+
+ const sycl::half2 dm_val = dm_base[bi];
+ const float dall = dm_val[0];
+ const float dmin = dm_val[1];
+
+ for (int im = 0; im < 2; ++im) {
+ const int q_offset = 32*im + l0;
+ const int s_offset = 8*im;
+ const int y_offset = 128*im + l0;
+
+ const float * y = yy + i * QK_K + y_offset;
+ const uint8_t * q = qs_base + bi * (QK_K / 4) + q_offset;
+
+ const uint32_t * a = (const uint32_t *)(scales_base + bi * (QK_K / 16) + s_offset);
+ aux[0] = a[0] & 0x0f0f0f0f;
+ aux[1] = a[1] & 0x0f0f0f0f;
+ aux[2] = (a[0] >> 4) & 0x0f0f0f0f;
+ aux[3] = (a[1] >> 4) & 0x0f0f0f0f;
+
+ float sum1 = 0, sum2 = 0;
+ for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) {
+ sum1 += y[l+ 0] * d[0] * ((q[l+ 0] >> 0) & 3)
+ + y[l+32] * d[2] * ((q[l+ 0] >> 2) & 3)
+ + y[l+64] * d[4] * ((q[l+ 0] >> 4) & 3)
+ + y[l+96] * d[6] * ((q[l+ 0] >> 6) & 3)
+ + y[l+16] * d[1] * ((q[l+16] >> 0) & 3)
+ + y[l+48] * d[3] * ((q[l+16] >> 2) & 3)
+ + y[l+80] * d[5] * ((q[l+16] >> 4) & 3)
+ +y[l+112] * d[7] * ((q[l+16] >> 6) & 3);
+ sum2 += y[l+ 0] * m[0] + y[l+32] * m[2] + y[l+64] * m[4] + y[ l+96] * m[6]
+ + y[l+16] * m[1] + y[l+48] * m[3] + y[l+80] * m[5] + y[l+112] * m[7];
+
+ }
+ tmp += dall * sum1 - dmin * sum2;
+ }
+ }
+#else
+ GGML_UNUSED(vx);
+ GGML_UNUSED(yy);
+ GGML_UNUSED(ncols);
+ GGML_UNUSED(item_ct1);
+ GGML_ABORT("Q2_K reorder DMMV not supported for QK_K != 256");
+#endif
+
+ // sum up partial sums and write back result
+#pragma unroll
+ for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
+ tmp +=
+ dpct::permute_sub_group_by_xor(item_ct1.get_sub_group(), tmp, mask);
+ }
+
+ if (item_ct1.get_local_id(2) == 0) {
+ dst[row] = tmp;
+ }
+}
+
static void dequantize_mul_mat_vec_q3_k(const void *__restrict__ vx,
const float *__restrict__ yy,
float *__restrict__ dst,
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
static void dequantize_mul_mat_vec_q5_k(const void *__restrict__ vx,
const float *__restrict__ yy,
float *__restrict__ dst,
- const int ncols,
+ const int ncols, int nrows,
const sycl::nd_item<3> &item_ct1) {
- const int row = item_ct1.get_group(2);
+ const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
+ item_ct1.get_local_id(1);
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
const int ncols, int nrows,
const sycl::nd_item<3> &item_ct1) {
- const int row = item_ct1.get_group(2);
+ const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
+ item_ct1.get_local_id(1);
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
const int tid = item_ct1.get_local_id(2) / 2; // 0...15
const int ix = item_ct1.get_local_id(2) % 2;
- const int il = tid/4; // 0...3
- const int ir = tid - 4*il;// 0...3
+ const int il_base = tid/4; // 0...3
+ const int ir = tid - 4*il_base;// 0...3
const int n = 2;
- const int im = il/2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224
- const int in = il%2;
+ const int in = il_base%2;
const int l0 = n*(2*ir + in);
- const int q_offset = 32*im + l0;
- const int y_offset = 64*im + l0;
-
- const uint8_t hm1 = 1 << (2*im);
- const uint8_t hm2 = hm1 << 4;
uint16_t aux[4];
const uint8_t * sc = (const uint8_t *)aux;
for (int i = ix; i < num_blocks_per_row; i += 2) {
const int bi = ib0 + i;
- const uint8_t * ql1 = qs_base + bi * (QK_K / 2) + q_offset;
const uint8_t * qh = qh_base + bi * (QK_K / 8) + l0;
- const float * y1 = yy + i*QK_K + y_offset;
- const float * y2 = y1 + 128;
-
const sycl::half2 dm_val = dm_base[bi];
const float dall = dm_val[0];
const float dmin = dm_val[1];
- const uint16_t * a = (const uint16_t *)(scales_base + bi * K_SCALE_SIZE);
- aux[0] = a[im+0] & kmask1;
- aux[1] = a[im+2] & kmask1;
- aux[2] = ((a[im+4] >> 0) & kmask2) | ((a[im+0] & kmask3) >> 2);
- aux[3] = ((a[im+4] >> 4) & kmask2) | ((a[im+2] & kmask3) >> 2);
-
- sycl::float4 sum = {0.f, 0.f, 0.f, 0.f};
- float smin = 0;
- const uint16_t * q1 = (const uint16_t *)ql1;
- const uint16_t * q2 = q1 + 32;
- q16[0] = q1[0] & 0x0f0f;
- q16[1] = q1[8] & 0x0f0f;
- q16[2] = (q1[0] >> 4) & 0x0f0f;
- q16[3] = (q1[8] >> 4) & 0x0f0f;
- q16[4] = q2[0] & 0x0f0f;
- q16[5] = q2[8] & 0x0f0f;
- q16[6] = (q2[0] >> 4) & 0x0f0f;
- q16[7] = (q2[8] >> 4) & 0x0f0f;
- for (int l = 0; l < n; ++l) {
- sum.x() +=
- y1[l + 0] * (q4[l + 0] + (qh[l + 0] & (hm1 << 0) ? 16 : 0)) +
- y1[l + 16] * (q4[l + 2] + (qh[l + 16] & (hm1 << 0) ? 16 : 0));
- sum.y() +=
- y1[l + 32] * (q4[l + 4] + (qh[l + 0] & (hm1 << 1) ? 16 : 0)) +
- y1[l + 48] * (q4[l + 6] + (qh[l + 16] & (hm1 << 1) ? 16 : 0));
- sum.z() +=
- y2[l + 0] * (q4[l + 8] + (qh[l + 0] & (hm2 << 0) ? 16 : 0)) +
- y2[l + 16] * (q4[l + 10] + (qh[l + 16] & (hm2 << 0) ? 16 : 0));
- sum.w() +=
- y2[l + 32] * (q4[l + 12] + (qh[l + 0] & (hm2 << 1) ? 16 : 0)) +
- y2[l + 48] * (q4[l + 14] + (qh[l + 16] & (hm2 << 1) ? 16 : 0));
- smin += (y1[l] + y1[l+16]) * sc[2] + (y1[l+32] + y1[l+48]) * sc[3]
- + (y2[l] + y2[l+16]) * sc[6] + (y2[l+32] + y2[l+48]) * sc[7];
+ for (int im = 0; im < 2; ++im) {
+ const int q_offset = 32*im + l0;
+ const int y_offset = 64*im + l0;
+
+ const uint8_t hm1 = 1 << (2*im);
+ const uint8_t hm2 = hm1 << 4;
+
+ const uint8_t * ql1 = qs_base + bi * (QK_K / 2) + q_offset;
+ const float * y1 = yy + i*QK_K + y_offset;
+ const float * y2 = y1 + 128;
+
+ const uint16_t * a = (const uint16_t *)(scales_base + bi * K_SCALE_SIZE);
+ aux[0] = a[im+0] & kmask1;
+ aux[1] = a[im+2] & kmask1;
+ aux[2] = ((a[im+4] >> 0) & kmask2) | ((a[im+0] & kmask3) >> 2);
+ aux[3] = ((a[im+4] >> 4) & kmask2) | ((a[im+2] & kmask3) >> 2);
+
+ sycl::float4 sum = {0.f, 0.f, 0.f, 0.f};
+ float smin = 0;
+ const uint16_t * q1 = (const uint16_t *)ql1;
+ const uint16_t * q2 = q1 + 32;
+ q16[0] = q1[0] & 0x0f0f;
+ q16[1] = q1[8] & 0x0f0f;
+ q16[2] = (q1[0] >> 4) & 0x0f0f;
+ q16[3] = (q1[8] >> 4) & 0x0f0f;
+ q16[4] = q2[0] & 0x0f0f;
+ q16[5] = q2[8] & 0x0f0f;
+ q16[6] = (q2[0] >> 4) & 0x0f0f;
+ q16[7] = (q2[8] >> 4) & 0x0f0f;
+ for (int l = 0; l < n; ++l) {
+ sum.x() +=
+ y1[l + 0] * (q4[l + 0] + (qh[l + 0] & (hm1 << 0) ? 16 : 0)) +
+ y1[l + 16] * (q4[l + 2] + (qh[l + 16] & (hm1 << 0) ? 16 : 0));
+ sum.y() +=
+ y1[l + 32] * (q4[l + 4] + (qh[l + 0] & (hm1 << 1) ? 16 : 0)) +
+ y1[l + 48] * (q4[l + 6] + (qh[l + 16] & (hm1 << 1) ? 16 : 0));
+ sum.z() +=
+ y2[l + 0] * (q4[l + 8] + (qh[l + 0] & (hm2 << 0) ? 16 : 0)) +
+ y2[l + 16] * (q4[l + 10] + (qh[l + 16] & (hm2 << 0) ? 16 : 0));
+ sum.w() +=
+ y2[l + 32] * (q4[l + 12] + (qh[l + 0] & (hm2 << 1) ? 16 : 0)) +
+ y2[l + 48] * (q4[l + 14] + (qh[l + 16] & (hm2 << 1) ? 16 : 0));
+ smin += (y1[l] + y1[l+16]) * sc[2] + (y1[l+32] + y1[l+48]) * sc[3]
+ + (y2[l] + y2[l+16]) * sc[6] + (y2[l+32] + y2[l+48]) * sc[7];
+ }
+ tmp += dall * (sum.x() * sc[0] + sum.y() * sc[1] + sum.z() * sc[4] +
+ sum.w() * sc[5]) -
+ dmin * smin;
}
- tmp += dall * (sum.x() * sc[0] + sum.y() * sc[1] + sum.z() * sc[4] +
- sum.w() * sc[5]) -
- dmin * smin;
}
#else
// The reordered Q5_K layout is only produced for QK_K == 256.
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
item_ct1.get_local_id(1);
- if (row > nrows) return;
+ if (row >= nrows) return;
const int num_blocks_per_row = ncols / QK_K;
const int ib0 = row*num_blocks_per_row;
});
}
+static void dequantize_mul_mat_vec_q2_K_sycl_reorder(const void *vx, const float *y,
+ float *dst, const int ncols,
+ const int nrows,
+ dpct::queue_ptr stream) {
+ GGML_ASSERT(ncols % QK_K == 0);
+ const int ny = 2 / K_QUANTS_PER_ITERATION;
+ const int block_num_y = (nrows + ny - 1) / ny;
+ const sycl::range<3> block_nums(1, 1, block_num_y);
+ const sycl::range<3> block_dims(1, ny, WARP_SIZE);
+ stream->parallel_for(
+ sycl::nd_range<3>(block_nums * block_dims, block_dims),
+ [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+ dequantize_mul_mat_vec_q2_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
+ });
+}
+
static void dequantize_mul_mat_vec_q3_K_sycl(const void *vx, const float *y,
float *dst, const int ncols,
const int nrows,
const int nrows,
dpct::queue_ptr stream) {
GGML_ASSERT(ncols % QK_K == 0);
- const sycl::range<3> block_dims(1, 1, WARP_SIZE);
+ const int ny = 2 / K_QUANTS_PER_ITERATION;
+ const int block_num_y = (nrows + ny - 1) / ny;
+ const sycl::range<3> block_nums(1, 1, block_num_y);
+ const sycl::range<3> block_dims(1, ny, WARP_SIZE);
stream->parallel_for(
- sycl::nd_range<3>(sycl::range<3>(1, 1, nrows) * block_dims, block_dims),
+ sycl::nd_range<3>(block_nums * block_dims, block_dims),
[=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- dequantize_mul_mat_vec_q5_k(vx, y, dst, ncols, item_ct1);
+ dequantize_mul_mat_vec_q5_k(vx, y, dst, ncols, nrows, item_ct1);
});
}
}
break;
case GGML_TYPE_Q2_K:
- dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+ if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
+ ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+ dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+ } else {
+ dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+ }
break;
case GGML_TYPE_Q3_K:
if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
}
}
+template<typename T>
+static __dpct_inline__ T op_xielu(T x, float alpha_n, float alpha_p, float beta, float eps) {
+ const float xi = static_cast<float>(x);
+ const float gate_pos = (xi > 0.0f);
+ const float y_pos = alpha_p * xi * xi + beta * xi;
+ const float min_v_eps = sycl::fmin(xi, eps);
+ const float y_neg = (sycl::expm1(min_v_eps) - xi) * alpha_n + beta * xi;
+ const float out = gate_pos * y_pos + (1.0f - gate_pos) * y_neg;
+ return static_cast<T>(out);
+}
+
template<typename T>
static __dpct_inline__ T op_sqr(T x) {
return x * x;
}
}
+template<typename T, typename F>
+static void unary_op_flat_kernel(const T * x, T * dst, const int k, const sycl::nd_item<1> & item_ct1, F func) {
+ SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+ dst[i] = func(x[i]);
+ }
+}
+
template<typename T, typename F>
static void unary_op_generic_kernel(
const T * x,
T * dst,
const int k,
- const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3,
+ const sycl::uint3 ne0_fd, const sycl::uint3 ne1_fd, const sycl::uint3 ne2_fd,
const size_t nb0, const size_t nb1, const size_t nb2, const size_t nb3,
const size_t nbd0, const size_t nbd1, const size_t nbd2, const size_t nbd3,
const sycl::nd_item<1> & item_ct1,
F func) {
- (void) ne3;
+ // 32-bit index math: k is int, so every logical index fits u32. 64-bit integer div/mod is
+ // emulated on Xe and dominates this kernel otherwise, and even the 32-bit divide is worth
+ // avoiding -- the divisors are launch-invariant, so the magic numbers are precomputed
+ // host-side and each division becomes a multiply-high plus a shift.
+ // Byte offsets are widened back to size_t only for the final address math.
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
- const int64_t i0 = i % ne0;
- const int64_t i1 = (i / ne0) % ne1;
- const int64_t i2 = (i / (ne0*ne1)) % ne2;
- const int64_t i3 = i / (ne0*ne1*ne2);
+ sycl::uint2 dm = fast_div_modulo((uint32_t) i, ne0_fd);
+ const uint32_t i0 = dm.y();
+ dm = fast_div_modulo(dm.x(), ne1_fd);
+ const uint32_t i1 = dm.y();
+ dm = fast_div_modulo(dm.x(), ne2_fd);
+ const uint32_t i2 = dm.y();
+ const uint32_t i3 = dm.x();
const char * src_base = (const char *) x;
char * dst_base = (char *) dst;
- const T * srcp = (const T *)(src_base + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3 );
- T * dstp = (T *)(dst_base + i0*nbd0 + i1*nbd1 + i2*nbd2 + i3*nbd3);
+ const T * srcp = (const T *)(src_base + (size_t) i0*nb0 + (size_t) i1*nb1 + (size_t) i2*nb2 + (size_t) i3*nb3 );
+ T * dstp = (T *)(dst_base + (size_t) i0*nbd0 + (size_t) i1*nbd1 + (size_t) i2*nbd2 + (size_t) i3*nbd3);
*dstp = func(*srcp);
}
}
}
+template<typename T>
+static void unary_op_xielu_kernel(const T * x, T * dst, const int k, float alpha_n, float alpha_p, float beta, float eps, const sycl::nd_item<1> &item_ct1) {
+ SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+ dst[i] = op_xielu(x[i], alpha_n, alpha_p, beta, eps);
+ }
+}
+
template<typename T>
static void unary_op_sqr_kernel(const T * x, T * dst, const int k, const sycl::nd_item<1> &item_ct1) {
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
}
template<typename T>
-static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
- const int64_t j0 = (i / n) * o0 + (i % n);
- const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+ const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+ const int64_t j0 = rc.x() * o0 + rc.y();
+ const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
dst[i] = op_gelu(x[j0]) * g[j1];
}
}
template<typename T>
-static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
- const int64_t j0 = (i / n) * o0 + (i % n);
- const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+ const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+ const int64_t j0 = rc.x() * o0 + rc.y();
+ const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
dst[i] = op_relu(x[j0]) * g[j1];
}
}
template<typename T>
-static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
- const int64_t j0 = (i / n) * o0 + (i % n);
- const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+ const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+ const int64_t j0 = rc.x() * o0 + rc.y();
+ const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
dst[i] = op_silu(x[j0]) * g[j1];
}
}
template<typename T>
-static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
- const int64_t j0 = (i / n) * o0 + (i % n);
- const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+ const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+ const int64_t j0 = rc.x() * o0 + rc.y();
+ const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
dst[i] = op_gelu_erf(x[j0]) * g[j1];
}
}
template<typename T>
-static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
- const int64_t j0 = (i / n) * o0 + (i % n);
- const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+ const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+ const int64_t j0 = rc.x() * o0 + rc.y();
+ const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
dst[i] = op_gelu_quick(x[j0]) * g[j1];
}
}
GGML_ASSERT(dst->ne[0] == nc);
GGML_ASSERT(ggml_is_contiguous_1(dst->src[0]));
GGML_ASSERT(ggml_is_contiguous(dst));
+ // The fused GLU kernels index with 32-bit fastdiv, which is exact only for indices below
+ // 2^31. A dst that large is ~8 GB at f32, and the grid sizing already narrows to 32 bits,
+ // so assert the bound rather than carry a second code path for it.
+ GGML_ASSERT(ggml_nelements(dst) < ((int64_t) 1 << 31));
const int32_t swapped = ((const int32_t *) dst->op_params)[1];
void * src0_d = src0->data;
void * src1_d = src1 ? src1->data : src0->data;
const int64_t ne0 = dst->ne[0];
const int64_t ne1 = dst->ne[1];
const int64_t ne2 = dst->ne[2];
- const int64_t ne3 = dst->ne[3];
const size_t nb0 = src0->nb[0];
const size_t nb1 = src0->nb[1];
const size_t nbd2 = dst->nb[2];
const size_t nbd3 = dst->nb[3];
+ // Hot unary ops (FFN/GDN silu, sigmoid, ...) run on contiguous tensors;
+ // skip the strided index math entirely for them.
+ const bool contiguous = ggml_is_contiguous(src0) && ggml_is_contiguous(dst);
+
ggml_sycl_detail::dispatch_ggml_sycl_op_unary(ctx, dst,
[=](const auto* src, auto* dst_ptr, int k_elements, queue_ptr stream) {
const int num_blocks = ceil_div(k_elements, 256);
- stream->parallel_for(
- sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256),
- sycl::range<1>(256)),
- [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- unary_op_generic_kernel(
- src, dst_ptr, k_elements,
- ne0, ne1, ne2, ne3,
- nb0, nb1, nb2, nb3,
- nbd0, nbd1, nbd2, nbd3,
- item_ct1,
- func
- );
- });
+ if (contiguous) {
+ stream->parallel_for(
+ sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256),
+ sycl::range<1>(256)),
+ [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+ unary_op_flat_kernel(src, dst_ptr, k_elements, item_ct1, func);
+ });
+ } else {
+ // Launch-invariant divisors: compute the magic numbers once on the host so the
+ // kernel never issues an integer divide. Only the strided path needs them.
+ const sycl::uint3 ne0_fd = init_fastdiv_values((uint32_t) ne0);
+ const sycl::uint3 ne1_fd = init_fastdiv_values((uint32_t) ne1);
+ const sycl::uint3 ne2_fd = init_fastdiv_values((uint32_t) ne2);
+ stream->parallel_for(
+ sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256),
+ sycl::range<1>(256)),
+ [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+ unary_op_generic_kernel(
+ src, dst_ptr, k_elements,
+ ne0_fd, ne1_fd, ne2_fd,
+ nb0, nb1, nb2, nb3,
+ nbd0, nbd1, nbd2, nbd3,
+ item_ct1,
+ func
+ );
+ });
+ }
});
}
}, min_val, max_val);
}
+static inline void ggml_sycl_op_xielu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+ const float alpha_n = ggml_get_op_params_f32(dst, 1);
+ const float alpha_p = ggml_get_op_params_f32(dst, 2);
+ const float beta = ggml_get_op_params_f32(dst, 3);
+ const float eps = ggml_get_op_params_f32(dst, 4);
+ ggml_sycl_detail::dispatch_ggml_sycl_op_unary(ctx, dst,
+ [](const auto* src, auto* dst_ptr, int k_elements, queue_ptr stream, float alpha_n_arg, float alpha_p_arg, float beta_arg, float eps_arg) {
+ const int num_blocks = ceil_div(k_elements, SYCL_RELU_BLOCK_SIZE);
+ stream->parallel_for(
+ sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(SYCL_RELU_BLOCK_SIZE),
+ sycl::range<1>(SYCL_RELU_BLOCK_SIZE)),
+ [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+ unary_op_xielu_kernel(src, dst_ptr, k_elements, alpha_n_arg, alpha_p_arg, beta_arg, eps_arg, item_ct1);
+ });
+ }, alpha_n, alpha_p, beta, eps);
+}
+
static inline void ggml_sycl_op_floor(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
ggml_sycl_detail::ggml_sycl_op_unary(ctx, dst, [](auto x) {
return op_floor(x);
ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
[](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+ const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
main_stream->parallel_for(
sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+ gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
});
});
}
ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
[](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_RELU_BLOCK_SIZE); // Using RELU block size for reglu
+ const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
main_stream->parallel_for(
sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_RELU_BLOCK_SIZE)),
sycl::range<1>(SYCL_RELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+ gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
});
});
}
ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
[](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_SILU_BLOCK_SIZE); // Using SILU block size for swiglu
+ const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
main_stream->parallel_for(
sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_SILU_BLOCK_SIZE)),
sycl::range<1>(SYCL_SILU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+ gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
});
});
}
ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
[](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+ const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
main_stream->parallel_for(
sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+ gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
});
});
}
ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
[](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+ const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
main_stream->parallel_for(
sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
- gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+ gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
});
});
}
ggml_sycl_op_clamp(ctx, dst);
}
+void ggml_sycl_xielu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+ scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
+ ggml_sycl_op_xielu(ctx, dst);
+}
+
void ggml_sycl_sgn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
ggml_sycl_op_sgn(ctx, dst);
void ggml_sycl_clamp(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+void ggml_sycl_xielu(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
void ggml_sycl_sgn(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
void ggml_sycl_abs(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
--- /dev/null
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "fattn-onednn.hpp"
+#include "fattn-tile.hpp"
+
+// set minimum query length to treat as prefill (32)
+#define GGML_SYCL_FA_ONEDNN_MIN_Q 32
+
+bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
+#if !GGML_SYCL_DNNL
+ GGML_UNUSED(dst);
+ return false;
+#else
+ if (!g_ggml_sycl_fa_onednn) {
+ return false;
+ }
+ // Battlemage (Xe2) only, for now. On other Intel archs oneDNN's fused SDPA returns wrong results
+ // for some shapes (e.g. head_dim=64 on Arc / xe_hpg) -- an oneDNN bug tracked upstream at
+ // https://github.com/uxlfoundation/oneDNN/issues/5510. Remove this hardware limitation once that
+ // is fixed; until then non-BMG archs fall back to the existing FA kernel.
+ const gpu_arch arch = ggml_sycl_info().devices[ggml_sycl_get_device()].hw_info.arch;
+ if (arch != gpu_arch::intel_gpu_bmg_g21 && arch != gpu_arch::intel_gpu_bmg_g31) {
+ return false;
+ }
+ const ggml_tensor * Q = dst->src[0];
+ const ggml_tensor * K = dst->src[1];
+ const ggml_tensor * V = dst->src[2];
+ const ggml_tensor * mask = dst->src[3];
+ const ggml_tensor * sinks = dst->src[4];
+
+ // gate for f16 KV only for now
+ // need to implement quantized KV
+ if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) {
+ return false;
+ }
+ // Optional KV-length ceiling (GGML_SYCL_FA_ONEDNN_MAX_KV, 0 = unlimited). Escape hatch:
+ // very long sequences make the fused SDPA slow enough to risk the xe driver watchdog on
+ // some stacks; past the cap we fall back to the native FA kernel instead.
+ if (g_ggml_sycl_fa_onednn_max_kv > 0 && K->ne[1] > g_ggml_sycl_fa_onednn_max_kv) {
+ return false;
+ }
+ // gate for the following cases
+ // 1. if the oneDNN graph Add node has no input --> skip
+ // 2. types other than f16 need different logical_tensor declaration
+ // 3. the mask must be shape [1, 1, q, seq]
+ // 4. sinks: excludes attention sink (Xiao et al., 2024) that can't be modeled by oneDNN graph
+ if (!mask || mask->type != GGML_TYPE_F16 || mask->ne[2] != 1 || mask->ne[3] != 1 || sinks) {
+ return false;
+ }
+ float max_bias = 0.0f, logit_softcap = 0.0f;
+ memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float));
+ memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
+ if (max_bias != 0.0f || logit_softcap != 0.0f) {
+ return false;
+ }
+ // K and V must share head_dim: the SDPA graph uses a single `d` for both.
+ const int64_t d = K->ne[0];
+ if (V->ne[0] != d || Q->ne[3] != 1) {
+ return false;
+ }
+ // GQA must divide evenly.
+ if (K->ne[2] == 0 || Q->ne[2] % K->ne[2] != 0) {
+ return false;
+ }
+ // Prefill only.
+ if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) {
+ return false;
+ }
+ return true;
+#endif
+}
+
+#if GGML_SYCL_DNNL
+
+#include "dnnl.hpp"
+#include "dnnl_sycl.hpp"
+#include "oneapi/dnnl/dnnl_graph.hpp" // graph API lives only under oneapi/dnnl/, not at the include root
+
+using namespace dnnl;
+using namespace dnnl::graph;
+
+// strided src (f16 or f32) -> contiguous f16 [ne0,ne1,ne2,ne3] (ne0 innermost). nb* are BYTE strides.
+template <typename src_t>
+static void cont_to_f16_sycl(const char * src, sycl::half * dst,
+ int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3,
+ size_t nb1, size_t nb2, size_t nb3, dpct::queue_ptr stream) {
+ const int64_t n = ne0 * ne1 * ne2 * ne3;
+ stream->parallel_for(sycl::range<1>(n), [=](sycl::id<1> ix) {
+ const int64_t gid = ix[0];
+ int64_t i = gid;
+ const int64_t i0 = i % ne0; i /= ne0;
+ const int64_t i1 = i % ne1; i /= ne1;
+ const int64_t i2 = i % ne2; const int64_t i3 = i / ne2;
+ const src_t * p = (const src_t *) (src + i1 * nb1 + i2 * nb2 + i3 * nb3) + i0;
+ dst[gid] = (sycl::half) (*p);
+ });
+}
+
+// oneDNN SDPA out (f16 contiguous [mb,H,q,d]) -> ggml dst (f32 [head_dim,H,n_tok,mb], contiguous).
+static void permute_sdpa_out_sycl(const sycl::half * out, float * dst,
+ int64_t mb, int64_t H, int64_t q, int64_t d, dpct::queue_ptr stream) {
+ const int64_t n = mb * H * q * d;
+ stream->parallel_for(sycl::range<1>(n), [=](sycl::id<1> ix) {
+ const int64_t gid = ix[0];
+ int64_t i = gid;
+ const int64_t e = i % d; i /= d;
+ const int64_t t = i % q; i /= q;
+ const int64_t h = i % H; const int64_t b = i / H;
+ dst[e + h * d + t * d * H + b * d * H * q] = (float) out[gid];
+ });
+}
+
+struct sdpa_partition {
+ compiled_partition cp;
+ std::vector<logical_tensor> ins;
+ logical_tensor out;
+ size_t id_q = 0, id_k = 0, id_v = 0, id_scale = 0, id_mask = 0;
+ bool ok = false;
+};
+
+// Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out.
+// Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t).
+static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) {
+ using ltype = logical_tensor::layout_type;
+ using dt = logical_tensor::data_type;
+ using ldims = logical_tensor::dims;
+ const dt fi = dt::f32, t = dt::f16;
+ const int rep = H / Hkv;
+ const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq},
+ sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d};
+ int64_t id = 0;
+ sdpa_partition E;
+
+ auto query = logical_tensor(id++, t, q_sz, ltype::strided);
+ auto key = logical_tensor(id++, t, kv_sz, ltype::strided);
+ auto score = logical_tensor(id++, fi, s_sz, ltype::strided);
+ auto bmm1 = op(id++, op::kind::MatMul, "bmm1");
+ bmm1.set_attr<bool>(op::attr::transpose_b, true); // key is [.., seq, d]
+ bmm1.add_inputs({query, key}); bmm1.add_outputs({score});
+
+ auto scale = logical_tensor(id++, t, sc, ltype::strided);
+ auto scaled = logical_tensor(id++, fi, s_sz, ltype::strided);
+ auto sdiv = op(id++, op::kind::Divide, "scale_div"); // score / (1/kq_scale) == score * kq_scale
+ sdiv.add_inputs({score, scale}); sdiv.add_outputs({scaled});
+
+ auto mask = logical_tensor(id++, t, msk, ltype::strided);
+ auto masked = logical_tensor(id++, fi, s_sz, ltype::strided);
+ auto madd = op(id++, op::kind::Add, "mask_add");
+ madd.add_inputs({scaled, mask}); madd.add_outputs({masked});
+
+ auto probs = logical_tensor(id++, t, s_sz, ltype::strided);
+ auto smax = op(id++, op::kind::SoftMax, "softmax");
+ smax.set_attr<int64_t>(op::attr::axis, -1);
+ smax.set_attr<std::string>(op::attr::mode, "inf_as_zero");
+ smax.add_inputs({masked}); smax.add_outputs({probs});
+
+ auto value = logical_tensor(id++, t, kv_sz, ltype::strided);
+ // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output
+ // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp).
+ // converted to the f32 ggml dst in the permute below.
+ auto output = logical_tensor(id++, t, o_sz, ltype::strided); // f16 contiguous [mb,Hkv,rep,q,d]
+ auto bmm2 = op(id++, op::kind::MatMul, "bmm2");
+ bmm2.add_inputs({probs, value}); bmm2.add_outputs({output});
+
+ dnnl::graph::graph g(eng.get_kind());
+ g.add_op(bmm1); g.add_op(sdiv); g.add_op(madd); g.add_op(smax); g.add_op(bmm2);
+ g.finalize();
+
+ auto parts = g.get_partitions();
+ if (parts.size() != 1 || !parts[0].is_supported()) {
+ return E; // ok stays false -> caller falls back to TILE
+ }
+ E.ins = parts[0].get_input_ports();
+ E.out = parts[0].get_output_ports()[0];
+ E.cp = parts[0].compile(E.ins, {E.out}, eng);
+ E.out = E.cp.query_logical_tensor(E.out.get_id());
+ E.id_q = query.get_id(); E.id_k = key.get_id(); E.id_v = value.get_id();
+ E.id_scale = scale.get_id(); E.id_mask = mask.get_id();
+ E.ok = true;
+ return E;
+}
+
+void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try {
+ const ggml_tensor * Q = dst->src[0];
+ const ggml_tensor * K = dst->src[1];
+ const ggml_tensor * V = dst->src[2];
+ const ggml_tensor * mask = dst->src[3];
+
+ const int64_t d = K->ne[0]; // head_dim
+ const int64_t seq = K->ne[1]; // n_kv
+ const int64_t Hkv = K->ne[2]; // n_head_kv
+ const int64_t H = Q->ne[2]; // n_head
+ const int64_t q = Q->ne[1]; // n_tok
+ const int64_t mb = Q->ne[3]; // batch (== 1, gated)
+
+ float kq_scale = 1.0f;
+ memcpy(&kq_scale, (const float *) dst->op_params + 0, sizeof(float));
+
+ dpct::queue_ptr stream = ctx.stream();
+ dnnl::engine eng = ctx.engine_dnnl(stream);
+ dnnl::stream strm = ctx.stream_dnnl(stream);
+
+ // cont/cast inputs to contiguous f16 (head-major) -- the layout the fast systolic path wants.
+ ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
+ ggml_sycl_pool_alloc<sycl::half> Kf(ctx.pool(), (size_t) Hkv * seq * d);
+ ggml_sycl_pool_alloc<sycl::half> Vf(ctx.pool(), (size_t) Hkv * seq * d);
+ cont_to_f16_sycl<float> ((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
+ cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf.get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
+ cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf.get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
+
+ // divide-by-(1/scale) reproduces ggml's score *= kq_scale on the proven probe graph.
+ //
+ // The scale must not be uploaded with an async memcpy from a stack local: on the in-order
+ // queue that copy waits behind the K/V staging kernels, and once those take long enough
+ // (n_kv >= ~26k on B70) the host frame is recycled before the copy runs, feeding the SDPA a
+ // garbage scale (output collapses to a repeated token). Write the scalar from a kernel
+ // instead -- the value is captured into the command, so no host memory has to outlive the
+ // call, and the enqueue stays async.
+ const sycl::half scale_h = (sycl::half) (1.0f / kq_scale);
+ ggml_sycl_pool_alloc<sycl::half> scbuf(ctx.pool(), 1);
+ sycl::half * const scale_dev = scbuf.get();
+ stream->single_task([=]() { *scale_dev = scale_h; });
+
+ ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d]
+
+ // compile once per (device, shape), reuse across layers/calls.
+ static std::unordered_map<std::string, sdpa_partition> cache;
+ char keyb[96];
+ snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(),
+ (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d);
+ auto it = cache.find(keyb);
+ if (it == cache.end()) {
+ it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first;
+ }
+ sdpa_partition & E = it->second;
+ // _supported() is authoritative: if it accepted this op the partition must build.
+ // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback.
+ GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape");
+
+ auto id2ptr = [&](size_t r) -> void * {
+ if (r == E.id_q) return Qf.get();
+ if (r == E.id_k) return Kf.get();
+ if (r == E.id_v) return Vf.get();
+ if (r == E.id_scale) return scale_dev;
+ if (r == E.id_mask) return (void *) mask->data;
+ return nullptr;
+ };
+ std::vector<tensor> ti;
+ ti.reserve(E.ins.size());
+ for (auto & lt : E.ins) {
+ ti.emplace_back(lt, eng, id2ptr(lt.get_id()));
+ }
+ tensor to(E.out, eng, outf.get());
+ E.cp.execute(strm, ti, {to});
+
+ permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream);
+ // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA
+ // serializes with the staging kernels before it and the permute/pool reuse after it. The
+ // garbage output formerly blamed on the missing sync here was the scale use-after-return
+ // fixed above. Keep the conservative wait for multi-GPU, where other devices' streams can
+ // race the pool:
+ if (ggml_sycl_info().device_count > 1) {
+ stream->wait_and_throw();
+ }
+}
+catch (const std::exception & e) {
+ // any oneDNN/SYCL failure is non-fatal: fall back to the existing kernel (strictly additive).
+ GGML_LOG_WARN("%s: oneDNN SDPA failed (%s); falling back to TILE kernel\n", __func__, e.what());
+ ggml_sycl_flash_attn_ext_tile(ctx, dst);
+}
+
+#endif // GGML_SYCL_DNNL
--- /dev/null
+#ifndef GGML_SYCL_FATTN_ONEDNN_HPP
+#define GGML_SYCL_FATTN_ONEDNN_HPP
+
+#include "common.hpp"
+
+// Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op
+// (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.)
+bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst);
+
+// Run flash attention through oneDNN's fused xmx SDPA
+// execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure.
+void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
+#endif // GGML_SYCL_FATTN_ONEDNN_HPP
namespace syclex = sycl::ext::oneapi::experimental;
-static int ggml_sycl_fattn_vec_get_nthreads_host(const int cc) {
- return 128;
- GGML_UNUSED(cc);
-}
-
-static constexpr int ggml_sycl_fattn_vec_get_nthreads_device() {
- return 128;
+static int ggml_sycl_fattn_vec_get_nthreads_device(gpu_arch arch) {
+ // Xe2 (Battlemage, Lunar Lake) runs the flash-attention vec kernel best with a 256-thread work group.
+ return (arch == gpu_arch::intel_gpu_bmg_g21 ||
+ arch == gpu_arch::intel_gpu_bmg_g31 ||
+ arch == gpu_arch::intel_gpu_lnl_m) ? 256 : 128;
}
// Currenlty llvm with the amdgcn target dose not support unrolling loops
int type_K,
int type_V,
bool use_logit_softcap,
- int warp_size> // D == head size
+ int warp_size,
+ int nthreads> // D == head size
static void flash_attn_ext_vec(const char* __restrict__ Q,
const char* __restrict__ K,
const char* __restrict__ V,
constexpr int nthreads_KQ_q = (D/4 < warp_size ? D/4 : warp_size);
constexpr int nthreads_V_q = (D/4 < warp_size ? D/4 : warp_size);
- constexpr int nthreads = ggml_sycl_fattn_vec_get_nthreads_device();
constexpr int nthreads_KQ = type_K == GGML_TYPE_F16 ? 128 / cpy_nb : nthreads_KQ_q;
constexpr int nthreads_V = type_V == GGML_TYPE_F16 ? 128 / cpy_nb : nthreads_V_q;
#endif // __clang__
+
template <int D, int cols_per_block, int type_K, int type_V, bool use_logit_softcap>
void ggml_sycl_flash_attn_ext_vec_case_impl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
- const int warp_size = WARP_16_SIZE; //better performance than WARP_32_SIZE
-
- const int cc = ggml_sycl_info().devices[ggml_sycl_get_device()].cc;
-
- const int nthreads = ggml_sycl_fattn_vec_get_nthreads_host(cc);
- const int nwarps = nthreads / warp_size;
+ constexpr int warp_size = WARP_16_SIZE; //better performance than WARP_32_SIZE
const bool need_f16_K = type_K == GGML_TYPE_F16;
const bool need_f16_V = type_V == GGML_TYPE_F16;
constexpr size_t nbytes_shared = 0;
- launch_fattn<D, cols_per_block, 1,
- flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
- use_logit_softcap, warp_size>, warp_size>(
- ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+ const auto arch = ggml_sycl_info().devices[ctx.device].hw_info.arch;
+ const int nthreads = ggml_sycl_fattn_vec_get_nthreads_device(arch);
+ // 256 threads would overflow the 64 KB work-group local memory at D == 512, so keep 128 there.
+ if (D <= 256 && nthreads == 256) {
+ constexpr int nthreads_hw = 256;
+ constexpr int nwarps = nthreads_hw / warp_size;
+ launch_fattn<D, cols_per_block, 1,
+ flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
+ use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
+ ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+ } else {
+ constexpr int nthreads_hw = 128;
+ constexpr int nwarps = nthreads_hw / warp_size;
+ launch_fattn<D, cols_per_block, 1,
+ flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
+ use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
+ ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+ }
}
template <int D, int type_K, int type_V>
#include "fattn-tile.hpp"
#include "fattn-vec.hpp"
#include "fattn.hpp"
+#include "fattn-onednn.hpp"
#define FATTN_VEC_CASE(D, type_K, type_V) \
enum best_fattn_kernel {
BEST_FATTN_KERNEL_NONE = 0,
BEST_FATTN_KERNEL_VEC = 100,
+ BEST_FATTN_KERNEL_ONEDNN = 150, // added enum for onednn==150
BEST_FATTN_KERNEL_TILE = 200,
};
// For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes:
const bool can_use_vector_kernel = Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0;
- // Todo: Use the XMX kernel if possible:
+ // Fused-XMX path: oneDNN Graph SDPA (flash attention). Strictly
+ // additive -- taken only when statically supported, otherwise falls through to VEC/TILE below.
+ if (ggml_sycl_flash_attn_ext_onednn_supported(dst)) {
+ return BEST_FATTN_KERNEL_ONEDNN;
+ }
// If there are no tensor cores available, use the generic tile kernel:
if (can_use_vector_kernel) {
switch (ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst)) {
case BEST_FATTN_KERNEL_NONE:
GGML_ABORT("Not support Flash-Attention");
+ case BEST_FATTN_KERNEL_ONEDNN:
+ // guarded: ggml_sycl_flash_attn_ext_onednn() is only defined under GGML_SYCL_DNNL;
+ // the reference must be compiled out here or the GGML_SYCL_DNNL=0 build fails to link.
+#if GGML_SYCL_DNNL
+ ggml_sycl_flash_attn_ext_onednn(ctx, dst);
+#endif
+ break;
case BEST_FATTN_KERNEL_TILE:
ggml_sycl_flash_attn_ext_tile(ctx, dst);
break;
dst_row[iybs + iqs + y_offset] = v.y();
}
+template<int qk, int qr, dequantize_kernel_f32_t dequantize_kernel, typename dst_t>
+static void k_get_rows_f32(
+ const void * src0, const int32_t * src1, dst_t * dst,
+ int64_t ne00,
+ int64_t ne12,
+ size_t s1, size_t s2, size_t s3,
+ size_t nb01, size_t nb02, size_t nb03,
+ size_t s10, size_t s11, size_t s12,
+ const sycl::nd_item<3> &item_ct1) {
+
+ const int i00 = (item_ct1.get_group(2) * item_ct1.get_local_range(2) +
+ item_ct1.get_local_id(2)) *
+ 2;
+ const int i10 = item_ct1.get_local_range(1) * item_ct1.get_group(1) +
+ item_ct1.get_local_id(1);
+ const int i11 = (item_ct1.get_group(0) * item_ct1.get_local_range(0) +
+ item_ct1.get_local_id(0)) /
+ ne12;
+ const int i12 = (item_ct1.get_group(0) * item_ct1.get_local_range(0) +
+ item_ct1.get_local_id(0)) %
+ ne12;
+
+ if (i00 >= ne00) {
+ return;
+ }
+
+ const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+ dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+ const void * src0_row = (const char *)src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+ const int ib = i00/qk;
+ const int iqs = (i00%qk)/qr;
+ const int iybs = i00 - i00%qk;
+ const int y_offset = qr == 1 ? 1 : qk/2;
+
+ float v0;
+ float v1;
+ dequantize_kernel(src0_row, ib, iqs, v0, v1);
+
+ dst_row[iybs + iqs + 0] = (dst_t) v0;
+ dst_row[iybs + iqs + y_offset] = (dst_t) v1;
+}
+
template<typename src0_t, typename dst_t>
static void k_get_rows_float(
const src0_t * src0, const int32_t * src1, dst_t * dst,
GGML_UNUSED(ctx);
}
+template <int qk, int qr, dequantize_kernel_f32_t dq>
+static void get_rows_sycl_f32(ggml_backend_sycl_context & ctx, const ggml_tensor *src0, const ggml_tensor *src1,
+ ggml_tensor *dst, const void *src0_dd,
+ const int32_t *src1_dd, float *dst_dd,
+ queue_ptr stream) {
+
+ GGML_TENSOR_BINARY_OP_LOCALS
+
+ const sycl::range<3> block_dims(1, 1, SYCL_GET_ROWS_BLOCK_SIZE);
+ const int block_num_x = (ne00 + 2*SYCL_GET_ROWS_BLOCK_SIZE - 1) / (2*SYCL_GET_ROWS_BLOCK_SIZE);
+ const sycl::range<3> block_nums(ne11 * ne12, ne10, block_num_x);
+
+ const size_t s1 = nb1 / ggml_element_size(dst);
+ const size_t s2 = nb2 / ggml_element_size(dst);
+ const size_t s3 = nb3 / ggml_element_size(dst);
+
+ const size_t s10 = nb10 / ggml_element_size(src1);
+ const size_t s11 = nb11 / ggml_element_size(src1);
+ const size_t s12 = nb12 / ggml_element_size(src1);
+
+ GGML_ASSERT(ne00 % 2 == 0);
+
+ stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
+ [=](sycl::nd_item<3> item_ct1) {
+ k_get_rows_f32<qk, qr, dq>(
+ src0_dd, src1_dd, dst_dd, ne00, ne12, s1, s2,
+ s3, nb01, nb02, nb03, s10, s11, s12, item_ct1);
+ });
+
+ GGML_UNUSED(dst);
+ GGML_UNUSED(ctx);
+}
+
template <typename src0_t, typename dst_t>
static void get_rows_sycl_float(ggml_backend_sycl_context & ctx, const ggml_tensor *src0,
const ggml_tensor *src1, ggml_tensor *dst,
src1_i32, (float *)dst->data, ctx.stream());
break;
case GGML_TYPE_Q2_K:
- get_rows_sycl<QK_K, 1, dequantize_q2_K>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
+ get_rows_sycl_f32<QK_K, 1, dequantize_q2_K_f32>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
src1_i32, (float *)dst->data, ctx.stream());
break;
case GGML_TYPE_Q3_K:
src1_i32, (float *)dst->data, ctx.stream());
break;
case GGML_TYPE_Q4_K:
- get_rows_sycl<QK_K, 1, dequantize_q4_K>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
+ get_rows_sycl_f32<QK_K, 1, dequantize_q4_K_f32>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
src1_i32, (float *)dst->data, ctx.stream());
break;
case GGML_TYPE_Q5_0:
src1_i32, (float *)dst->data, ctx.stream());
break;
case GGML_TYPE_Q5_K:
- get_rows_sycl<QK_K, 1, dequantize_q5_K>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
+ get_rows_sycl_f32<QK_K, 1, dequantize_q5_K_f32>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
src1_i32, (float *)dst->data, ctx.stream());
break;
case GGML_TYPE_Q6_K:
int g_ggml_sycl_enable_optimize = 1;
int g_ggml_sycl_enable_graph = 0;
int g_ggml_sycl_enable_dnn = 1;
+int g_ggml_sycl_fa_onednn = 1;
+int g_ggml_sycl_fa_onednn_max_kv = 0;
int g_ggml_sycl_enable_vmm = 1;
+int g_ggml_sycl_enable_fusion = 1;
int g_ggml_sycl_prioritize_dmmv = 0;
int g_ggml_sycl_use_async_mem_op = 0;
int g_ggml_sycl_use_async_mem_op_requested = 1;
g_ggml_sycl_enable_optimize = ggml_sycl_get_env("GGML_SYCL_ENABLE_OPT", 1);
g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0);
g_ggml_sycl_enable_dnn = ggml_sycl_get_env("GGML_SYCL_ENABLE_DNN", 1);
+ g_ggml_sycl_fa_onednn = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN", 1);
+ g_ggml_sycl_fa_onednn_max_kv = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN_MAX_KV", 0);
g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1);
+ g_ggml_sycl_enable_fusion = ggml_sycl_get_env("GGML_SYCL_ENABLE_FUSION", 1);
g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
g_ggml_sycl_dev2dev_memcpy = ggml_sycl_get_env("GGML_SYCL_DEV2DEV_MEMCPY", DEV2DEV_MEMCPY_SYCL);
#if defined(GGML_SYCL_DNNL)
GGML_LOG_INFO(" GGML_SYCL_ENABLE_DNN: %d\n", g_ggml_sycl_enable_dnn);
+ GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn);
#else
GGML_LOG_INFO(" GGML_SYCL_ENABLE_DNN: DNN disabled by compile flag\n");
+ GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn);
#endif
-
+ GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN_MAX_KV: %d\n", g_ggml_sycl_fa_onednn_max_kv);
#ifdef SYCL_FLASH_ATTN
GGML_LOG_INFO(" GGML_SYCL_ENABLE_FLASH_ATTN: %d\n", g_ggml_sycl_enable_flash_attention);
#else
GGML_LOG_INFO(" GGML_SYCL_ENABLE_VMM: virtual memory extension is not available\n");
#endif
+ GGML_LOG_INFO(" GGML_SYCL_ENABLE_FUSION: %d\n", g_ggml_sycl_enable_fusion);
+
GGML_LOG_INFO(" GGML_SYCL_PRIORITIZE_DMMV: %d\n", g_ggml_sycl_prioritize_dmmv);
g_ggml_sycl_use_async_mem_op_requested = ggml_sycl_get_env("GGML_SYCL_USE_ASYNC_MEM_OP", 1);
switch (tensor->type) {
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q8_0:
+ case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
}
static bool check_usm_system(int device, size_t size) {
- bool use_usm_system = g_ggml_sycl_usm_system && size >= MEM_SIZE_1G;
+ bool use_usm_system = g_ggml_sycl_usm_system && size >= ((size_t)4 * MEM_SIZE_1G);
if (use_usm_system && !ggml_sycl_info().devices[device].usm_system_support) {
GGML_LOG_INFO("Device does not support USM system allocations\n");
void * dev_ptr;
if (use_usm_system) {
+ GGML_SYCL_DEBUG("[SYCL] allocating %lu Bytes with USM system\n", size);
dev_ptr = (void *)aligned_malloc_host(alignment, aligned_size);
if (!dev_ptr) {
GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on host\n", __func__, size);
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q8_0:
return true;
+ case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q8_0:
+ case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
return true;
}
+static bool reorder_qw_q2_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
+ GGML_ASSERT(size % sizeof(block_q2_K) == 0);
+ GGML_ASSERT(offset % sizeof(block_q2_K) == 0);
+
+ const int nblocks = size / sizeof(block_q2_K);
+
+ sycl_reorder_temp_buffer tmp(stream, size);
+ if (!tmp) {
+ GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, size);
+ return false;
+ }
+ uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
+
+ sycl::event copy_event;
+ SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, size)));
+ if (!g_ggml_sycl_use_async_mem_op) {
+ copy_event.wait();
+ }
+
+ auto * qs_ptr = data_device;
+ auto * scales_ptr = qs_ptr + (QK_K / 4) * nblocks;
+ sycl::half2 * dm_ptr = (sycl::half2 *) (scales_ptr + (QK_K / 16) * nblocks);
+
+ auto reorder_event = stream->parallel_for(nblocks, [=](auto i) {
+ const block_q2_K * x = (const block_q2_K *) tmp_buf;
+ const int ib = i;
+
+ for (int j = 0; j < QK_K / 4; ++j) {
+ qs_ptr[ib * (QK_K / 4) + j] = x[ib].qs[j];
+ }
+
+ for (int j = 0; j < QK_K / 16; ++j) {
+ scales_ptr[ib * (QK_K / 16) + j] = x[ib].scales[j];
+ }
+
+ dm_ptr[ib] = x[ib].dm;
+ });
+ if (!g_ggml_sycl_use_async_mem_op) {
+ reorder_event.wait_and_throw();
+ }
+ return true;
+}
+
static bool reorder_qw_q3_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
GGML_ASSERT(size % sizeof(block_q3_K) == 0);
GGML_ASSERT(offset % sizeof(block_q3_K) == 0);
return reorder_qw_q4_0(data_device, ncols, nrows, size, 0, stream);
case GGML_TYPE_Q8_0:
return reorder_qw_q8_0(data_device, ncols, nrows, size, 0, stream);
+ case GGML_TYPE_Q2_K:
+ return reorder_qw_q2_k(data_device, size, 0, stream);
case GGML_TYPE_Q3_K:
return reorder_qw_q3_k(data_device, size, 0, stream);
case GGML_TYPE_Q4_K:
case GGML_UNARY_OP_ELU:
ggml_sycl_elu(ctx, dst);
break;
+ case GGML_UNARY_OP_XIELU:
+ ggml_sycl_xielu(ctx, dst);
+ break;
case GGML_UNARY_OP_FLOOR:
ggml_sycl_floor(ctx, dst);
break;
if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
continue;
}
+
+ const int nodes_to_skip = ggml_sycl_fuse(*sycl_ctx, cgraph, i);
+ if (nodes_to_skip != 0) {
+ i += nodes_to_skip;
+ continue;
+ }
#ifndef NDEBUG
assert(node->buffer->buft == ggml_backend_sycl_buffer_type(sycl_ctx->device));
for (int j = 0; j < GGML_MAX_SRC; j++) {
case GGML_UNARY_OP_EXPM1:
case GGML_UNARY_OP_SOFTPLUS:
case GGML_UNARY_OP_ELU:
+ case GGML_UNARY_OP_XIELU:
case GGML_UNARY_OP_CEIL:
return true;
case GGML_UNARY_OP_FLOOR:
--- /dev/null
+#include <cfloat>
+#include <initializer_list>
+#include <vector>
+
+#include "ggml.h"
+#include "ggml-impl.h"
+#include "ggml-backend-impl.h"
+#include "topk-moe.hpp"
+
+// SYCL port of ggml-cuda/topk-moe.cu. The kernel is a translation of the CUDA no-bias, no-PDL
+// path of topk_moe_cuda; the fusion-detection helpers below are ported near-verbatim from
+// ggml-cuda.cu (pure graph / pointer inspection, backend-agnostic). Bias is not implemented here:
+// if a routing bias is detected, the fusion is declined and the eager path runs unchanged.
+
+struct ggml_sycl_topk_moe_args {
+ bool sigmoid{};
+ bool softmax{};
+ bool delayed_softmax{};
+ bool prob_bias{};
+ bool norm{};
+ bool scale{};
+};
+
+struct topk_moe_config {
+ bool use_sigmoid;
+ bool with_norm;
+ bool delayed_softmax;
+};
+
+// warp-local softmax used for both the pre-top-k logits and the post-top-k delayed path
+template <int experts_per_thread, bool use_limit>
+static inline void softmax_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
+ float max_val = -INFINITY;
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ const int idx = lane + i * WARP_SIZE;
+ const bool active = !use_limit || (idx < limit);
+ if (active) {
+ max_val = sycl::fmax(max_val, vals[i]);
+ }
+ }
+ max_val = warp_reduce_max<WARP_SIZE>(max_val);
+
+ float sum = 0.f;
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ const int idx = lane + i * WARP_SIZE;
+ const bool active = !use_limit || (idx < limit);
+ if (active) {
+ const float val = sycl::exp(vals[i] - max_val);
+ vals[i] = val;
+ sum += val;
+ } else {
+ vals[i] = 0.f;
+ }
+ }
+ sum = warp_reduce_sum<WARP_SIZE>(sum);
+
+ const float inv_sum = 1.0f / sum;
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ const int idx = lane + i * WARP_SIZE;
+ if (!use_limit || idx < limit) {
+ vals[i] *= inv_sum;
+ }
+ }
+}
+
+template <int experts_per_thread, bool use_limit>
+static inline void sigmoid_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ const int idx = lane + i * WARP_SIZE;
+ const bool active = !use_limit || (idx < limit);
+ vals[i] = active ? 1.f / (1.f + sycl::exp(-vals[i])) : -INFINITY;
+ }
+}
+
+/*
+ This kernel does the following:
+ 1. optionally softmax/sigmoid over the logits per token [n_experts, n_tokens]
+ 2. argmax reduce over the top-k (n_experts_used) logits
+ 3. write weights + ids to global memory
+ 4. optionally normalize the weights or apply softmax over the selected logits
+
+ It is intended as a fusion of the softmax->top-k->get_rows pipeline for MoE models.
+ One sub-group handles one row/token, mirroring topk_moe_cuda's one-warp-per-row layout.
+*/
+template <int n_experts>
+static void topk_moe_kernel(const float * __restrict__ logits,
+ float * __restrict__ weights,
+ int32_t * __restrict__ ids,
+ const int n_rows,
+ const int n_expert_used,
+ const float clamp_val,
+ const float scale_val,
+ const topk_moe_config config) {
+ auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<1>();
+ const int row = item_ct1.get_group(0);
+ if (row >= n_rows) {
+ return;
+ }
+ const int lane = item_ct1.get_local_id(0);
+
+ logits += (size_t) n_experts * row;
+ weights += (size_t) n_expert_used * row;
+ ids += (size_t) n_experts * row; // ids row stride is n_experts (matches the argsort tensor)
+
+ constexpr int experts_per_thread = (n_experts > WARP_SIZE) ? n_experts / WARP_SIZE : 1;
+
+ float wt[experts_per_thread];
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ wt[i] = -INFINITY;
+ }
+#pragma unroll
+ for (int i = 0; i < n_experts; i += WARP_SIZE) {
+ const int expert = i + lane;
+ wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY;
+ }
+
+ if (!config.delayed_softmax) {
+ if (config.use_sigmoid) {
+ sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, lane);
+ } else {
+ softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, lane);
+ }
+ }
+
+ // Sanitize NaN to -FLT_MAX so the iterative argmax produces unique expert IDs. NaN comparisons
+ // always return false, which would cause the same expert to be selected repeatedly.
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ if (sycl::isnan(wt[i])) {
+ wt[i] = -FLT_MAX;
+ }
+ }
+
+ // each thread now holds either a portion of the softmax distribution or the raw logits. Do the
+ // argmax reduce over n_expert_used, each time marking the selected expert as -inf to exclude it
+ // from the next iteration.
+
+ float wt_sum = 0.f;
+ float output_weights[experts_per_thread];
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ output_weights[i] = 0.f;
+ }
+
+ const sycl::sub_group sg = item_ct1.get_sub_group();
+
+ for (int k = 0; k < n_expert_used; k++) {
+ float max_val = wt[0];
+ int max_expert = lane;
+#pragma unroll
+ for (int i = 1; i < experts_per_thread; i++) {
+ const int expert = lane + i * WARP_SIZE;
+ if ((n_experts % WARP_SIZE == 0 || expert < n_experts) && wt[i] > max_val) {
+ max_val = wt[i];
+ max_expert = expert;
+ }
+ }
+#pragma unroll
+ for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
+ const float val = dpct::permute_sub_group_by_xor(sg, max_val, mask);
+ const int expert = dpct::permute_sub_group_by_xor(sg, max_expert, mask);
+ if (val > max_val || (val == max_val && expert < max_expert)) {
+ max_val = val;
+ max_expert = expert;
+ }
+ }
+
+ if ((max_expert & (WARP_SIZE - 1)) == lane) {
+ wt[max_expert / WARP_SIZE] = -INFINITY;
+ }
+ if ((k & (WARP_SIZE - 1)) == lane) {
+ output_weights[k / WARP_SIZE] = max_val;
+ }
+ if ((max_expert & (WARP_SIZE - 1)) == lane) {
+ ids[k] = max_expert;
+ if (config.with_norm) {
+ wt_sum += max_val;
+ }
+ }
+ }
+
+ if (config.with_norm) {
+ wt_sum = warp_reduce_sum<WARP_SIZE>(wt_sum);
+ wt_sum = sycl::fmax(wt_sum, clamp_val);
+ const float inv = 1.0f / wt_sum;
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ output_weights[i] *= inv;
+ }
+ }
+
+ if (config.delayed_softmax) {
+ softmax_warp_inplace<experts_per_thread, true>(output_weights, n_expert_used, lane);
+ }
+
+#pragma unroll
+ for (int i = 0; i < experts_per_thread; i++) {
+ const int idx = i * WARP_SIZE + lane;
+ if (idx < n_expert_used) {
+ weights[idx] = output_weights[i] * scale_val;
+ }
+ }
+}
+
+template <int n_experts>
+static void launch_topk_moe(queue_ptr stream, const float * logits, float * weights, int32_t * ids, int n_rows,
+ int n_expert_used, float clamp_val, float scale_val, const topk_moe_config & config) {
+ const sycl::range<1> block_dims(WARP_SIZE);
+ const sycl::range<1> block_nums(n_rows);
+ stream->parallel_for(sycl::nd_range<1>(block_nums * block_dims, block_dims),
+ [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+ topk_moe_kernel<n_experts>(logits, weights, ids, n_rows, n_expert_used, clamp_val,
+ scale_val, config);
+ GGML_UNUSED(item_ct1);
+ });
+}
+
+static void ggml_sycl_op_topk_moe(ggml_backend_sycl_context & ctx,
+ const ggml_tensor * logits,
+ ggml_tensor * weights,
+ ggml_tensor * ids,
+ const ggml_tensor * clamp,
+ const ggml_tensor * scale,
+ const ggml_sycl_topk_moe_args & args) {
+ GGML_ASSERT(logits->type == GGML_TYPE_F32);
+ GGML_ASSERT(weights->type == GGML_TYPE_F32);
+ GGML_ASSERT(ids->type == GGML_TYPE_I32);
+
+ const int n_experts = logits->ne[0];
+ const int n_rows = logits->ne[1];
+ const int n_expert_used = weights->ne[1];
+
+ GGML_ASSERT(ids->nb[1] / ggml_type_size(ids->type) == (size_t) n_experts);
+
+ const float * logits_d = (const float *) logits->data;
+ float * weights_d = (float *) weights->data;
+ int32_t * ids_d = (int32_t *) ids->data;
+
+ const bool with_norm = clamp != nullptr;
+ const float clamp_val = clamp ? ggml_get_op_params_f32(clamp, 0) : -INFINITY;
+ const float scale_val = scale ? ggml_get_op_params_f32(scale, 0) : 1.0f;
+
+ topk_moe_config config;
+ config.use_sigmoid = args.sigmoid;
+ config.with_norm = with_norm;
+ config.delayed_softmax = args.delayed_softmax;
+
+ queue_ptr stream = ctx.stream();
+ ggml_sycl_set_device(ctx.device);
+
+ switch (n_experts) {
+ case 1:
+ launch_topk_moe<1>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 2:
+ launch_topk_moe<2>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 4:
+ launch_topk_moe<4>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 8:
+ launch_topk_moe<8>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 16:
+ launch_topk_moe<16>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 32:
+ launch_topk_moe<32>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 64:
+ launch_topk_moe<64>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 128:
+ launch_topk_moe<128>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 256:
+ launch_topk_moe<256>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ case 512:
+ launch_topk_moe<512>(stream, logits_d, weights_d, ids_d, n_rows, n_expert_used, clamp_val, scale_val,
+ config);
+ break;
+ default:
+ GGML_ASSERT(false && "fatal error");
+ break;
+ }
+}
+
+static bool ggml_sycl_should_use_topk_moe(const ggml_tensor * gating_op, const ggml_tensor * weights,
+ const ggml_tensor * logits, const ggml_tensor * ids) {
+ const int n_expert = ids->nb[1] / ids->nb[0];
+ if ((n_expert & (n_expert - 1)) != 0 || n_expert > 512) {
+ return false;
+ }
+
+ if (!ggml_is_contiguous(weights) || !ggml_is_contiguous(logits)) {
+ return false;
+ }
+
+ if (gating_op->op == GGML_OP_SOFT_MAX) {
+ float scale = 1.0f;
+ float max_bias = 0.0f;
+
+ memcpy(&scale, (const float *) gating_op->op_params + 0, sizeof(float));
+ memcpy(&max_bias, (const float *) gating_op->op_params + 1, sizeof(float));
+
+ if (!ggml_is_contiguous(gating_op->src[0])) {
+ return false;
+ }
+ if (scale != 1.0f || max_bias != 0.0f) {
+ return false;
+ }
+ // don't fuse when masks or sinks are present
+ if (gating_op->src[1] || gating_op->src[2]) {
+ return false;
+ }
+ } else if (gating_op->op == GGML_OP_UNARY) {
+ if (ggml_get_unary_op(gating_op) != GGML_UNARY_OP_SIGMOID) {
+ return false;
+ }
+ }
+
+ return true;
+}
+
+// ported from ggml_cuda_topk_moe_fusion - pure graph inspection, backend-agnostic
+static bool ggml_sycl_topk_moe_fusion(const ggml_cgraph * cgraph, int node_idx, ggml_sycl_topk_moe_args & args) {
+ args = ggml_sycl_topk_moe_args{};
+
+ const int n_nodes = cgraph->n_nodes;
+ ggml_tensor ** nodes = cgraph->nodes;
+
+ if (nodes[node_idx]->op == GGML_OP_SOFT_MAX) {
+ args.softmax = true;
+ }
+
+ if (nodes[node_idx]->op == GGML_OP_UNARY) {
+ if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) {
+ return false;
+ }
+ args.sigmoid = true;
+ }
+
+ if (nodes[node_idx]->op == GGML_OP_ARGSORT) {
+ args.delayed_softmax = true;
+ }
+
+ node_idx++;
+
+ if (args.sigmoid || args.softmax) {
+ // SOFTMAX -> RESHAPE
+ if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
+ nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
+ return false;
+ }
+ ggml_tensor * probs_reshaped = nodes[node_idx];
+ node_idx++;
+
+ if (node_idx >= n_nodes) {
+ return false;
+ }
+
+ // src of bias add is the unreshaped probs (-2 instead of -1)
+ if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == nodes[node_idx - 2]) {
+ args.prob_bias = true;
+ node_idx++;
+ }
+ // RESHAPE/ADD -> ARGSORT
+ if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) {
+ return false;
+ }
+
+ if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
+ return false;
+ } else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) {
+ return false;
+ }
+
+ node_idx++;
+
+ // ARGSORT -> VIEW
+ if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW ||
+ nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
+ return false;
+ }
+ node_idx++;
+
+ if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_GET_ROWS) {
+ return false;
+ }
+
+ // GET_ROWS
+ if (nodes[node_idx]->src[0] != probs_reshaped || nodes[node_idx]->src[1] != nodes[node_idx - 1]) {
+ return false;
+ }
+ node_idx++;
+ } else if (args.delayed_softmax) {
+ if (node_idx - 2 < 0) {
+ return false;
+ }
+ ggml_tensor * probs_reshaped = nodes[node_idx - 2];
+
+ // VIEW -> ARGSORT
+ if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW ||
+ nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
+ return false;
+ }
+ node_idx++;
+
+ // GET_ROWS
+ if (node_idx >= n_nodes || nodes[node_idx]->src[1] != nodes[node_idx - 1] ||
+ nodes[node_idx]->src[0] != probs_reshaped) {
+ return false;
+ }
+ node_idx++;
+
+ static const std::vector<ggml_op> remaining_ops = { GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE };
+
+ for (const ggml_op op : remaining_ops) {
+ if (node_idx >= n_nodes || nodes[node_idx]->op != op || nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
+ return false;
+ }
+ node_idx++;
+ }
+ }
+
+ // at this point we can check for norm + scale; everything is now at least valid up to the norm
+ if (node_idx >= n_nodes) {
+ return true;
+ }
+
+ if (nodes[node_idx]->op == GGML_OP_RESHAPE) {
+ // check RESHAPE -> SUM_ROWS -> CLAMP -> DIV -> RESHAPE
+ static const std::vector<ggml_op> norm_ops = { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP };
+
+ args.norm = true;
+ for (const ggml_op op : norm_ops) {
+ if (nodes[node_idx]->op == op && nodes[node_idx]->src[0] == nodes[node_idx - 1]) {
+ node_idx++;
+ } else {
+ args.norm = false;
+ return true;
+ }
+ }
+
+ // DIV <- CLAMP, RESHAPE
+ if (nodes[node_idx]->op != GGML_OP_DIV || nodes[node_idx]->src[1] != nodes[node_idx - 1] ||
+ nodes[node_idx]->src[0] != nodes[node_idx - 3]) {
+ args.norm = false;
+ return true;
+ }
+ node_idx++;
+
+ if (nodes[node_idx]->op != GGML_OP_RESHAPE || nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
+ args.norm = false;
+ return true;
+ }
+ node_idx++;
+ }
+
+ if (nodes[node_idx]->op == GGML_OP_SCALE && nodes[node_idx]->src[0] == nodes[node_idx - 1]) {
+ args.scale = true;
+ }
+
+ return true;
+}
+
+// returns whether the write (out) nodes overwrite the read nodes in operation
+// ported from ggml_cuda_check_fusion_memory_ranges - pure pointer/range inspection
+static bool ggml_sycl_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx,
+ const int node_count, const int * out_nodes, const int out_count,
+ const bool is_topk_moe = false) {
+ auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) {
+ const int64_t a_start = (int64_t) a->data;
+ const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a);
+
+ const int64_t b_start = (int64_t) b->data;
+ const int64_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b);
+
+ if ((b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end)) {
+ return true;
+ }
+
+ return false;
+ };
+
+ bool is_ok = true;
+ // exception for topk-moe, as each row is read entirely before writing
+ if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) {
+ return true;
+ }
+
+ for (int i = 0; i < out_count; ++i) {
+ const ggml_tensor * dst = cgraph->nodes[out_nodes[i]];
+
+ for (int j = node_idx; j < node_idx + node_count; ++j) {
+ // loop over all srcs of all nodes in the fusion. If the src overlaps the destination and
+ // the src is not an intermediate node that's being elided, then disable fusion.
+ for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
+ const ggml_tensor * src = cgraph->nodes[j]->src[src_idx];
+
+ if (!src || src->op == GGML_OP_NONE) {
+ continue;
+ }
+
+ if (nodes_overlap(dst, src)) {
+ bool found = false;
+
+ for (int k = node_idx; k < j; ++k) {
+ if (cgraph->nodes[k] == src) {
+ found = true;
+ break;
+ }
+ }
+
+ if (!found) {
+ is_ok = false;
+ break;
+ }
+ }
+ }
+ }
+ }
+
+ return is_ok;
+}
+
+int ggml_sycl_fuse(ggml_backend_sycl_context & ctx, ggml_cgraph * cgraph, int i) {
+ if (!g_ggml_sycl_enable_fusion) {
+ return 0;
+ }
+
+ return ggml_sycl_fuse_topk_moe(ctx, cgraph, i);
+}
+
+int ggml_sycl_fuse_topk_moe(ggml_backend_sycl_context & ctx, ggml_cgraph * cgraph, int i) {
+ ggml_tensor * node = cgraph->nodes[i];
+
+ if (node->op != GGML_OP_UNARY && node->op != GGML_OP_SOFT_MAX && node->op != GGML_OP_ARGSORT) {
+ return 0;
+ }
+
+ ggml_sycl_topk_moe_args args;
+ if (!ggml_sycl_topk_moe_fusion(cgraph, i, args)) {
+ return 0;
+ }
+
+ // this kernel implements the no-bias path only; decline anything with a routing bias
+ if (args.prob_bias) {
+ return 0;
+ }
+
+ const ggml_tensor * logits = node->src[0];
+ ggml_tensor * weights = nullptr;
+ ggml_tensor * ids = nullptr;
+ const ggml_tensor * clamp = nullptr;
+ const ggml_tensor * scale = nullptr;
+
+ std::vector<ggml_op> ops;
+ int out_nodes[2];
+
+ if (!args.delayed_softmax) {
+ const ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX;
+ ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
+ out_nodes[0] = i + 3;
+ ids = cgraph->nodes[i + 3];
+
+ if (args.norm) {
+ ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE });
+ clamp = cgraph->nodes[i + (int) ops.size() - 3];
+ }
+ if (args.scale) {
+ ops.insert(ops.end(), { GGML_OP_SCALE });
+ scale = cgraph->nodes[i + (int) ops.size() - 1];
+ }
+
+ weights = cgraph->nodes[i + (int) ops.size() - 1];
+ out_nodes[1] = i + (int) ops.size() - 1;
+
+ if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) &&
+ ggml_sycl_should_use_topk_moe(node, weights, logits, ids) &&
+ ggml_sycl_check_fusion_memory_ranges(cgraph, i, (int) ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) {
+ ggml_sycl_op_topk_moe(ctx, logits, weights, ids, clamp, scale, args);
+ return (int) ops.size() - 1;
+ }
+ } else if (!args.norm && !args.prob_bias) {
+ // gpt-oss style: argsort -> view -> get_rows -> reshape -> softmax -> reshape, no norm/bias
+ ops.insert(ops.end(),
+ { GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, GGML_OP_SOFT_MAX,
+ GGML_OP_RESHAPE });
+ weights = cgraph->nodes[i + 5];
+ ids = cgraph->nodes[i + 1];
+ const ggml_tensor * softmax = cgraph->nodes[i + 4];
+ out_nodes[0] = i + 1;
+ out_nodes[1] = i + 5;
+
+ if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) &&
+ ggml_sycl_should_use_topk_moe(softmax, weights, logits, ids) &&
+ ggml_sycl_check_fusion_memory_ranges(cgraph, i, (int) ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) {
+ ggml_sycl_op_topk_moe(ctx, logits, weights, ids, clamp, scale, args);
+ return (int) ops.size() - 1;
+ }
+ }
+
+ return 0;
+}
--- /dev/null
+#ifndef GGML_SYCL_TOPK_MOE_HPP
+#define GGML_SYCL_TOPK_MOE_HPP
+
+#include "common.hpp"
+
+// Detect a fusable op subgraph starting at cgraph node `i` and, if found, dispatch the fused
+// kernel. Returns the number of *following* nodes consumed (0 = no fusion applies at i).
+int ggml_sycl_fuse(ggml_backend_sycl_context & ctx, ggml_cgraph * cgraph, int i);
+
+int ggml_sycl_fuse_topk_moe(ggml_backend_sycl_context & ctx, ggml_cgraph * cgraph, int i);
+
+#endif // GGML_SYCL_TOPK_MOE_HPP
"GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT"
)
+ test_shader_extension_support(
+ "GL_EXT_float_e2m1"
+ "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/feature-tests/float_e2m1.comp"
+ "GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT"
+ )
+
+ test_shader_extension_support(
+ "GL_EXT_float_e4m3"
+ "${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/feature-tests/float_e4m3.comp"
+ "GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT"
+ )
+
target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan)
target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
} VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE;
#endif
+#if !defined(VK_EXT_shader_ocp_microscaling_types)
+#define VK_EXT_shader_ocp_microscaling_types 1
+#define VK_EXT_SHADER_OCP_MICROSCALING_TYPES_SPEC_VERSION 1
+#define VK_EXT_SHADER_OCP_MICROSCALING_TYPES_EXTENSION_NAME "VK_EXT_shader_ocp_microscaling_types"
+#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_OCP_MICROSCALING_TYPES_FEATURES_EXT ((VkStructureType)1000672000)
+typedef struct VkPhysicalDeviceShaderOCPMicroscalingTypesFeaturesEXT {
+ VkStructureType sType;
+ void* pNext;
+ VkBool32 shaderFloat4;
+ VkBool32 shaderFloat6;
+ VkBool32 shaderFloat8UnsignedE8M0;
+ VkBool32 shaderMXInt8;
+} VkPhysicalDeviceShaderOCPMicroscalingTypesFeaturesEXT;
+#endif
+
+#if !defined(VK_EXT_shader_float8)
+#define VK_EXT_shader_float8 1
+#define VK_EXT_SHADER_FLOAT8_SPEC_VERSION 1
+#define VK_EXT_SHADER_FLOAT8_EXTENSION_NAME "VK_EXT_shader_float8"
+#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT ((VkStructureType)1000567000)
+typedef struct VkPhysicalDeviceShaderFloat8FeaturesEXT {
+ VkStructureType sType;
+ void* pNext;
+ VkBool32 shaderFloat8;
+ VkBool32 shaderFloat8CooperativeMatrix;
+} VkPhysicalDeviceShaderFloat8FeaturesEXT;
+#endif
+
+#ifndef VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME
+#define VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME "VK_KHR_internally_synchronized_queues"
+#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR ((VkStructureType)1000504000)
+#define VK_DEVICE_QUEUE_CREATE_INTERNALLY_SYNCHRONIZED_BIT_KHR ((VkDeviceQueueCreateFlagBits)0x00000004)
+
+// Compile-time constant guaranteed; no runtime initialization overhead
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR =
+ static_cast<vk::DeviceQueueCreateFlagBits>(0x00000004);
+
+typedef struct VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR {
+ VkStructureType sType;
+ void* pNext;
+ VkBool32 internallySynchronizedQueues;
+} VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR;
+#else
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR = vk::DeviceQueueCreateFlagBits::eInternallySynchronizedKHR;
+#endif
+
#define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
#define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
};
// Prevent simultaneous submissions to the same queue.
-// This could be per vk_queue if we stopped having two vk_queue structures
-// sharing the same vk::Queue.
-static std::mutex queue_mutex;
+struct vk_queue_handle {
+ vk::Queue queue;
+ virtual void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) = 0;
+ virtual void lock() {} // no-op by default (internally synchronized case)
+ virtual void unlock() {}
+ virtual ~vk_queue_handle() = default;
+};
+
+struct vk_queue_handle_synchronized : vk_queue_handle {
+ std::mutex mutex;
+ void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+ std::lock_guard<std::mutex> guard(mutex);
+ queue.submit(submits, fence);
+ }
+ void lock() override { mutex.lock(); }
+ void unlock() override { mutex.unlock(); }
+};
+
+struct vk_queue_handle_unsynchronized : vk_queue_handle {
+ void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+ // Driver guarantees internal synchronization via VK_KHR_internally_synchronized_queues
+ queue.submit(submits, fence);
+ }
+ // lock()/unlock() inherited no-ops
+};
struct vk_queue {
uint32_t queue_family_index;
- vk::Queue queue;
+ std::shared_ptr<vk_queue_handle> handle;
vk_command_pool cmd_pool;
vk::PipelineStageFlags stage_flags;
bool transfer_only;
-
- // copy everything except the cmd_pool
- void copyFrom(vk_queue &other) {
- queue_family_index = other.queue_family_index;
- queue = other.queue;
- stage_flags = other.stage_flags;
- transfer_only = other.transfer_only;
- }
};
static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft);
uint32_t vendor_id;
vk::DriverId driver_id;
vk_device_architecture architecture;
- vk_queue compute_queue;
- vk_queue transfer_queue;
+ std::unique_ptr<vk_queue> compute_queue;
+ std::unique_ptr<vk_queue> transfer_queue;
bool single_queue;
bool support_async;
bool async_use_transfer_queue;
+ bool has_internally_synchronized_queues = false;
uint32_t subgroup_size;
uint32_t subgroup_size_log2;
uint32_t shader_core_count;
bool uma;
bool prefer_host_memory;
bool float_controls_rte_fp16;
+ bool float_controls_denorm_preserve_fp16;
bool subgroup_basic;
bool subgroup_arithmetic;
bool subgroup_shuffle;
bool coopmat2_decode_vector;
bool dot2_f16 {};
+ bool ocp_fp4 {};
bool pipeline_executable_properties_support {};
vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT];
vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT];
vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32;
- vk_pipeline pipeline_set_rows_i32[GGML_TYPE_COUNT];
- vk_pipeline pipeline_set_rows_i64[GGML_TYPE_COUNT];
+ // [src0 0=fp32,1=fp16][dst]
+ vk_pipeline pipeline_set_rows_i32[2][GGML_TYPE_COUNT];
+ vk_pipeline pipeline_set_rows_i64[2][GGML_TYPE_COUNT];
vk_pipeline pipeline_norm_f32;
vk_pipeline pipeline_group_norm_f32;
vk_pipeline pipeline_rms_norm_f32;
vk_pipeline pipeline_col2im_1d_f32;
vk_pipeline pipeline_col2im_1d_f16;
vk_pipeline pipeline_col2im_1d_bf16;
+ vk_pipeline pipeline_out_prod_f32;
vk_pipeline pipeline_snake_f32;
vk_pipeline pipeline_snake_f16;
vk_pipeline pipeline_snake_bf16;
ggml_vk_destroy_buffer(sync_staging);
- compute_queue.cmd_pool.destroy(device);
- transfer_queue.cmd_pool.destroy(device);
+ if (compute_queue) compute_queue->cmd_pool.destroy(device);
+ if (transfer_queue) transfer_queue->cmd_pool.destroy(device);
+
+ // Explicitly clear to ensure queues drop their shared_ptrs to handles
+ // before the Vulkan logical device instance is destroyed
+ compute_queue.reset();
+ transfer_queue.reset();
for (auto& pipeline : all_pipelines) {
if (pipeline.expired()) {
vk::ShaderModuleCreateInfo shader_module_create_info({}, spv_size, reinterpret_cast<const uint32_t *>(spv_data));
- // Patch SPIR-V to enable RTE rounding for FP16, avoiding the need for
- // separate shader variants compiled with -DRTE16.
+ // Patch SPIR-V to enable supported FP16 float controls, avoiding the need
+ // for separate shader variants.
std::vector<uint32_t> spirv;
- if (device->float_controls_rte_fp16) {
+ if (device->float_controls_rte_fp16 || device->float_controls_denorm_preserve_fp16) {
const uint32_t* spv_words = reinterpret_cast<const uint32_t *>(spv_data);
size_t word_count = spv_size / sizeof(uint32_t);
spirv.assign(spv_words, spv_words + word_count);
// Insert from latest position first so earlier indices stay valid.
- // OpExecutionMode %entrypoint RoundingModeRTE 16
- uint32_t exec_mode[] = { (4u << spv::WordCountShift) | spv::OpExecutionMode, entry_point_id, spv::ExecutionModeRoundingModeRTE, 16 };
- spirv.insert(spirv.begin() + exec_insert_pos, std::begin(exec_mode), std::end(exec_mode));
+ if (device->float_controls_rte_fp16) {
+ // OpExecutionMode %entrypoint RoundingModeRTE 16
+ uint32_t exec_mode[] = { (4u << spv::WordCountShift) | spv::OpExecutionMode, entry_point_id, spv::ExecutionModeRoundingModeRTE, 16 };
+ spirv.insert(spirv.begin() + exec_insert_pos, std::begin(exec_mode), std::end(exec_mode));
+ }
+
+ if (device->float_controls_denorm_preserve_fp16) {
+ // OpExecutionMode %entrypoint DenormPreserve 16
+ uint32_t exec_mode[] = { (4u << spv::WordCountShift) | spv::OpExecutionMode, entry_point_id, spv::ExecutionModeDenormPreserve, 16 };
+ spirv.insert(spirv.begin() + exec_insert_pos, std::begin(exec_mode), std::end(exec_mode));
+ }
// OpExtension "SPV_KHR_float_controls"
const char ext_str[] = "SPV_KHR_float_controls";
memcpy(&extension[1], ext_str, sizeof(ext_str));
spirv.insert(spirv.begin() + ext_insert_pos, extension.begin(), extension.end());
- // OpCapability RoundingModeRTE
- uint32_t capability[] = { (2u << spv::WordCountShift) | spv::OpCapability, spv::CapabilityRoundingModeRTE };
- spirv.insert(spirv.begin() + cap_insert_pos, std::begin(capability), std::end(capability));
+ if (device->float_controls_rte_fp16) {
+ // OpCapability RoundingModeRTE
+ uint32_t capability[] = { (2u << spv::WordCountShift) | spv::OpCapability, spv::CapabilityRoundingModeRTE };
+ spirv.insert(spirv.begin() + cap_insert_pos, std::begin(capability), std::end(capability));
+ }
+
+ if (device->float_controls_denorm_preserve_fp16) {
+ // OpCapability DenormPreserve
+ uint32_t capability[] = { (2u << spv::WordCountShift) | spv::OpCapability, spv::CapabilityDenormPreserve };
+ spirv.insert(spirv.begin() + cap_insert_pos, std::begin(capability), std::end(capability));
+ }
shader_module_create_info = vk::ShaderModuleCreateInfo({}, spirv.size() * sizeof(uint32_t), spirv.data());
}
static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
if (ctx->seqs.empty()) {
if (fence) {
- std::lock_guard<std::mutex> guard(queue_mutex);
- ctx->p->q->queue.submit({}, fence);
+ ctx->p->q->handle->submit({}, fence);
}
return;
}
}
}
- std::lock_guard<std::mutex> guard(queue_mutex);
- ctx->p->q->queue.submit(submit_infos, fence);
+ ctx->p->q->handle->submit(submit_infos, fence);
ctx->seqs.clear();
}
abort();
}
-static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
+static std::unique_ptr<vk_queue> ggml_vk_create_queue(vk_device& device, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
VK_LOG_DEBUG("ggml_vk_create_queue()");
std::lock_guard<std::recursive_mutex> guard(device->mutex);
- q.queue_family_index = queue_family_index;
- q.transfer_only = transfer_only;
+ auto q = std::make_unique<vk_queue>();
+ q->queue_family_index = queue_family_index;
+ q->transfer_only = transfer_only;
- q.cmd_pool.init(device, &q);
+ std::shared_ptr<vk_queue_handle> h;
+ vk::DeviceQueueInfo2 queue_info2{};
+ queue_info2.queueFamilyIndex = queue_family_index;
+ queue_info2.queueIndex = queue_index;
- q.queue = device->device.getQueue(queue_family_index, queue_index);
+ if (device->has_internally_synchronized_queues) {
+ h = std::make_shared<vk_queue_handle_unsynchronized>();
+ queue_info2.flags = eInternallySynchronizedKHR;
+ } else {
+ h = std::make_shared<vk_queue_handle_synchronized>();
+ }
+
+ h->queue = device->device.getQueue2(queue_info2);
+ q->handle = h;
+
+ q->cmd_pool.init(device, q.get());
+
+ q->stage_flags = stage_flags;
+ return q;
+}
- q.stage_flags = stage_flags;
+static std::unique_ptr<vk_queue> ggml_vk_create_aliased_queue(vk_device& device, const std::unique_ptr<vk_queue>& source) {
+ std::lock_guard<std::recursive_mutex> guard(device->mutex);
+ auto q = std::make_unique<vk_queue>();
+ q->handle = source->handle;
+ q->queue_family_index = source->queue_family_index;
+ q->stage_flags = source->stage_flags;
+ q->transfer_only = source->transfer_only;
+ q->cmd_pool.init(device, q.get());
+ return q;
}
static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_command_pool& p) {
// Arbitrary frequency to cleanup/reuse command buffers
static constexpr uint32_t cleanup_frequency = 10;
- if (device->compute_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
- ggml_vk_command_pool_cleanup(device, device->compute_queue.cmd_pool);
+ if (device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+ ggml_vk_command_pool_cleanup(device, device->compute_queue->cmd_pool);
}
- if (device->transfer_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
- ggml_vk_command_pool_cleanup(device, device->transfer_queue.cmd_pool);
+ if (device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+ ggml_vk_command_pool_cleanup(device, device->transfer_queue->cmd_pool);
}
}
};
static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type);
-static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16);
+static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16, ggml_type v_type = GGML_TYPE_F16);
static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) {
bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) ||
(!f32acc && device->coopmat_support_16x16x16_f16acc);
const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
- bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type);
+ bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type, v_type);
if (!shape_ok || !shmem_ok) {
path = FA_SCALAR;
path = FA_SCALAR;
}
- // Q1_0 K/V is only implemented on coopmat2 (flash_attn_cm2); there is no scalar FA shader for it.
- if ((k_type == GGML_TYPE_Q1_0 || v_type == GGML_TYPE_Q1_0) && device->coopmat2) {
- path = FA_COOPMAT2;
- }
-
switch (path) {
case FA_SCALAR:
return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
uint32_t block_a_size = 0;
switch (src0_type) {
+ case GGML_TYPE_Q2_0: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + dm
case GGML_TYPE_Q4_0: block_a_size = std430_size({{16, 4}, {fp_size, fp_align}}); break; // qs[16/4] + dm
case GGML_TYPE_Q4_1: block_a_size = std430_size({{16, 4}, {fp2_size, fp2_align}}); break; // qs[16/4] + dm(vec2)
case GGML_TYPE_Q5_0: block_a_size = std430_size({{16, 4}, {4, 4}, {fp_size, fp_align}}); break; // qs[16/4] + qh + dm
return 0; // If no matching configuration is found
}
-// Whether scalar flash attention will use the MMQ path for the given k_type.
-static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type) {
+// Whether scalar flash attention will use the MMQ path for the given K/V types.
+static bool ggml_vk_fa_type_needs_shmem(ggml_type type) {
+ switch (type) {
+ case GGML_TYPE_IQ4_NL:
+ return true;
+ default:
+ return false;
+ }
+}
+
+static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type, ggml_type v_type) {
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
return device->integer_dot_product && device->subgroup_clustered &&
+ !ggml_vk_fa_type_needs_shmem(v_type) &&
(k_type == GGML_TYPE_Q4_0 || k_type == GGML_TYPE_Q4_1 ||
k_type == GGML_TYPE_Q5_0 || k_type == GGML_TYPE_Q5_1 ||
k_type == GGML_TYPE_Q8_0);
#else
GGML_UNUSED(device);
GGML_UNUSED(k_type);
+ GGML_UNUSED(v_type);
return false;
#endif
}
const bool fa_ds = fa.first.subgroup_size == 0;
const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
- const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type);
+ const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type, fa.first.v_type);
const void * spv_data = nullptr;
size_t spv_size = 0;
const char *name = nullptr;
}
#endif
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q1_0], matmul_q1_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_0], matmul_q2_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_S], matmul_iq3_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
- CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
- CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_NVFP4], matmul_nvfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ if (device->ocp_fp4) {
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_NVFP4], matmul_nvfp4_f16_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+ } else
+#endif
+ {
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_NVFP4], matmul_nvfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+ }
GGML_ASSERT(device->subgroup_ballot);
}
#endif
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
- CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
- CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ if (device->ocp_fp4) {
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+ } else
+#endif
+ {
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+ CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+ }
#undef CREATE_MM
#undef CREATE_MM2
} else
}
#endif
- if (device->coopmat_acc_f16_support) {
- CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
-
- CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K], matmul_q6_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S], matmul_iq1_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M], matmul_iq1_m_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_S], matmul_iq2_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_S], matmul_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0], matmul_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+
+ CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K], matmul_q6_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S], matmul_iq1_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M], matmul_iq1_m_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_S], matmul_iq2_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_S], matmul_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ if (device->ocp_fp4) {
+ CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4], matmul_mxfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4], matmul_nvfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+ } else
+#endif
+ {
CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4], matmul_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4], matmul_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- } else {
- CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0].f32acc, matmul_q1_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
-
- CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S].f32acc, matmul_iq1_s_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M].f32acc, matmul_iq1_m_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS].f32acc, matmul_iq2_xxs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XS].f32acc, matmul_iq2_xs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_S].f32acc, matmul_iq2_s_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_XXS].f32acc, matmul_iq3_xxs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_S].f32acc, matmul_iq3_s_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS].f32acc, matmul_iq4_xs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f32acc, matmul_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4].f32acc, matmul_mxfp4_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
- CREATE_MM(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4].f32acc, matmul_nvfp4_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
}
GGML_ASSERT(device->subgroup_ballot);
#endif
CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+ CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
- CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
- CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ if (device->ocp_fp4) {
+ CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+ CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+ } else
+#endif
+ {
+ CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+ CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+ }
#undef CREATE_MM2
#undef CREATE_MM
} else
CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
+ CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0], matmul_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->integer_dot_product) {
+ CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_0], matmul_q2_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0], matmul_q4_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_1], matmul_q4_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_0], matmul_q5_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_subgroup_f16_f32, wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16);
CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16);
CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
+ CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->integer_dot_product) {
+ CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
+ CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_q2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_q5_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->integer_dot_product) {
+ CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_0], matmul_id_q2_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0], matmul_id_q4_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_1], matmul_id_q4_1_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_0], matmul_id_q5_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0].f32acc, matmul_q1_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
+ CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0].f32acc, matmul_q2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->integer_dot_product) {
+ CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_0].f32acc, matmul_q2_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16);
CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0].f32acc, matmul_id_subgroup_q1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
+ CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0].f32acc, matmul_id_subgroup_q2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_subgroup_q4_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_subgroup_q4_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_subgroup_q5_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0].f32acc, matmul_id_q1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
+ CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0].f32acc, matmul_id_q2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
static constexpr uint32_t mul_mat_vec_num_bindings = 5;
static constexpr uint32_t mul_mat_vec_id_num_bindings = 6;
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+#define OCP_DMMV_LEN(NAME, REDUC) (device->ocp_fp4 ? NAME ## _ocp_len[REDUC] : NAME ## _len[REDUC])
+#define OCP_DMMV_DATA(NAME, REDUC) (device->ocp_fp4 ? NAME ## _ocp_data[REDUC] : NAME ## _data[REDUC])
+#else
+#define OCP_DMMV_LEN(NAME, REDUC) NAME ## _len[REDUC]
+#define OCP_DMMV_DATA(NAME, REDUC) NAME ## _data[REDUC]
+#endif
+
for (uint32_t w = 0; w < DMMV_WG_SIZE_COUNT; ++w) {
const uint32_t wg_size_subgroup = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size : (subgroup_size * 4);
const uint32_t wg_size_subgroup16 = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size16 : (subgroup_size16 * 4);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f32_f32", arr_dmmv_f16_f32_f32_len[reduc], arr_dmmv_f16_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f32_f32", arr_dmmv_bf16_f32_f32_len[reduc], arr_dmmv_bf16_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f32_f32", arr_dmmv_q1_0_f32_f32_len[reduc], arr_dmmv_q1_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f32_f32", arr_dmmv_q2_0_f32_f32_len[reduc], arr_dmmv_q2_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32", arr_dmmv_q4_0_f32_f32_len[reduc], arr_dmmv_q4_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32", arr_dmmv_q4_1_f32_f32_len[reduc], arr_dmmv_q4_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32", arr_dmmv_q5_0_f32_f32_len[reduc], arr_dmmv_q5_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ3_S][i], "mul_mat_vec_iq3_s_f32_f32", arr_dmmv_iq3_s_f32_f32_len[reduc16], arr_dmmv_iq3_s_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ4_XS][i], "mul_mat_vec_iq4_xs_f32_f32", arr_dmmv_iq4_xs_f32_f32_len[reduc16], arr_dmmv_iq4_xs_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f32_f32", arr_dmmv_iq4_nl_f32_f32_len[reduc16], arr_dmmv_iq4_nl_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
- ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_f32_f32", arr_dmmv_mxfp4_f32_f32_len[reduc16], arr_dmmv_mxfp4_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
- ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_NVFP4][i], "mul_mat_vec_nvfp4_f32_f32", arr_dmmv_nvfp4_f32_f32_len[reduc16], arr_dmmv_nvfp4_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_f32_f32", OCP_DMMV_LEN(arr_dmmv_mxfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_mxfp4_f32_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_NVFP4][i], "mul_mat_vec_nvfp4_f32_f32", OCP_DMMV_LEN(arr_dmmv_nvfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_nvfp4_f32_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_F32 ][i], "mul_mat_vec_f32_f16_f32", arr_dmmv_f32_f16_f32_len[reduc], arr_dmmv_f32_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1, 1, 1}, {wg_size_subgroup, 1, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f16_f32", arr_dmmv_f16_f16_f32_len[reduc], arr_dmmv_f16_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f16_f32", arr_dmmv_bf16_f16_f32_len[reduc], arr_dmmv_bf16_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f16_f32", arr_dmmv_q1_0_f16_f32_len[reduc], arr_dmmv_q1_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f16_f32", arr_dmmv_q2_0_f16_f32_len[reduc], arr_dmmv_q2_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32", arr_dmmv_q4_0_f16_f32_len[reduc], arr_dmmv_q4_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32", arr_dmmv_q4_1_f16_f32_len[reduc], arr_dmmv_q4_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32", arr_dmmv_q5_0_f16_f32_len[reduc], arr_dmmv_q5_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ3_S][i], "mul_mat_vec_iq3_s_f16_f32", arr_dmmv_iq3_s_f16_f32_len[reduc16], arr_dmmv_iq3_s_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ4_XS][i], "mul_mat_vec_iq4_xs_f16_f32", arr_dmmv_iq4_xs_f16_f32_len[reduc16], arr_dmmv_iq4_xs_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f16_f32", arr_dmmv_iq4_nl_f16_f32_len[reduc16], arr_dmmv_iq4_nl_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
- ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_f16_f32", arr_dmmv_mxfp4_f16_f32_len[reduc16], arr_dmmv_mxfp4_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
- ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_NVFP4][i], "mul_mat_vec_nvfp4_f16_f32", arr_dmmv_nvfp4_f16_f32_len[reduc16], arr_dmmv_nvfp4_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_f16_f32", OCP_DMMV_LEN(arr_dmmv_mxfp4_f16_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_mxfp4_f16_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_NVFP4][i], "mul_mat_vec_nvfp4_f16_f32", OCP_DMMV_LEN(arr_dmmv_nvfp4_f16_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_nvfp4_f16_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->integer_dot_product) {
const uint32_t subgroup_size_int = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
const uint32_t wg_size_subgroup_int = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size_int : (subgroup_size_int * 4);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_q8_1_f32", arr_dmmv_q2_0_q8_1_f32_len[reduc], arr_dmmv_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_q8_1_f32", arr_dmmv_q4_0_q8_1_f32_len[reduc], arr_dmmv_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_q8_1_f32", arr_dmmv_q4_1_q8_1_f32_len[reduc], arr_dmmv_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_q8_1_f32", arr_dmmv_q5_0_q8_1_f32_len[reduc], arr_dmmv_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32", arr_dmmv_id_f16_f32_f32_len[reduc], arr_dmmv_id_f16_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {wg_size_subgroup, 2}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_BF16], "mul_mat_vec_id_bf16_f32", arr_dmmv_id_bf16_f32_f32_len[reduc], arr_dmmv_id_bf16_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {wg_size_subgroup, 2}, 1, false, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q1_0], "mul_mat_vec_id_q1_0_f32", arr_dmmv_id_q1_0_f32_f32_len[reduc], arr_dmmv_id_q1_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q2_0], "mul_mat_vec_id_q2_0_f32", arr_dmmv_id_q2_0_f32_f32_len[reduc], arr_dmmv_id_q2_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_f32", arr_dmmv_id_q4_0_f32_f32_len[reduc], arr_dmmv_id_q4_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_f32", arr_dmmv_id_q4_1_f32_f32_len[reduc], arr_dmmv_id_q4_1_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_f32", arr_dmmv_id_q5_0_f32_f32_len[reduc], arr_dmmv_id_q5_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ3_S], "mul_mat_vec_id_iq3_s_f32", arr_dmmv_id_iq3_s_f32_f32_len[reduc16], arr_dmmv_id_iq3_s_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ4_XS], "mul_mat_vec_id_iq4_xs_f32", arr_dmmv_id_iq4_xs_f32_f32_len[reduc16], arr_dmmv_id_iq4_xs_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ4_NL], "mul_mat_vec_id_iq4_nl_f32", arr_dmmv_id_iq4_nl_f32_f32_len[reduc16], arr_dmmv_id_iq4_nl_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
- ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_MXFP4], "mul_mat_vec_id_mxfp4_f32", arr_dmmv_id_mxfp4_f32_f32_len[reduc16], arr_dmmv_id_mxfp4_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
- ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_NVFP4], "mul_mat_vec_id_nvfp4_f32", arr_dmmv_id_nvfp4_f32_f32_len[reduc16], arr_dmmv_id_nvfp4_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_MXFP4], "mul_mat_vec_id_mxfp4_f32", OCP_DMMV_LEN(arr_dmmv_id_mxfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_id_mxfp4_f32_f32, reduc16), "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_NVFP4], "mul_mat_vec_id_nvfp4_f32", OCP_DMMV_LEN(arr_dmmv_id_nvfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_id_nvfp4_f32_f32, reduc16), "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16);
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->integer_dot_product) {
const uint32_t subgroup_size_int = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
const uint32_t wg_size_subgroup_int = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size_int : (subgroup_size_int * 4);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_0], "mul_mat_vec_id_q2_0_q8_1_f32", arr_dmmv_id_q2_0_q8_1_f32_len[reduc], arr_dmmv_id_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_q8_1_f32", arr_dmmv_id_q4_0_q8_1_f32_len[reduc], arr_dmmv_id_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_q8_1_f32", arr_dmmv_id_q4_1_q8_1_f32_len[reduc], arr_dmmv_id_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_q8_1_f32", arr_dmmv_id_q5_0_q8_1_f32_len[reduc], arr_dmmv_id_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
#endif // GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT
}
+#undef OCP_DMMV_DATA
+#undef OCP_DMMV_LEN
+
#if !defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
GGML_UNUSED(rm_stdq_int);
GGML_UNUSED(rm_kq_int);
// dequant shaders
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16", dequant_f32_len, dequant_f32_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q1_0], "dequant_q1_0", dequant_q1_0_len, dequant_q1_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 8, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_0], "dequant_q2_0", dequant_q2_0_len, dequant_q2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F16 ], "get_rows_f16", get_rows_f16_len, get_rows_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_BF16], "get_rows_bf16", get_rows_bf16_len, get_rows_bf16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q1_0], "get_rows_q1_0", get_rows_q1_0_len, get_rows_q1_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_0], "get_rows_q2_0", get_rows_q2_0_len, get_rows_q2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F16 ], "get_rows_f16_f32", get_rows_f16_f32_len, get_rows_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_BF16], "get_rows_bf16_f32", get_rows_bf16_f32_len, get_rows_bf16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q1_0], "get_rows_q1_0_f32", get_rows_q1_0_f32_len, get_rows_q1_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_0], "get_rows_q2_0_f32", get_rows_q2_0_f32_len, get_rows_q2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_16, "cpy_transpose_16", cpy_transpose_16_len, cpy_transpose_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q1_0], "cpy_f32_q1_0", cpy_f32_q1_0_len, cpy_f32_q1_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q2_0], "cpy_f32_q2_0", cpy_f32_q2_0_len, cpy_f32_q2_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q4_0], "cpy_f32_q4_0", cpy_f32_q4_0_len, cpy_f32_q4_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q4_1], "cpy_f32_q4_1", cpy_f32_q4_1_len, cpy_f32_q4_1_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q5_0], "cpy_f32_q5_0", cpy_f32_q5_0_len, cpy_f32_q5_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q8_0], "cpy_f32_q8_0", cpy_f32_q8_0_len, cpy_f32_q8_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_IQ4_NL], "cpy_f32_iq4_nl", cpy_f32_iq4_nl_len, cpy_f32_iq4_nl_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
-#define SET_ROWS(itype) \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_F32], "set_rows_f32" #itype, set_rows_f32 ## itype ## _len, set_rows_f32 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_F16], "set_rows_f16" #itype, set_rows_f16 ## itype ## _len, set_rows_f16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_BF16], "set_rows_bf16" #itype, set_rows_bf16 ## itype ## _len, set_rows_bf16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q1_0], "set_rows_q1_0" #itype, set_rows_q1_0 ## itype ## _len, set_rows_q1_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q4_0], "set_rows_q4_0" #itype, set_rows_q4_0 ## itype ## _len, set_rows_q4_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q4_1], "set_rows_q4_1" #itype, set_rows_q4_1 ## itype ## _len, set_rows_q4_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q5_0], "set_rows_q5_0" #itype, set_rows_q5_0 ## itype ## _len, set_rows_q5_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q5_1], "set_rows_q5_1" #itype, set_rows_q5_1 ## itype ## _len, set_rows_q5_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q8_0], "set_rows_q8_0" #itype, set_rows_q8_0 ## itype ## _len, set_rows_q8_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
- ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_IQ4_NL], "set_rows_iq4_nl" #itype, set_rows_iq4_nl ## itype ## _len, set_rows_iq4_nl ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true);
-
- SET_ROWS(_i32)
- SET_ROWS(_i64)
+#define SET_ROWS(src_idx, src, itype) \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_F32], "set_rows_" #src "_f32" #itype, set_rows_ ## src ## _f32 ## itype ## _len, set_rows_ ## src ## _f32 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_F16], "set_rows_" #src "_f16" #itype, set_rows_ ## src ## _f16 ## itype ## _len, set_rows_ ## src ## _f16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_BF16], "set_rows_" #src "_bf16" #itype, set_rows_ ## src ## _bf16 ## itype ## _len, set_rows_ ## src ## _bf16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q1_0], "set_rows_" #src "_q1_0" #itype, set_rows_ ## src ## _q1_0 ## itype ## _len, set_rows_ ## src ## _q1_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q2_0], "set_rows_" #src "_q2_0" #itype, set_rows_ ## src ## _q2_0 ## itype ## _len, set_rows_ ## src ## _q2_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_0], "set_rows_" #src "_q4_0" #itype, set_rows_ ## src ## _q4_0 ## itype ## _len, set_rows_ ## src ## _q4_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_1], "set_rows_" #src "_q4_1" #itype, set_rows_ ## src ## _q4_1 ## itype ## _len, set_rows_ ## src ## _q4_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_0], "set_rows_" #src "_q5_0" #itype, set_rows_ ## src ## _q5_0 ## itype ## _len, set_rows_ ## src ## _q5_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_1], "set_rows_" #src "_q5_1" #itype, set_rows_ ## src ## _q5_1 ## itype ## _len, set_rows_ ## src ## _q5_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q8_0], "set_rows_" #src "_q8_0" #itype, set_rows_ ## src ## _q8_0 ## itype ## _len, set_rows_ ## src ## _q8_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_IQ4_NL], "set_rows_" #src "_iq4_nl" #itype, set_rows_ ## src ## _iq4_nl ## itype ## _len, set_rows_ ## src ## _iq4_nl ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true);
+
+ SET_ROWS(0, f32, _i32)
+ SET_ROWS(0, f32, _i64)
+ SET_ROWS(1, f16, _i32)
+ SET_ROWS(1, f16, _i64)
#undef SET_ROWS
ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q1_0], "cpy_q1_0_f32", cpy_q1_0_f32_len, cpy_q1_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q1_0), 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q2_0], "cpy_q2_0_f32", cpy_q2_0_f32_len, cpy_q2_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q2_0), 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q4_0], "cpy_q4_0_f32", cpy_q4_0_f32_len, cpy_q4_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q4_0), 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q4_1], "cpy_q4_1_f32", cpy_q4_1_f32_len, cpy_q4_1_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q4_1), 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q5_0], "cpy_q5_0_f32", cpy_q5_0_f32_len, cpy_q5_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q5_0), 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_col2im_1d_f16, "col2im_1d_f16", col2im_1d_f16_len, col2im_1d_f16_data, "main", 2, sizeof(vk_op_col2im_1d_push_constants), {256, 1, 1}, {}, 1, true);
ggml_vk_create_pipeline(device, device->pipeline_col2im_1d_bf16, "col2im_1d_bf16", col2im_1d_bf16_len, col2im_1d_bf16_data, "main", 2, sizeof(vk_op_col2im_1d_push_constants), {256, 1, 1}, {}, 1, true);
+ ggml_vk_create_pipeline(device, device->pipeline_out_prod_f32, "out_prod_f32", out_prod_f32_len, out_prod_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {256, 1, 1}, {}, 1);
+
ggml_vk_create_pipeline(device, device->pipeline_snake_f32, "snake_f32", snake_f32_len, snake_f32_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_snake_f16, "snake_f16", snake_f16_len, snake_f16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_snake_bf16, "snake_bf16", snake_bf16_len, snake_bf16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
bool coopmat2_support = false;
bool coopmat2_decode_vector_support = false;
bool pipeline_executable_properties_support = false;
+ bool internally_sync_support = false;
device->coopmat_support = false;
device->integer_dot_product = false;
device->shader_64b_indexing = false;
bool bfloat16_support = false;
bool dot2_f16_support = false;
+ bool ocp_microscaling_extension = false;
+ bool shader_float8_extension = false;
for (const auto& properties : ext_props) {
if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) {
} else if (strcmp("VK_KHR_shader_bfloat16", properties.extensionName) == 0 &&
!getenv("GGML_VK_DISABLE_BFLOAT16")) {
bfloat16_support = true;
+#endif
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
+ } else if (strcmp(VK_EXT_SHADER_OCP_MICROSCALING_TYPES_EXTENSION_NAME, properties.extensionName) == 0) {
+ ocp_microscaling_extension = true;
+#endif
+#if defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ } else if (strcmp(VK_EXT_SHADER_FLOAT8_EXTENSION_NAME, properties.extensionName) == 0) {
+ shader_float8_extension = true;
#endif
} else if (strcmp("VK_VALVE_shader_mixed_float_dot_product", properties.extensionName) == 0 &&
!getenv("GGML_VK_DISABLE_DOT2")) {
} else if (strcmp("VK_EXT_shader_64bit_indexing", properties.extensionName) == 0) {
device->shader_64b_indexing = true;
#endif
+ } else if (strcmp(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME, properties.extensionName) == 0) {
+ internally_sync_support = true;
}
}
device->shader_core_count = 0;
}
device->float_controls_rte_fp16 = vk12_props.shaderRoundingModeRTEFloat16;
+ device->float_controls_denorm_preserve_fp16 = vk12_props.shaderDenormPreserveFloat16;
device->subgroup_basic = (vk11_props.subgroupSupportedStages & vk::ShaderStageFlagBits::eCompute) &&
(vk11_props.subgroupSupportedOperations & vk::SubgroupFeatureFlagBits::eBasic);
device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1;
std::vector<vk::DeviceQueueCreateInfo> device_queue_create_infos;
- if (compute_queue_family_index != transfer_queue_family_index) {
- device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
- device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1});
- } else if(!device->single_queue) {
- device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities});
- } else {
- device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
- }
vk::DeviceCreateInfo device_create_info{};
std::vector<const char *> device_extensions;
vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures();
last_struct = (VkBaseOutStructure *)&vk12_features;
+ VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR internally_synchronized_queues_features{};
+ internally_synchronized_queues_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR;
+ internally_synchronized_queues_features.pNext = nullptr;
+ internally_synchronized_queues_features.internallySynchronizedQueues = VK_FALSE;
+
+ if (internally_sync_support) {
+ last_struct->pNext = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+ last_struct = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+ device_extensions.push_back(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME);
+ }
+
VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features;
pl_robustness_features.pNext = nullptr;
pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT;
}
#endif
+ VkPhysicalDeviceShaderOCPMicroscalingTypesFeaturesEXT ocp_microscaling_features {};
+ ocp_microscaling_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_OCP_MICROSCALING_TYPES_FEATURES_EXT;
+ if (ocp_microscaling_extension) {
+ last_struct->pNext = (VkBaseOutStructure *)&ocp_microscaling_features;
+ last_struct = (VkBaseOutStructure *)&ocp_microscaling_features;
+ device_extensions.push_back(VK_EXT_SHADER_OCP_MICROSCALING_TYPES_EXTENSION_NAME);
+ }
+
+ VkPhysicalDeviceShaderFloat8FeaturesEXT shader_float8_features {};
+ shader_float8_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT;
+ if (shader_float8_extension) {
+ last_struct->pNext = (VkBaseOutStructure *)&shader_float8_features;
+ last_struct = (VkBaseOutStructure *)&shader_float8_features;
+ device_extensions.push_back(VK_EXT_SHADER_FLOAT8_EXTENSION_NAME);
+ }
+
VkPhysicalDeviceMaintenance4Features maint4_features {};
maint4_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MAINTENANCE_4_FEATURES;
if (maintenance4_support) {
vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2);
+ device->has_internally_synchronized_queues = internally_synchronized_queues_features.internallySynchronizedQueues;
+
+ // Build queue create infos only after querying whether internally synchronized queues are enabled.
+ // getQueue2() later uses the same flag, so creation/retrieval must stay consistent.
+ vk::DeviceQueueCreateFlags queue_flags = device->has_internally_synchronized_queues ?
+ eInternallySynchronizedKHR :
+ vk::DeviceQueueCreateFlags();
+
+ if (compute_queue_family_index != transfer_queue_family_index) {
+ device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+ device_queue_create_infos.push_back({queue_flags, transfer_queue_family_index, 1, priorities + 1});
+ } else if(!device->single_queue) {
+ device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 2, priorities});
+ } else {
+ device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+ }
+
device->pipeline_executable_properties_support = pipeline_executable_properties_support;
device->fp16 = device->fp16 && vk12_features.shaderFloat16;
#endif
device->dot2_f16 = dot2_f16_support && dot2_features.shaderMixedFloatDotProductFloat16AccFloat32;
+ device->ocp_fp4 = ocp_microscaling_extension && ocp_microscaling_features.shaderFloat4 &&
+ shader_float8_extension && shader_float8_features.shaderFloat8 &&
+ !getenv("GGML_VK_DISABLE_OCP_FP4");
device->pipeline_robustness = pl_robustness_features.pipelineRobustness;
device->device = device->physical_device.createDevice(device_create_info);
// Queues
- ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
+ device->compute_queue = ggml_vk_create_queue(device, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
// Shaders
// Disable matmul tile sizes early if performance low or not supported
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;
if (!device->single_queue) {
const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0;
- ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
+ device->transfer_queue = ggml_vk_create_queue(device, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
device->async_use_transfer_queue = prefers_transfer_queue || (getenv("GGML_VK_ASYNC_USE_TRANSFER_QUEUE") != nullptr);
} else {
- // TODO: Use pointer or reference to avoid copy
- device->transfer_queue.copyFrom(device->compute_queue);
- device->transfer_queue.cmd_pool.init(device, &device->transfer_queue);
+ device->transfer_queue = ggml_vk_create_aliased_queue(device, device->compute_queue);
device->async_use_transfer_queue = false;
}
bool integer_dot_product = false;
bool bfloat16_support = false;
bool dot2_f16_support = false;
+ bool ocp_microscaling_extension = false;
+ bool shader_float8_extension = false;
for (auto properties : ext_props) {
if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) {
} else if (strcmp("VK_KHR_shader_bfloat16", properties.extensionName) == 0 &&
!getenv("GGML_VK_DISABLE_BFLOAT16")) {
bfloat16_support = true;
+#endif
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
+ } else if (strcmp(VK_EXT_SHADER_OCP_MICROSCALING_TYPES_EXTENSION_NAME, properties.extensionName) == 0) {
+ ocp_microscaling_extension = true;
+#endif
+#if defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ } else if (strcmp(VK_EXT_SHADER_FLOAT8_EXTENSION_NAME, properties.extensionName) == 0) {
+ shader_float8_extension = true;
#endif
} else if (strcmp("VK_VALVE_shader_mixed_float_dot_product", properties.extensionName) == 0 &&
!getenv("GGML_VK_DISABLE_DOT2")) {
last_struct = (VkBaseOutStructure *)&dot2_features;
}
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ VkPhysicalDeviceShaderOCPMicroscalingTypesFeaturesEXT ocp_microscaling_features {};
+ ocp_microscaling_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_OCP_MICROSCALING_TYPES_FEATURES_EXT;
+ VkPhysicalDeviceShaderFloat8FeaturesEXT shader_float8_features {};
+ shader_float8_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT;
+ if (ocp_microscaling_extension) {
+ last_struct->pNext = (VkBaseOutStructure *)&ocp_microscaling_features;
+ last_struct = (VkBaseOutStructure *)&ocp_microscaling_features;
+ }
+ if (shader_float8_extension) {
+ last_struct->pNext = (VkBaseOutStructure *)&shader_float8_features;
+ last_struct = (VkBaseOutStructure *)&shader_float8_features;
+ }
+#endif
+
vkGetPhysicalDeviceFeatures2(physical_device, &device_features2);
fp16 = fp16 && vk12_features.shaderFloat16;
bool dot2_f16 = dot2_f16_support && dot2_features.shaderMixedFloatDotProductFloat16AccFloat32;
const char *fp16_str = fp16 ? (dot2_f16 ? "dot2" : "1") : "0";
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ const bool fp4 = ocp_microscaling_extension && ocp_microscaling_features.shaderFloat4 &&
+ shader_float8_extension && shader_float8_features.shaderFloat8 &&
+ !getenv("GGML_VK_DISABLE_OCP_FP4");
+#else
+ GGML_UNUSED(ocp_microscaling_extension);
+ GGML_UNUSED(shader_float8_extension);
+ const bool fp4 = false;
+#endif
std::string device_name = props2.properties.deviceName.data();
- GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %s | bf16: %d | warp size: %zu | shared memory: %d | int dot: %d | matrix cores: %s\n",
- idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16_str, bf16, subgroup_size,
+ GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %s | bf16: %d | fp4: %d | warp size: %zu | shared memory: %d | int dot: %d | matrix cores: %s\n",
+ idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16_str, bf16, fp4, subgroup_size,
props2.properties.limits.maxComputeSharedMemorySize, integer_dot_product, matrix_cores.c_str());
if (props2.properties.deviceType == vk::PhysicalDeviceType::eCpu) {
ctx->fence = ctx->device->device.createFence({});
ctx->almost_ready_fence = ctx->device->device.createFence({});
- ctx->compute_cmd_pool.init(ctx->device, &ctx->device->compute_queue);
+ ctx->compute_cmd_pool.init(ctx->device, ctx->device->compute_queue.get());
if (ctx->device->async_use_transfer_queue) {
vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 };
vk::SemaphoreCreateInfo ci{};
ctx->transfer_semaphore.s = ctx->device->device.createSemaphore(ci);
ctx->transfer_semaphore.value = 0;
- ctx->transfer_cmd_pool.init(ctx->device, &ctx->device->transfer_queue);
+ ctx->transfer_cmd_pool.init(ctx->device, ctx->device->transfer_queue.get());
}
if (vk_perf_logger_enabled) {
switch (type) {
case GGML_TYPE_F32:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
switch (src0_type) {
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
if (b_type == GGML_TYPE_Q8_1) {
switch (a_type) {
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
switch (src0_type) {
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
if (b_type == GGML_TYPE_Q8_1) {
switch (a_type) {
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
return result;
}
+static vk_context ggml_vk_get_transfer_ctx(ggml_backend_vk_context * ctx) {
+ vk_context result;
+ if (!ctx->transfer_ctx.expired()) {
+ result = ctx->transfer_ctx.lock();
+ } else {
+ result = ggml_vk_create_context(ctx, ctx->transfer_cmd_pool);
+
+ ctx->transfer_ctx = result;
+ ggml_vk_ctx_begin(ctx->device, result);
+ }
+
+ return result;
+}
+
// Submit any pending transfer queue work and signal the transfer semaphore.
// The next compute context created via ggml_vk_get_compute_ctx will wait on this semaphore.
// Returns true if work was submitted.
} else {
std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
- vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+ vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
ggml_vk_ctx_begin(dst->device, subctx);
bool ret = ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, dpitch, width, height, true);
GGML_ASSERT(ret);
GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent);
std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
- vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue.cmd_pool);
+ vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue->cmd_pool);
ggml_vk_ctx_begin(src->device, subctx);
subctx->s->buffer->buf.pipelineBarrier(
vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer,
} else {
std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
- vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+ vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
ggml_vk_ctx_begin(src->device, subctx);
bool ret = ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, spitch, dpitch, width, height, true);
GGML_ASSERT(ret);
std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")");
// Copy within the device
- vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+ vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
ggml_vk_ctx_begin(src->device, subctx);
ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size);
ggml_vk_ctx_end(subctx);
}
std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
- vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+ vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
ggml_vk_ctx_begin(dst->device, subctx);
subctx->s->buffer->buf.fillBuffer(dst->buffer, offset, size, c);
ggml_vk_ctx_end(subctx);
if (src->type == GGML_TYPE_F32) {
switch (to) {
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
if (to == GGML_TYPE_F32) {
switch (src->type) {
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
// Quantization overhead is not worth it for small k
switch (device->vendor_id) {
case VK_VENDOR_ID_NVIDIA:
- if (src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_IQ1_S || src0_type == GGML_TYPE_IQ1_M) {
+ if (src0_type == GGML_TYPE_Q2_0 || src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_IQ1_S || src0_type == GGML_TYPE_IQ1_M) {
return true;
}
}
case VK_VENDOR_ID_INTEL:
if (device->architecture == vk_device_architecture::INTEL_XE2) {
- if (src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_Q6_K) {
+ if (src0_type == GGML_TYPE_Q2_0 || src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_Q6_K) {
return true;
}
}
static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) {
GGML_UNUSED(f32acc);
- GGML_UNUSED(v_type);
// Needs to be kept up to date on shader changes
const uint32_t wg_size = params.workgroup_size;
const uint32_t Br = params.block_rows;
// BF16 uses the fp32 shader (FLOAT_TYPE=float)
const uint32_t float_type_size = (device->fp16 && k_type != GGML_TYPE_BF16) ? sizeof(ggml_fp16_t) : sizeof(float);
- const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type);
+ const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type, v_type);
// tmpsh is overestimated slightly
const uint32_t tmpsh = wg_size * sizeof(float);
const uint32_t tmpshv4 = wg_size * 4 * float_type_size;
const uint32_t masksh = Bc * (Br + 1) * float_type_size;
+ // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated.
+ const uint32_t iq_shmem = 16 * float_type_size;
uint32_t Qf, kvsh, kblocksh_size;
if (mmq) {
kblocksh_size = 0;
}
- const uint32_t total_size = tmpsh + tmpshv4 + masksh + Qf + kvsh + kblocksh_size;
+ const uint32_t total_size = tmpsh + tmpshv4 + masksh + iq_shmem + Qf + kvsh + kblocksh_size;
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
VK_LOG_DEBUG("ggml_vk_flash_attn_scalar_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", mmq=" << mmq << ", total_size=" << total_size << ", supported=" << supported);
return supported;
}
-static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type) {
+static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) {
+ GGML_UNUSED(v_type);
// Needs to be kept up to date on shader changes
const uint32_t Br = params.block_rows;
const uint32_t Bc = params.block_cols;
const uint32_t f16vec4 = 8;
const uint32_t tmpsh = (Bc / MatBc) * sizeof(float);
+ // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated.
+ const uint32_t iq_shmem = 16 * sizeof(ggml_fp16_t);
const uint32_t qstride = hsk_pad / 4 + 2;
const uint32_t Qf = Br * qstride * f16vec4;
const uint32_t slope = Br * acctype;
- const uint32_t total_size = tmpsh + Qf + Psh + sfsh + ksh + pvsh + slope;
+ const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope;
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported);
return ctx->device->pipeline_add_id_f32;
}
return nullptr;
+ case GGML_OP_OUT_PROD:
+ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+ return ctx->device->pipeline_out_prod_f32;
+ }
+ return nullptr;
case GGML_OP_CONCAT: {
if (src0->type != src1->type || src0->type != dst->type) {
return nullptr;
case GGML_OP_DUP:
return ggml_vk_get_cpy_pipeline(ctx, src0, dst, dst->type);
case GGML_OP_SET_ROWS:
- if (src1->type == GGML_TYPE_I64) {
- return ctx->device->pipeline_set_rows_i64[dst->type];
- } else {
- return ctx->device->pipeline_set_rows_i32[dst->type];
+ {
+ if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) {
+ return nullptr;
+ }
+ const int src_idx = src0->type == GGML_TYPE_F16;
+ if (src1->type == GGML_TYPE_I64) {
+ return ctx->device->pipeline_set_rows_i64[src_idx][dst->type];
+ } else if (src1->type == GGML_TYPE_I32) {
+ return ctx->device->pipeline_set_rows_i32[src_idx][dst->type];
+ }
+ return nullptr;
}
case GGML_OP_SILU_BACK:
if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
case GGML_OP_DIV:
case GGML_OP_MUL:
case GGML_OP_ADD1:
+ case GGML_OP_OUT_PROD:
case GGML_OP_ARANGE:
case GGML_OP_FILL:
case GGML_OP_SCALE:
});
}
+static void ggml_vk_out_prod(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+ const uint32_t src0_type_size = ggml_type_size(src0->type);
+ const uint32_t src1_type_size = ggml_type_size(src1->type);
+ const uint32_t dst_type_size = ggml_type_size(dst->type);
+
+ ggml_vk_op_f32<vk_op_binary_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_OUT_PROD, {
+ (uint32_t)ggml_nelements(dst),
+ (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3],
+ (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size,
+ (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3],
+ (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size,
+ (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3],
+ (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size,
+ 0,
+ 0.0f, 0.0f, 0,
+ });
+}
+
static void ggml_vk_sub(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
const uint32_t src0_type_size = ggml_type_size(src0->type);
const uint32_t src1_type_size = ggml_type_size(src1->type);
ggml_vk_add(ctx, compute_ctx, src0, src1, node);
}
break;
+ case GGML_OP_OUT_PROD:
+ ggml_vk_out_prod(ctx, compute_ctx, src0, src1, node);
+ break;
case GGML_OP_SUB:
ggml_vk_sub(ctx, compute_ctx, src0, src1, node);
vk_context cpy_ctx;
if (ctx->device->async_use_transfer_queue) {
- if (ctx->transfer_ctx.expired()) {
- cpy_ctx = ggml_vk_create_context(ctx, ctx->transfer_cmd_pool);
- ctx->transfer_ctx = cpy_ctx;
- ggml_vk_ctx_begin(ctx->device, cpy_ctx);
- } else {
- cpy_ctx = ctx->transfer_ctx.lock();
- }
+ cpy_ctx = ggml_vk_get_transfer_ctx(ctx);
} else {
cpy_ctx = ggml_vk_get_compute_ctx(ctx);
}
vk_context cpy_ctx;
if (ctx->device->async_use_transfer_queue) {
- if (ctx->transfer_ctx.expired()) {
- cpy_ctx = ggml_vk_create_context(ctx, ctx->transfer_cmd_pool);
- ctx->transfer_ctx = cpy_ctx;
- ggml_vk_ctx_begin(ctx->device, cpy_ctx);
- } else {
- cpy_ctx = ctx->transfer_ctx.lock();
- }
+ cpy_ctx = ggml_vk_get_transfer_ctx(ctx);
} else {
cpy_ctx = ggml_vk_get_compute_ctx(ctx);
}
1, &ctx->transfer_semaphore.value,
0, nullptr,
};
- vk::PipelineStageFlags stage = ctx->device->transfer_queue.stage_flags;
+ vk::PipelineStageFlags stage = ctx->device->transfer_queue->stage_flags;
vk::SubmitInfo si{
1, &ctx->transfer_semaphore.s, &stage,
0, nullptr,
0, nullptr,
};
si.setPNext(&tl_info);
- std::lock_guard<std::mutex> guard(queue_mutex);
- ctx->device->compute_queue.queue.submit({ si }, ctx->fence);
+ ctx->device->compute_queue->handle->submit({ si }, ctx->fence);
ctx->transfer_semaphore_last_submitted = ctx->transfer_semaphore.value;
} else {
- std::lock_guard<std::mutex> guard(queue_mutex);
- ctx->device->compute_queue.queue.submit({}, ctx->fence);
+ ctx->device->compute_queue->handle->submit({}, ctx->fence);
}
ggml_vk_wait_for_fence(ctx);
ctx->submit_pending = false;
vk::DebugUtilsLabelEXT dul = {};
dul.pLabelName = "ggml_backend_vk_graph_compute";
dul.color = std::array<float,4>{1.0f, 1.0f, 1.0f, 1.0f};
- vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue.queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
+
+ std::lock_guard<vk_queue_handle> guard(*ctx->device->compute_queue->handle);
+ vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue->handle->queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
}
ctx->prealloc_size_add_rms_partials_offset = 0;
if (vkev->has_event) {
// Wait for latest event
ggml_vk_wait_events(compute_ctx, { vkev->event });
+
+ if (ctx->device->async_use_transfer_queue) {
+ vk_context transfer_ctx = ggml_vk_get_transfer_ctx(ctx);
+ transfer_ctx->s->wait_semaphores.push_back(vkev->tl_semaphore);
+ }
}
}
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) {
return false;
}
- auto fa_kv_ok = [coopmat2](ggml_type t) {
+ auto fa_kv_ok = [](ggml_type t) {
switch (t) {
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q4_0:
+ case GGML_TYPE_IQ4_NL:
return true;
- case GGML_TYPE_Q1_0:
- return coopmat2;
default:
return false;
}
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_SET_ROWS:
{
- if (op->src[0]->type == GGML_TYPE_F32) {
- switch (op->type) {
- case GGML_TYPE_F32:
- case GGML_TYPE_F16:
- case GGML_TYPE_BF16:
- case GGML_TYPE_Q1_0:
- case GGML_TYPE_Q4_0:
- case GGML_TYPE_Q4_1:
- case GGML_TYPE_Q5_0:
- case GGML_TYPE_Q5_1:
- case GGML_TYPE_Q8_0:
- case GGML_TYPE_IQ4_NL:
- return true;
- default:
- return false;
- }
+ if ((op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) ||
+ (op->src[1]->type != GGML_TYPE_I32 && op->src[1]->type != GGML_TYPE_I64)) {
+ return false;
+ }
+ switch (op->type) {
+ case GGML_TYPE_F32:
+ case GGML_TYPE_F16:
+ case GGML_TYPE_BF16:
+ case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
+ case GGML_TYPE_Q4_0:
+ case GGML_TYPE_Q4_1:
+ case GGML_TYPE_Q5_0:
+ case GGML_TYPE_Q5_1:
+ case GGML_TYPE_Q8_0:
+ case GGML_TYPE_IQ4_NL:
+ return true;
+ default:
+ return false;
}
- return false;
}
case GGML_OP_CONT:
case GGML_OP_CPY:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
+ case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
case GGML_OP_OPT_STEP_ADAMW:
case GGML_OP_OPT_STEP_SGD:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
+ case GGML_OP_OUT_PROD:
+ return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
+ && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+ && op->type == GGML_TYPE_F32;
case GGML_OP_LOG:
case GGML_OP_TRI:
case GGML_OP_DIAG:
add_compile_definitions(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
message(STATUS "Enabling bfloat16 glslc support")
endif()
+if (GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
+ add_compile_definitions(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
+ message(STATUS "Enabling E2M1 glslc support")
+endif()
+if (GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ add_compile_definitions(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ message(STATUS "Enabling E4M3 glslc support")
+endif()
if (GGML_VULKAN_SHADER_DEBUG_INFO)
add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO)
message(STATUS "Enabling shader debug info")
const uint BLOCK_SIZE = 32;
#endif
-layout (binding = 0) readonly buffer S {float data_s[];};
+layout (binding = 0) readonly buffer S {S_TYPE data_s[];};
#if defined(SET_ROWS)
#include "generic_binary_head.glsl"
float vmax = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_Q4_0; ++j) {
- const float v = data_s[src_idx + j];
+ const float v = float(data_s[src_idx + j]);
if (amax < abs(v)) {
amax = abs(v);
vmax = v;
data_q[dst_idx].d = float16_t(d);
[[unroll]] for (int j = 0; j < QUANT_K_Q4_0/2; ++j) {
- const float x0 = data_s[src_idx + 0 + j]*id;
- const float x1 = data_s[src_idx + QUANT_K_Q4_0/2 + j]*id;
+ const float x0 = float(data_s[src_idx + 0 + j])*id;
+ const float x1 = float(data_s[src_idx + QUANT_K_Q4_0/2 + j])*id;
const uint xi0 = min(15, int(x0 + 8.5));
const uint xi1 = min(15, int(x1 + 8.5));
float vmax = -vmin;
[[unroll]] for (int j = 0; j < QUANT_K_Q4_1; ++j) {
- const float v = data_s[src_idx + j];
+ const float v = float(data_s[src_idx + j]);
if (v < vmin) vmin = v;
if (v > vmax) vmax = v;
data_q[dst_idx].m = float16_t(vmin);
[[unroll]] for (int j = 0; j < QUANT_K_Q4_1/2; ++j) {
- const float x0 = (data_s[src_idx + 0 + j] - vmin)*id;
- const float x1 = (data_s[src_idx + QUANT_K_Q4_1/2 + j] - vmin)*id;
+ const float x0 = (float(data_s[src_idx + 0 + j]) - vmin)*id;
+ const float x1 = (float(data_s[src_idx + QUANT_K_Q4_1/2 + j]) - vmin)*id;
const uint xi0 = min(15, int(x0 + 0.5));
const uint xi1 = min(15, int(x1 + 0.5));
float vmax = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_Q5_0; ++j) {
- const float v = data_s[src_idx + j];
+ const float v = float(data_s[src_idx + j]);
if (amax < abs(v)) {
amax = abs(v);
vmax = v;
uint32_t qh = 0;
[[unroll]] for (int j = 0; j < QUANT_K_Q5_0/2; ++j) {
- const float x0 = data_s[src_idx + 0 + j]*id;
- const float x1 = data_s[src_idx + QUANT_K_Q5_0/2 + j]*id;
+ const float x0 = float(data_s[src_idx + 0 + j])*id;
+ const float x1 = float(data_s[src_idx + QUANT_K_Q5_0/2 + j])*id;
const uint xi0 = min(31, int(x0 + 16.5));
const uint xi1 = min(31, int(x1 + 16.5));
#if defined(DATA_A_Q5_1)
void quantize(uint dst_idx, uint src_idx)
{
- float min = data_s[src_idx + 0];
+ float min = float(data_s[src_idx + 0]);
float max = min;
[[unroll]] for (int j = 1; j < QUANT_K_Q5_1; ++j) {
- const float v = data_s[src_idx + j];
+ const float v = float(data_s[src_idx + j]);
min = v < min ? v : min;
max = v > max ? v : max;
}
uint32_t qh = 0;
[[unroll]] for (int j = 0; j < QUANT_K_Q5_1/2; ++j) {
- const float x0 = (data_s[src_idx + 0 + j] - min)*id;
- const float x1 = (data_s[src_idx + QUANT_K_Q5_1/2 + j] - min)*id;
+ const float x0 = (float(data_s[src_idx + 0 + j]) - min)*id;
+ const float x1 = (float(data_s[src_idx + QUANT_K_Q5_1/2 + j]) - min)*id;
const uint xi0 = uint(x0 + 0.5);
const uint xi1 = uint(x1 + 0.5);
float amax = 0.0; // absolute max
[[unroll]] for (int j = 0; j < QUANT_K_Q8_0; j++) {
- const float v = data_s[src_idx + j];
+ const float v = float(data_s[src_idx + j]);
amax = max(amax, abs(v));
}
data_q[dst_idx].d = float16_t(d);
[[unroll]] for (int j = 0; j < QUANT_K_Q8_0; ++j) {
- const float x0 = data_s[src_idx + j]*id;
+ const float x0 = float(data_s[src_idx + j])*id;
data_q[dst_idx].qs[j] = int8_t(round(x0));
}
float sum_abs = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_Q1_0; j++) {
- sum_abs += abs(data_s[src_idx + j]);
+ sum_abs += abs(float(data_s[src_idx + j]));
}
const float d = sum_abs / QUANT_K_Q1_0;
}
[[unroll]] for (int j = 0; j < QUANT_K_Q1_0; ++j) {
- if (data_s[src_idx + j] >= 0.0) {
+ if (float(data_s[src_idx + j]) >= 0.0) {
data_q[dst_idx].qs[j / 8] |= uint8_t(1 << (j % 8));
}
}
}
#endif
+#if defined(DATA_A_Q2_0)
+uint quantize_q2_0(float x)
+{
+ const int q = int(x >= 0.0f ? floor(x + 0.5f) : ceil(x - 0.5f)) + 1;
+ return uint(clamp(q, 0, 3));
+}
+
+void quantize(uint dst_idx, uint src_idx)
+{
+ float amax = 0.0f;
+
+ [[unroll]] for (int j = 0; j < QUANT_K_Q2_0; ++j) {
+ amax = max(amax, abs(float(data_s[src_idx + j])));
+ }
+
+ const float d = amax;
+ const float id = d != 0.0f ? 1.0f / d : 0.0f;
+
+ data_q[dst_idx].d = float16_t(d);
+
+ [[unroll]] for (int j = 0; j < QUANT_K_Q2_0 / 4; ++j) {
+ const uint q0 = quantize_q2_0(float(data_s[src_idx + 4*j ]) * id);
+ const uint q1 = quantize_q2_0(float(data_s[src_idx + 4*j + 1]) * id);
+ const uint q2 = quantize_q2_0(float(data_s[src_idx + 4*j + 2]) * id);
+ const uint q3 = quantize_q2_0(float(data_s[src_idx + 4*j + 3]) * id);
+ data_q[dst_idx].qs[j] = uint8_t(q0 | (q1 << 2u) | (q2 << 4u) | (q3 << 6u));
+ }
+}
+#endif
+
#if defined(DATA_A_IQ4_NL)
uint best_index(float x) {
if (x <= kvalues_iq4nl[0]) return 0;
float vmax = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_IQ4_NL; ++j) {
- const float v = data_s[src_idx + j];
+ const float v = float(data_s[src_idx + j]);
if (amax < abs(v)) {
amax = abs(v);
vmax = v;
float sumqx = 0, sumq2 = 0;
[[unroll]] for (int j = 0; j < QUANT_K_IQ4_NL/2; ++j) {
- const float x0 = data_s[src_idx + 0 + j]*id;
- const float x1 = data_s[src_idx + QUANT_K_IQ4_NL/2 + j]*id;
+ const float x0 = float(data_s[src_idx + 0 + j])*id;
+ const float x1 = float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j])*id;
const uint xi0 = best_index(x0);
const uint xi1 = best_index(x1);
data_q[dst_idx].qs[j] = uint8_t(xi0 | (xi1 << 4));
const float v0 = kvalues_iq4nl[xi0];
const float v1 = kvalues_iq4nl[xi1];
- const float w0 = data_s[src_idx + 0 + j]*data_s[src_idx + 0 + j];
- const float w1 = data_s[src_idx + QUANT_K_IQ4_NL/2 + j]*data_s[src_idx + QUANT_K_IQ4_NL/2 + j];
- sumqx += w0*v0*data_s[src_idx + j] + w1*v1*data_s[src_idx + QUANT_K_IQ4_NL/2 + j];
+ const float w0 = float(data_s[src_idx + 0 + j])*float(data_s[src_idx + 0 + j]);
+ const float w1 = float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j])*float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j]);
+ sumqx += w0*v0*float(data_s[src_idx + j]) + w1*v1*float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j]);
sumq2 += w0*v0*v0 + w1*v1*v1;
}
#if defined(DATA_A_F32) || defined(DATA_A_F16)
void quantize(uint dst_idx, uint src_idx)
{
- data_q[dst_idx] = A_TYPE(data_s[src_idx]);
+ data_q[dst_idx] = A_TYPE(float(data_s[src_idx]));
}
#endif
#if defined(DATA_A_BF16)
void quantize(uint dst_idx, uint src_idx)
{
- data_q[dst_idx] = A_TYPE(fp32_to_bf16(data_s[src_idx]));
+ data_q[dst_idx] = A_TYPE(fp32_to_bf16(float(data_s[src_idx])));
}
#endif
}
#endif
+#if defined(DATA_A_Q2_0)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]) >> (2u * (iqs % 4u));
+ return vec2(bits & 3u, (bits >> 2u) & 3u) - 1.0f;
+}
+vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
+ const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]);
+ return vec4(bits & 3u, (bits >> 2u) & 3u, (bits >> 4u) & 3u, bits >> 6u) - 1.0f;
+}
+#endif
+
#if defined(DATA_A_IQ1_S)
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
const uint ib32 = iqs / 32;
#if defined(DATA_A_MXFP4)
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+#ifdef USE_OCP_FP4
+ return vec2(unpackFloat2xfe2m1EXT(uint8_t(vui)));
+#else
return vec2(kvalues_mxfp4[vui & 0xF], kvalues_mxfp4[vui >> 4]) * 0.5;
+#endif
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
+#ifdef USE_OCP_FP4
+ const uint16_t vui = uint16_t(uint(data_a[a_offset + ib].qs[iqs]) |
+ uint(data_a[a_offset + ib].qs[iqs + 1]) << 8);
+ return vec4(unpackFloat4xfe2m1EXT(vui));
+#else
vec2 v0 = dequantize(ib, iqs, a_offset);
vec2 v1 = dequantize(ib, iqs + 1, a_offset);
return vec4(v0.x, v0.y, v1.x, v1.y);
+#endif
}
#endif
const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
const uint j = iqs & 7;
const uint shift = (iqs & 8) >> 1; // 0 or 4
+#ifdef USE_OCP_FP4
+ const uint vui = uint(data_a_packed16[a_offset + ib].qs[(sub * 8u + j) / 2u]);
+ return vec2(bitcastExtractfe2m1EXT(unpack8(vui).xy, shift)) * d;
+#else
const uint vui0 = uint(data_a[a_offset + ib].qs[sub * 8u + j]);
const uint vui1 = uint(data_a[a_offset + ib].qs[sub * 8u + j + 1]);
const uint qs0 = (vui0 >> shift) & 0xF;
const uint qs1 = (vui1 >> shift) & 0xF;
return vec2(float(kvalues_mxfp4[qs0]), float(kvalues_mxfp4[qs1])) * d * 0.5;
+#endif
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
+#ifdef USE_OCP_FP4
+ const uint sub = iqs >> 4;
+ const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
+ const uint j = iqs & 7;
+ const uint shift = (iqs & 8) >> 1; // 0 or 4
+ const uint vui = data_a_packed32[a_offset + ib].qs[(sub * 8u + j) / 4u];
+ return vec4(bitcastExtractfe2m1EXT(unpack8(vui), shift)) * d;
+#else
const vec2 v0 = dequantize(ib, iqs, a_offset);
const vec2 v1 = dequantize(ib, iqs + 2u, a_offset);
return vec4(v0.x, v0.y, v1.x, v1.y);
+#endif
}
#endif
}
#endif
-#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
vec2 get_dm(uint ib, uint a_offset) {
return vec2(float(data_a[a_offset + ib].d), 0);
}
(qs_nib & 8u) != 0u ? d : md);
}
+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ2_0 {
+ block_q2_0 block;
+};
+
+float16_t dequantFuncQ2_0(const in decodeBufQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
+{
+ const float16_t d = bl.block.d;
+ const uint idx = coordInBlock[1];
+ const uint bits = uint(bl.block.qs[idx >> 2]) >> (2u * (idx & 3u));
+ return (float16_t(bits & 3u) - float16_t(1.0)) * d;
+}
+
+f16vec4 dequantFuncQ2_0_v(const in decodeBufQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
+{
+ const float16_t d = bl.block.d;
+ const uint idx = coordInBlock[1];
+ const uint bits = uint(bl.block.qs[idx >> 2]);
+ return f16vec4((vec4(bits & 3u, (bits >> 2u) & 3u, (bits >> 4u) & 3u, bits >> 6u) - 1.0f) * float(d));
+}
+
layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 {
block_q4_0_packed16 block;
};
const uint idx = coordInBlock[1];
const uint iqs = idx & 0xF;
const uint shift = (idx & 0x10) >> 2;
+#ifdef USE_OCP_FP4
+ return float16_t(bitcastExtractfe2m1EXT(bl.block.qs[iqs], shift)) * float16_t(d);
+#else
uint32_t qs = bl.block.qs[iqs];
qs >>= shift;
qs &= 0xF;
float16_t ret = float16_t(kvalues_mxfp4[qs] * d * 0.5);
return ret;
+#endif
}
f16vec4 dequantFuncMXFP4_v(const in decodeBufMXFP4 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
const uint idx = coordInBlock[1];
const uint iqs = idx & 0xF;
const uint shift = (idx & 0x10) >> 2;
+#ifdef USE_OCP_FP4
+ const fe2m1vec4 qv = bitcastExtractfe2m1EXT(
+ u8vec4(
+ bl.block.qs[iqs],
+ bl.block.qs[iqs + 1u],
+ bl.block.qs[iqs + 2u],
+ bl.block.qs[iqs + 3u]),
+ shift);
+ return f16vec4(qv) * float16_t(d);
+#else
uvec4 qv = uvec4(
uint(bl.block.qs[iqs]),
uint(bl.block.qs[iqs + 1u]),
float(kvalues_mxfp4[qv.z]),
float(kvalues_mxfp4[qv.w])) * d * 0.5f;
return f16vec4(ret);
+#endif
}
#endif
const uint sub = (idx & 0x30) >> 4;
const uint iqs = ((idx & 0x30) >> 1) + (idx & 0x7);
const uint shift = (idx & 0x8) >> 1;
+#ifdef USE_OCP_FP4
+ const float16_t d = float16_t(ue4m3_from_bits(bl.block.d[sub]));
+ return float16_t(bitcastExtractfe2m1EXT(bl.block.qs[iqs], shift)) * d;
+#else
const float d = ue4m3_to_fp32(bl.block.d[sub]);
uint qs = uint(bl.block.qs[iqs]);
qs = (qs >> shift) & 0xF;
return float16_t(kvalues_mxfp4[qs] * d * 0.5);
+#endif
}
f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
const uint sub = idx >> 4;
const uint qs_w = ((idx & 0x30) >> 3) + ((idx & 0x4u) >> 2); // iqs / 4, in [0,8)
const uint shift = (idx & 0x8) >> 1;
- const float d = ue4m3_to_fp32(bl.block.d[sub]);
const uint qsw = uint32_t(bl32.block.qs[qs_w]);
+#ifdef USE_OCP_FP4
+ const float16_t d = float16_t(ue4m3_from_bits(bl.block.d[sub]));
+ const fe2m1vec4 qv = bitcastExtractfe2m1EXT(unpack8(qsw), shift);
+ return f16vec4(qv) * d;
+#else
+ const float d = ue4m3_to_fp32(bl.block.d[sub]);
const u8vec4 qv = unpack8((qsw >> shift) & 0x0F0F0F0Fu);
const vec4 ret = vec4(
float(kvalues_mxfp4[qv.x]),
float(kvalues_mxfp4[qv.z]),
float(kvalues_mxfp4[qv.w])) * d * 0.5f;
return f16vec4(ret);
+#endif
}
#endif
#if defined(DATA_A_Q1_0)
#define dequantFuncA dequantFuncQ1_0
#define dequantFuncA_v dequantFuncQ1_0_v
+#elif defined(DATA_A_Q2_0)
+#define dequantFuncA dequantFuncQ2_0
+#define dequantFuncA_v dequantFuncQ2_0_v
#elif defined(DATA_A_Q4_0)
#define dequantFuncA dequantFuncQ4_0
#define dequantFuncA_v dequantFuncQ4_0_v
--- /dev/null
+#version 450
+
+#include "dequant_head.glsl"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q2_0 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid / 4;
+ const uint ir = tid % 4;
+ const uint ib = 4*i + ir;
+ if (ib >= p.nel / QUANT_K_Q2_0) {
+ return;
+ }
+
+ const uint b_idx = 256*i + QUANT_K_Q2_0*ir + 4*il;
+ const uint bits = uint(data_a[ib].qs[il]);
+ const float d = float(data_a[ib].d);
+
+ data_b[b_idx ] = D_TYPE(d * (float(bits & 3u) - 1.0f));
+ data_b[b_idx + 1] = D_TYPE(d * (float((bits >> 2u) & 3u) - 1.0f));
+ data_b[b_idx + 2] = D_TYPE(d * (float((bits >> 4u) & 3u) - 1.0f));
+ data_b[b_idx + 3] = D_TYPE(d * (float(bits >> 6u) - 1.0f));
+}
--- /dev/null
+#version 460
+
+#extension GL_EXT_float_e2m1 : require
+
+void main()
+{
+}
--- /dev/null
+#version 460
+
+#extension GL_EXT_float_e4m3 : require
+
+void main()
+{
+}
void main() {
#ifdef NEEDS_INIT_IQ_SHMEM
- init_iq_shmem(gl_WorkGroupSize);
+ if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
+ init_iq_shmem(gl_WorkGroupSize);
+ }
#endif
init_indices();
#define FA_TYPE_Q5_0 6u
#define FA_TYPE_Q5_1 7u
#define FA_TYPE_Q8_0 8u
+#define FA_TYPE_IQ4_NL 20u
#define FA_TYPE_BF16 30u
-#define FA_TYPE_Q1_0 41u
#if defined(BFLOAT16)
#define O_TYPE float
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
+ case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
case FA_TYPE_BF16: return 1u;
- case FA_TYPE_Q1_0: return uint(QUANT_K_Q1_0); // cm2-only, harmless elsewhere
default: return 1u;
}
}
}
}
+bool fa_type_needs_shmem(uint ty) {
+ switch (ty) {
+ case FA_TYPE_IQ4_NL: return true;
+ default: return false;
+ }
+}
+
// These can't be `const` globals because GLSL forbids function calls in global
// const initializers, even when the spec constants would let the driver fold
// them. Macros expand at the use site and fold after specialization.
void main() {
#ifdef NEEDS_INIT_IQ_SHMEM
- init_iq_shmem(gl_WorkGroupSize);
+ if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
+ init_iq_shmem(gl_WorkGroupSize);
+ }
#endif
init_indices();
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
- case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
default: return float16_t(0);
}
}
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
- case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
default: return float16_t(0);
}
}
// V=4 vector decode for K/V; dispatches to per-format _v decoders.
f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
switch (FaTypeK) {
- case 0u: return f16vec4(decodeBufF32(bl_in).block);
- case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
- case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
- case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
- case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
- case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
- case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
+ case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
default: return f16vec4(0);
}
}
f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
switch (FaTypeV) {
- case 0u: return f16vec4(decodeBufF32(bl_in).block);
- case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
- case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
- case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
- case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
- case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
- case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
+ case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
+ case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
default: return f16vec4(0);
}
}
}
void main() {
+#ifdef NEEDS_INIT_IQ_SHMEM
+ if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
+ init_iq_shmem(gl_WorkGroupSize);
+ }
+#endif
+
init_indices();
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, HSK_pad, Bc, gl_MatrixUseB> K_T;
uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13;
- // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Q4/Q8 family: bs_k==32. Q1_0: bs_k==128.
+ // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32.
#if defined(BFLOAT16)
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
#else
layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1;
layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0;
layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0;
+layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl;
+layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl;
layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16;
layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16;
return FLOAT_TYPE(BUF.data[a_offset + ib].d) * FLOAT_TYPEV4(v0.x, v0.y, v1.x, v1.y); \
}
+#define FA_DEQUANT4_IQ4_NL(BUF) { \
+ const uint shift = (iqs & 0x10) >> 2; \
+ const uint qs_i = (iqs & 0xC) >> 1; \
+ const uint qsw = uint(BUF.data[a_offset + ib].qs[qs_i]) \
+ | (uint(BUF.data[a_offset + ib].qs[qs_i + 1u]) << 16); \
+ const FLOAT_TYPE d = FLOAT_TYPE(BUF.data[a_offset + ib].d); \
+ const u8vec4 q = unpack8((qsw >> shift) & 0x0F0F0F0Fu); \
+ return d * FLOAT_TYPEV4(kvalues_iq4nl[q.x], kvalues_iq4nl[q.y], \
+ kvalues_iq4nl[q.z], kvalues_iq4nl[q.w]); \
+}
+
#define FA_DEQUANT4_BF16(BUF) \
return FLOAT_TYPEV4(bf16_to_fp32(uvec4(BUF.data[(a_offset + ib) / 4])));
case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
+ case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
}
} else {
case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
+ case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
}
}
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
-#if defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
-#define K_PER_ITER 8
-#elif defined(DATA_A_QUANT_K)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_QUANT_K)
#define K_PER_ITER 16
+#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
+#define K_PER_ITER 8
#elif defined(DATA_A_IQ1_S) || defined(DATA_A_IQ1_M)
#define K_PER_ITER 32
#else
#include "types.glsl"
-#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
+#if defined(DATA_A_Q2_0)
+FLOAT_TYPE get_dm(uint ib) {
+ return FLOAT_TYPE(data_a[ib / 2].d);
+}
+#elif defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
FLOAT_TYPE get_dm(uint ib) {
return FLOAT_TYPE(data_a[ib].d);
}
#endif
// Each iqs value maps to a 32-bit integer
+#if defined(DATA_A_Q2_0)
+uint unpack_q2_0(uint bits) {
+ // Move bit pairs [1:0], [3:2], [5:4], [7:6] to [1:0], [9:8], [17:16], [25:24].
+ bits &= 0xffu;
+ bits = (bits | (bits << 12u)) & 0x000f000fu;
+ return (bits | (bits << 6u)) & 0x03030303u;
+}
+
+i32vec4 repack4(uint ib, uint iqs) {
+ const uint qs_idx = (ib & 1u) * 4u + iqs * 2u;
+ const uint bits = pack32(u16vec2(data_a_packed16[ib / 2].qs[qs_idx],
+ data_a_packed16[ib / 2].qs[qs_idx + 1]));
+ return i32vec4(unpack_q2_0(bits), unpack_q2_0(bits >> 8u),
+ unpack_q2_0(bits >> 16u), unpack_q2_0(bits >> 24u));
+}
+
+FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const int32_t sum_divisor) {
+ return FLOAT_TYPE(da * (float(q_sum) * dsb.x - dsb.y / float(sum_divisor)));
+}
+#endif
+
#if defined(DATA_A_Q4_0)
// 2-byte loads for Q4_0 blocks (18 bytes)
i32vec2 repack(uint ib, uint iqs) {
}
#endif
-#if defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
+#if defined(DATA_A_Q2_0)
+FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) {
+ int32_t q_sum = 0;
+ const i32vec4 qs_a = repack4(ib_a, iqs);
+ q_sum += dotPacked4x8EXT(qs_a.x, cache_b_qs[0]);
+ q_sum += dotPacked4x8EXT(qs_a.y, cache_b_qs[1]);
+ q_sum += dotPacked4x8EXT(qs_a.z, cache_b_qs[2]);
+ q_sum += dotPacked4x8EXT(qs_a.w, cache_b_qs[3]);
+
+ // 16 quants per call => divide sums by 32/16 = 2
+ return mul_q8_1(q_sum, get_dm(ib_a), cache_b_ds, 2);
+}
+#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) {
int32_t q_sum = 0;
#if QUANT_R == 2
buf_a[buf_idx + 1] = FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d);
buf_a[buf_idx + 2] = FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d);
buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d);
+#elif defined(DATA_A_Q2_0)
+ const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
+ const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
+
+ const uint ib = idx / 16;
+ const uint iqs = idx & 0xfu;
+
+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib].d);
+ const uint bits = uint(data_a[ib].qs[iqs]);
+
+ buf_a[buf_idx ] = d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f));
+ buf_a[buf_idx + 1] = d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f));
#elif defined(DATA_A_Q2_K)
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
const uint ib = idx / 8;
const uint iqs = (idx & 0x07) * 2;
- const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
const uint vui = uint(data_a[ib].qs[iqs]);
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
+#ifdef USE_OCP_FP4
+ const float d = e8m0_to_fp32(data_a[ib].e);
+ const u8vec2 packed = u8vec2(vui, vui2);
+ buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d);
+ buf_a[buf_idx + 8] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d);
+#else
+ const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
kvalues_mxfp4[vui2 & 0xF] * d);
buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
kvalues_mxfp4[vui2 >> 4] * d);
+#endif
#elif defined(DATA_A_NVFP4)
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
// lo and hi nibbles are 8 elements apart, which doesn't quite line up with
const uint ib = idx / 16u;
const uint sub = (idx & 0xC) >> 2;
const uint iqs = (idx & 0xF) * 2;
- const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
const uint vui = uint(data_a[ib].qs[iqs]);
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
+#ifdef USE_OCP_FP4
+ const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub]));
+ const u8vec2 packed = u8vec2(vui, vui2);
+ buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d;
+ buf_a[buf_idx + 4] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d;
+#else
+ const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
kvalues_mxfp4[vui2 & 0xF] * d);
buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
kvalues_mxfp4[vui2 >> 4] * d);
#endif
+#endif
}
#if !defined(MUL_MAT_ID)
// Each iqs value maps to a 32-bit integer
+#if defined(DATA_A_Q2_0)
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
+ const uint block_idx = ib / 2;
+ const uint byte_idx = (ib & 1u) * 8u + iqs;
+ const uint bits = uint(data_a[block_idx].qs[byte_idx]);
+ buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
+ int8_t(bits & 3u),
+ int8_t((bits >> 2u) & 3u),
+ int8_t((bits >> 4u) & 3u),
+ int8_t(bits >> 6u)));
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
+ }
+}
+
+void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
+ cache_a[reg_ib].dm = buf_a[buf_ib].dm;
+
+ [[unroll]] for (uint iqs = 0; iqs < 8; ++iqs) {
+ cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
+ }
+}
+
+ACC_TYPE mmq_dot_product(const uint ib_a) {
+ int32_t q_sum = 0;
+ [[unroll]] for (uint iqs = 0; iqs < 8; ++iqs) {
+ q_sum += dotPacked4x8EXT(cache_a[ib_a].qs[iqs], cache_b.qs[iqs]);
+ }
+
+ return ACC_TYPE(float(cache_a[ib_a].dm) * (float(q_sum) * float(cache_b.ds.x) - float(cache_b.ds.y)));
+}
+#endif
+
#if defined(DATA_A_Q4_0) || defined(DATA_A_Q4_1)
// 2-byte loads for Q4_0 blocks (18 bytes)
// 4-byte loads for Q4_1 blocks (20 bytes)
uint32_t qs[16/4];
FLOAT_TYPE dm;
};
+#elif defined(DATA_A_Q2_0)
+#define QUANT_R_MMQ 1
+struct block_a_cache {
+ int32_t qs[8];
+ FLOAT_TYPE dm;
+};
#elif defined(DATA_A_Q4_1)
#define QUANT_R_MMQ 2
struct block_a_cache {
--- /dev/null
+#version 450
+
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+ uint ne;
+ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
+ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
+ uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23;
+ uint misalign_offsets;
+ float param1; float param2; int param3;
+} p;
+
+layout (binding = 0) readonly buffer A {float data_a[];};
+layout (binding = 1) readonly buffer B {float data_b[];};
+layout (binding = 2) writeonly buffer D {float data_d[];};
+
+uint get_idx() {
+ return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
+}
+
+uint get_aoffset() { return p.misalign_offsets >> 16; }
+uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; }
+uint get_doffset() { return p.misalign_offsets & 0xFF; }
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+void main() {
+ uint idx = get_idx();
+ if (idx >= p.ne) {
+ return;
+ }
+
+ uint tmp = idx;
+ uint i0 = tmp % p.ne20; tmp /= p.ne20;
+ uint i1 = tmp % p.ne21; tmp /= p.ne21;
+ uint i2 = tmp % p.ne22; tmp /= p.ne22;
+ uint i3 = tmp;
+
+ uint a_i0 = i0 % p.ne00;
+ uint a_i2 = i2 / (p.ne22 / p.ne02);
+ uint a_i3 = i3 / (p.ne23 / p.ne03);
+
+ uint b_i0 = i1 % p.ne10;
+ uint b_i2 = i2;
+ uint b_i3 = i3;
+
+ float sum = 0.0f;
+ uint K = p.ne01;
+ for (uint k = 0; k < K; k++) {
+ uint aoff = get_aoffset() + a_i3*p.nb03 + a_i2*p.nb02 + k*p.nb01 + a_i0*p.nb00;
+ uint boff = get_boffset() + b_i3*p.nb13 + b_i2*p.nb12 + k*p.nb11 + b_i0*p.nb10;
+ sum += data_a[aoff] * data_b[boff];
+ }
+
+ uint doff = get_doffset() + i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20;
+ data_d[doff] = sum;
+}
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
#extension GL_EXT_shader_16bit_storage : require
+#ifdef USE_OCP_FP4
+#extension GL_EXT_float_e2m1 : require
+#extension GL_EXT_float_e4m3 : require
+#endif
+
#if defined(DATA_A_F32)
#define QUANT_K 1
#define QUANT_R 1
#define A_TYPE block_q1_0
#endif
+#define QUANT_K_Q2_0 64
+#define QUANT_R_Q2_0 1
+
+struct block_q2_0
+{
+ float16_t d;
+ uint8_t qs[QUANT_K_Q2_0 / 4];
+};
+
+struct block_q2_0_packed16
+{
+ float16_t d;
+ uint16_t qs[QUANT_K_Q2_0 / 8];
+};
+
+#if defined(DATA_A_Q2_0)
+#define QUANT_K QUANT_K_Q2_0
+#define QUANT_R QUANT_R_Q2_0
+#define QUANT_AUXF 1
+#define A_TYPE block_q2_0
+#define A_TYPE_PACKED16 block_q2_0_packed16
+#define DATA_A_QUANT_LEGACY
+#endif
+
#define QUANT_K_Q8_1 32
#define QUANT_R_Q8_1 1
uint8_t qs[QUANT_K_NVFP4 / 2];
};
+struct block_nvfp4_packed16
+{
+ uint16_t d[QUANT_K_NVFP4 / 16 / 2];
+ uint16_t qs[QUANT_K_NVFP4 / 2 / 2];
+};
+
struct block_nvfp4_packed32
{
uint32_t d[QUANT_K_NVFP4 / 16 / 4];
#define QUANT_R QUANT_R_NVFP4
#define QUANT_AUXF 1
#define A_TYPE block_nvfp4
+#define A_TYPE_PACKED16 block_nvfp4_packed16
#define A_TYPE_PACKED32 block_nvfp4_packed32
#endif
#endif
#if defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4)
+#if !defined(USE_OCP_FP4)
const int8_t kvalues_mxfp4_const[16] = {
int8_t(0), int8_t(1), int8_t(2), int8_t(3), int8_t(4), int8_t(6), int8_t(8), int8_t(12),
int8_t(0), int8_t(-1), int8_t(-2), int8_t(-3), int8_t(-4), int8_t(-6), int8_t(-8), int8_t(-12),
};
shared int8_t kvalues_mxfp4[16];
+#endif
-#if defined(DATA_A_NVFP4)
+#if defined(DATA_A_NVFP4) && !defined(USE_OCP_FP4)
// UE4M3 scale in NVFP4 blocks use only 7 bits; sign (bit 7) is always zero.
shared float ue4m3_fp32_lut[128];
}
#endif
+#if !defined(USE_OCP_FP4)
#define NEEDS_INIT_IQ_SHMEM
void init_iq_shmem(uvec3 wgsize)
{
barrier();
}
#endif
+#endif
// returns the bfloat value in the low 16b.
// See ggml_compute_fp32_to_bf16
}
#if defined(DATA_A_NVFP4)
+#if defined(USE_OCP_FP4)
+floate4m3_t ue4m3_from_bits(uint8_t x) {
+ if (x == uint8_t(0x7F)) {
+ return floate4m3_t(0.0);
+ }
+ return uintBitsToFloate4m3EXT(x);
+}
+#endif
+
float ue4m3_to_fp32(uint8_t x) {
+#if defined(USE_OCP_FP4)
+ return float(ue4m3_from_bits(x));
+#else
return ue4m3_fp32_lut[uint(x)];
+#endif
}
#endif
"f32",
"f16",
"q1_0",
+ "q2_0",
"q4_0",
"q4_1",
"q5_0",
}
bool is_legacy_quant(const std::string& type_name) {
- return type_name == "q4_0" || type_name == "q4_1" || type_name == "q5_0" || type_name == "q5_1" || type_name == "q8_0";
+ return type_name == "q2_0" || type_name == "q4_0" || type_name == "q4_1" || type_name == "q5_0" || type_name == "q5_1" || type_name == "q8_0";
}
bool is_k_quant(const std::string& type_name) {
std::string load_vec_quant = "2";
if ((tname == "q1_0") || (tname == "q4_0") || (tname == "q4_1") || (tname == "q5_1") || (tname == "iq1_s") || (tname == "iq1_m") || (tname == "iq2_xxs") || (tname == "iq2_xs") || (tname == "iq2_s"))
load_vec_quant = "8";
- else if ((tname == "q5_0") || (tname == "q8_0") || (tname == "q2_k") || (tname == "q4_k") || (tname == "q5_k") || (tname == "iq3_xxs") || (tname == "iq3_s") || (tname == "iq4_xs") || (tname == "iq4_nl") || (tname == "mxfp4") || (tname == "nvfp4"))
+ else if ((tname == "q2_0") || (tname == "q5_0") || (tname == "q8_0") || (tname == "q2_k") || (tname == "q4_k") || (tname == "q5_k") || (tname == "iq3_xxs") || (tname == "iq3_s") || (tname == "iq4_xs") || (tname == "iq4_nl") || (tname == "mxfp4") || (tname == "nvfp4"))
load_vec_quant = "4";
if (tname == "bf16") {
string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
}
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ if ((coopmat || coopmat2) && (tname == "mxfp4" || tname == "nvfp4")) {
+ if (!coopmat2) {
+ string_to_spv(shader_name + "_" + tname + "_f32_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+ }
+ string_to_spv(shader_name + "_" + tname + "_f16_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+ }
+#endif
+
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
// Integer dot mmq performs better with f32 accumulators (different shader, skip for dot2)
if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) {
fa_base_dict["ACC_TYPE"] = fp16 && f16acc ? "float16_t" : "float";
fa_base_dict["ACC_TYPEV2"] = fp16 && f16acc ? "f16vec2" : "vec2";
fa_base_dict["ACC_TYPEV4"] = fp16 && f16acc ? "f16vec4" : "vec4";
+ // Compile IQ4_NL support into all FA variants so its shared LUT is available when K or V uses it.
+ fa_base_dict["DATA_A_IQ4_NL"] = "1";
if (fp16 && f16acc) {
fa_base_dict["ACC_TYPE_MAX"] = "float16_t(65504.0)";
}
string_to_spv("mul_mat_vec_" + tname + "_f32_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
string_to_spv("mul_mat_vec_" + tname + "_f16_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ if (tname == "mxfp4" || tname == "nvfp4") {
+ string_to_spv("mul_mat_vec_" + tname + "_f32_f32_ocp", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
+ string_to_spv("mul_mat_vec_" + tname + "_f16_f32_ocp", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}));
+ string_to_spv("mul_mat_vec_" + tname + "_f32_f32_ocp_subgroup", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
+ string_to_spv("mul_mat_vec_" + tname + "_f16_f32_ocp_subgroup", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
+ string_to_spv("mul_mat_vec_" + tname + "_f32_f32_ocp_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
+ string_to_spv("mul_mat_vec_" + tname + "_f16_f32_ocp_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
+ string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_ocp", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
+ string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_ocp_subgroup", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
+ string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_ocp_subgroup_no_shmem", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
+ }
+#endif
+
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
string_to_spv("cpy_transpose_16", "copy_transpose.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
- for (std::string t : {"q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
- string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
+ string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
string_to_spv("cpy_" + t + "_f32", "copy_from_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
}
- for (std::string t : {"f32", "f16", "bf16", "q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
- string_to_spv("set_rows_" + t + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(t), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
- string_to_spv("set_rows_" + t + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(t), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ for (auto src : {std::pair{"f32", "float"}, std::pair{"f16", "float16_t"}}) {
+ for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
+ string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }
}
auto get_type_str = [](bool f16) {
}
}
+ string_to_spv("out_prod_f32", "out_prod.comp", {});
+
string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("conv_transpose_1d_f32", "conv_transpose_1d.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
}
}
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+ for (const std::string& btype : {"f16", "f32"}) {
+ for (const std::string& tname : {"mxfp4", "nvfp4"}) {
+ hdr << "extern const void * arr_dmmv_" << tname << "_" << btype << "_f32_ocp_data[3];\n";
+ hdr << "extern const uint64_t arr_dmmv_" << tname << "_" << btype << "_f32_ocp_len[3];\n";
+ if (basename(input_filepath) == "mul_mat_vec.comp") {
+ src << "const void * arr_dmmv_" << tname << "_" << btype << "_f32_ocp_data[3] = {mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_data, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_data, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_no_shmem_data};\n";
+ src << "const uint64_t arr_dmmv_" << tname << "_" << btype << "_f32_ocp_len[3] = {mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_len, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_len, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_no_shmem_len};\n";
+ }
+ }
+ }
+ for (const std::string& tname : {"mxfp4", "nvfp4"}) {
+ hdr << "extern const void * arr_dmmv_id_" << tname << "_f32_f32_ocp_data[3];\n";
+ hdr << "extern const uint64_t arr_dmmv_id_" << tname << "_f32_f32_ocp_len[3];\n";
+ if (basename(input_filepath) == "mul_mat_vec.comp") {
+ src << "const void * arr_dmmv_id_" << tname << "_f32_f32_ocp_data[3] = {mul_mat_vec_id_" << tname << "_f32_f32_ocp_data, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_data, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_no_shmem_data};\n";
+ src << "const uint64_t arr_dmmv_id_" << tname << "_f32_f32_ocp_len[3] = {mul_mat_vec_id_" << tname << "_f32_f32_ocp_len, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_len, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_no_shmem_len};\n";
+ }
+ }
+#endif
+
if (input_filepath == "") {
write_file_if_changed(target_hpp, hdr.str());
}
return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) == ggml_webgpu_tensor_addr(b);
}
-inline bool ggml_webgpu_tensor_overlap(const ggml_tensor * a, const ggml_tensor * b) {
- return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) < ggml_webgpu_tensor_addr(b) + ggml_nbytes(b) &&
- ggml_webgpu_tensor_addr(b) < ggml_webgpu_tensor_addr(a) + ggml_nbytes(a);
-}
-
struct ggml_webgpu_shader_lib_context {
ggml_tensor * src0;
ggml_tensor * src1;
bool src_overlap = false;
};
+struct ggml_webgpu_glu_shader_decisions {
+ uint32_t wg_size = 0;
+ bool src_overlap = false;
+};
+
struct ggml_webgpu_processed_shader {
std::string wgsl;
std::string variant;
int type;
int d_state;
bool xbc_overlap;
+ bool a_overlap;
+ bool ids_overlap;
bool operator==(const ggml_webgpu_ssm_scan_pipeline_key & other) const {
- return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap;
+ return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap &&
+ a_overlap == other.a_overlap && ids_overlap == other.ids_overlap;
}
};
ggml_webgpu_hash_combine(seed, key.type);
ggml_webgpu_hash_combine(seed, key.d_state);
ggml_webgpu_hash_combine(seed, key.xbc_overlap);
+ ggml_webgpu_hash_combine(seed, key.a_overlap);
+ ggml_webgpu_hash_combine(seed, key.ids_overlap);
return seed;
}
};
uint32_t wg_size;
uint32_t tokens_per_tile;
bool xbc_overlap = false;
+ bool a_overlap = false;
+ bool ids_overlap = false;
};
/** Argsort **/
struct ggml_webgpu_rms_norm_mul_pipeline_key {
bool inplace; // rn_src == dst
bool overlap; // mul_src == dst
- bool src_overlap; // rn_src == mul_src
+ bool src_overlap; // rn_src binding overlaps mul_src binding
bool operator==(const ggml_webgpu_rms_norm_mul_pipeline_key & other) const {
return inplace == other.inplace && overlap == other.overlap && src_overlap == other.src_overlap;
}
};
+// Same type fields as conv2d plus the input layout (WHCN vs CWHN).
+struct ggml_webgpu_conv2d_dw_pipeline_key {
+ ggml_type weight_type;
+ ggml_type input_type;
+ ggml_type output_type;
+ bool whcn;
+
+ bool operator==(const ggml_webgpu_conv2d_dw_pipeline_key & other) const {
+ return weight_type == other.weight_type && input_type == other.input_type && output_type == other.output_type &&
+ whcn == other.whcn;
+ }
+};
+
+struct ggml_webgpu_conv2d_dw_pipeline_key_hash {
+ size_t operator()(const ggml_webgpu_conv2d_dw_pipeline_key & key) const {
+ size_t seed = 0;
+ ggml_webgpu_hash_combine(seed, key.weight_type);
+ ggml_webgpu_hash_combine(seed, key.input_type);
+ ggml_webgpu_hash_combine(seed, key.output_type);
+ ggml_webgpu_hash_combine(seed, key.whcn);
+ return seed;
+ }
+};
+
/** Im2Col **/
struct ggml_webgpu_im2col_pipeline_key {
ggml_type input_type;
inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key(
const ggml_webgpu_shader_lib_context & context,
- uint32_t kv_direct_align) {
+ uint32_t kv_direct_align,
+ bool kv_overlap) {
ggml_webgpu_flash_attn_common_pipeline_key key = {};
key.q_type = context.src0->type;
key.k_type = context.src1->type;
key.head_dim_qk = (uint32_t) context.src0->ne[0];
key.head_dim_v = (uint32_t) context.src2->ne[0];
key.kv_direct = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align);
- key.kv_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src2);
+ key.kv_overlap = kv_overlap;
key.has_mask = context.src3 != nullptr;
key.has_sinks = context.src4 != nullptr;
key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f;
ggml_glu_op glu_op;
ggml_type type;
bool split;
+ bool src_overlap;
bool operator==(const ggml_webgpu_glu_pipeline_key & other) const {
- return glu_op == other.glu_op && type == other.type && split == other.split;
+ return glu_op == other.glu_op && type == other.type && split == other.split && src_overlap == other.src_overlap;
}
};
ggml_webgpu_hash_combine(seed, key.glu_op);
ggml_webgpu_hash_combine(seed, key.type);
ggml_webgpu_hash_combine(seed, key.split);
+ ggml_webgpu_hash_combine(seed, key.src_overlap);
return seed;
}
};
soft_max_pipelines;
std::unordered_map<ggml_webgpu_conv2d_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_pipeline_key_hash>
conv2d_pipelines;
+ std::unordered_map<ggml_webgpu_conv2d_dw_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_dw_pipeline_key_hash>
+ conv2d_dw_pipelines;
std::unordered_map<ggml_webgpu_im2col_pipeline_key, webgpu_pipeline, ggml_webgpu_im2col_pipeline_key_hash>
im2col_pipelines;
return ssm_conv_pipelines[key];
}
- webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context,
+ bool xbc_overlap,
+ bool a_overlap,
+ bool ids_overlap) {
ggml_webgpu_ssm_scan_pipeline_key key = {};
key.type = context.dst->type;
key.d_state = (int) context.src0->ne[0];
- key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) &&
- ggml_webgpu_tensor_overlap(context.src1, context.src5);
+ key.xbc_overlap = xbc_overlap;
+ key.a_overlap = a_overlap;
+ key.ids_overlap = ids_overlap;
auto it = ssm_scan_pipelines.find(key);
if (it != ssm_scan_pipelines.end()) {
if (key.xbc_overlap) {
defines.push_back("XBC_OVERLAP");
}
-
+ if (key.a_overlap) {
+ defines.push_back("A_OVERLAP");
+ }
+ if (key.ids_overlap) {
+ defines.push_back("IDS_OVERLAP");
+ }
variant += "_d" + std::to_string(key.d_state);
auto processed = preprocessor.preprocess(wgsl_ssm_scan, defines);
decisions->wg_size = wg_size;
decisions->tokens_per_tile = tokens_per_tile;
decisions->xbc_overlap = key.xbc_overlap;
+ decisions->a_overlap = key.a_overlap;
+ decisions->ids_overlap = key.ids_overlap;
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
pipeline.context = decisions;
ssm_scan_pipelines[key] = pipeline;
return unary_pipelines[key];
}
- webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
ggml_webgpu_rms_norm_mul_pipeline_key key = {};
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst);
- key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1);
+ key.src_overlap = src_overlap;
auto it = rms_norm_mul_pipelines.find(key);
if (it != rms_norm_mul_pipelines.end()) {
return rms_norm_mul_pipelines[key];
}
- webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
ggml_webgpu_binary_pipeline_key key = {};
key.type = context.dst->type;
key.op = context.dst->op;
key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst);
key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst);
- key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1);
+ key.src_overlap = src_overlap;
auto it = binary_pipelines.find(key);
if (it != binary_pipelines.end()) {
return pipeline;
}
- webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
ggml_webgpu_concat_pipeline_key key = {};
key.type = context.dst->type;
- key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1);
+ key.src_overlap = src_overlap;
auto it = concat_pipelines.find(key);
if (it != concat_pipelines.end()) {
return repeat_pipelines[key];
}
- webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context, bool kv_overlap) {
const bool can_use_subgroup_matrix = ggml_webgpu_flash_attn_can_use_subgroup_matrix_path(
context.supports_subgroup_matrix, context.sg_mat_k, context.sg_mat_n, context.src0, context.src2);
ggml_webgpu_flash_attn_decisions decisions = {};
decisions.q_tile = decisions.use_sg_matrix ? context.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE;
ggml_webgpu_flash_attn_pipeline_key key = {};
- key.common =
- ggml_webgpu_flash_attn_make_common_pipeline_key(context, decisions.use_sg_matrix ? context.sg_mat_k : 1u);
+ key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(
+ context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap);
key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct;
key.use_sg_matrix = decisions.use_sg_matrix;
return flash_attn_pipelines[key];
}
- webgpu_pipeline get_flash_attn_vec_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_flash_attn_vec_pipeline(const ggml_webgpu_shader_lib_context & context, bool kv_overlap) {
ggml_webgpu_flash_attn_vec_pipeline_key key = {};
- key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(context, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
+ key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(context, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH,
+ kv_overlap);
auto it = flash_attn_vec_pipelines.find(key);
if (it != flash_attn_vec_pipelines.end()) {
return cpy_pipelines[key];
}
- webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) {
ggml_webgpu_glu_pipeline_key key = {};
key.glu_op = ggml_get_glu_op(context.dst);
key.type = context.dst->type;
key.split = (context.src1 != nullptr);
+ key.src_overlap = src_overlap;
auto it = glu_pipelines.find(key);
if (it != glu_pipelines.end()) {
GGML_ABORT("Unsupported type for GLU shader");
}
- if (key.split) {
+ if (key.src_overlap) {
+ defines.push_back("SRC_OVERLAP");
+ variant += "_src_overlap";
+ } else if (key.split) {
variant += "_split";
} else {
defines.push_back("NO_SPLIT");
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
auto processed = preprocessor.preprocess(wgsl_glu, defines);
- auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+ auto decisions = std::make_shared<ggml_webgpu_glu_shader_decisions>();
decisions->wg_size = context.max_wg_size;
+ decisions->src_overlap = key.src_overlap;
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
pipeline.context = decisions;
glu_pipelines[key] = pipeline;
return conv2d_pipelines[key];
}
+ // whcn selects the input layout: contiguous WHCN vs contiguous-channels CWHN
+ webgpu_pipeline get_conv2d_dw_pipeline(const ggml_webgpu_shader_lib_context & context, bool whcn) {
+ ggml_webgpu_conv2d_dw_pipeline_key key = {};
+ key.weight_type = context.src0->type;
+ key.input_type = context.src1->type;
+ key.output_type = context.dst->type;
+ key.whcn = whcn;
+
+ auto it = conv2d_dw_pipelines.find(key);
+ if (it != conv2d_dw_pipelines.end()) {
+ return it->second;
+ }
+
+ std::vector<std::string> defines;
+ std::string variant = whcn ? "conv_2d_dw_whcn" : "conv_2d_dw_cwhn";
+
+ auto push_type_defines = [&](const char * prefix, ggml_type type) {
+ std::string s_prefix = prefix;
+ if (type == GGML_TYPE_F32) {
+ defines.push_back(s_prefix + "_F32");
+ } else if (type == GGML_TYPE_F16) {
+ defines.push_back(s_prefix + "_F16");
+ } else {
+ GGML_ABORT("Unsupported type for CONV_2D_DW shader");
+ }
+ };
+
+ push_type_defines("WEIGHT", key.weight_type);
+ push_type_defines("INPUT", key.input_type);
+ push_type_defines("OUTPUT", key.output_type);
+ if (whcn) {
+ defines.push_back("WHCN");
+ }
+ defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+
+ auto processed = preprocessor.preprocess(wgsl_conv2d_dw, defines);
+ auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+ decisions->wg_size = context.max_wg_size;
+ webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
+ pipeline.context = decisions;
+ conv2d_dw_pipelines[key] = pipeline;
+ return conv2d_dw_pipelines[key];
+ }
+
webgpu_pipeline get_im2col_pipeline(const ggml_webgpu_shader_lib_context & context) {
ggml_webgpu_im2col_pipeline_key key = {};
key.input_type = context.src1->type;
return ctx->buffer;
}
+static size_t ggml_webgpu_tensor_misalignment(const ggml_tensor * t, size_t alignment) {
+ size_t offset = ggml_webgpu_tensor_offset(t);
+ return offset & (alignment - 1);
+}
+
static size_t ggml_webgpu_tensor_misalignment(webgpu_context & ctx, const ggml_tensor * t) {
+ return ggml_webgpu_tensor_misalignment(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment);
+}
+
+static size_t ggml_webgpu_tensor_align_offset(const ggml_tensor * t, size_t alignment) {
size_t offset = ggml_webgpu_tensor_offset(t);
- return offset & (ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1);
+ return offset & ~(alignment - 1);
}
static size_t ggml_webgpu_tensor_align_offset(webgpu_context & ctx, const ggml_tensor * t) {
- size_t offset = ggml_webgpu_tensor_offset(t);
- return offset & ~(ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1);
+ return ggml_webgpu_tensor_align_offset(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment);
+}
+
+static size_t ggml_webgpu_tensor_binding_size(const ggml_tensor * t, size_t alignment) {
+ return ROUNDUP_POW2(ggml_nbytes(t) + ggml_webgpu_tensor_misalignment(t, alignment),
+ WEBGPU_STORAGE_BUF_BINDING_MULT);
}
-static size_t ggml_webgpu_tensor_binding_size(webgpu_context & ctx, ggml_tensor * t) {
- return ROUNDUP_POW2(ggml_nbytes(t) + ggml_webgpu_tensor_misalignment(ctx, t), WEBGPU_STORAGE_BUF_BINDING_MULT);
+static size_t ggml_webgpu_tensor_binding_size(webgpu_context & ctx, const ggml_tensor * t) {
+ return ggml_webgpu_tensor_binding_size(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment);
+}
+
+static bool ggml_webgpu_tensor_binding_overlap(const webgpu_global_context & global_ctx,
+ const ggml_tensor * a,
+ const ggml_tensor * b) {
+ if (a->buffer != b->buffer) {
+ return false;
+ }
+
+ const size_t alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment;
+ const size_t a_offset = ggml_webgpu_tensor_align_offset(a, alignment);
+ const size_t b_offset = ggml_webgpu_tensor_align_offset(b, alignment);
+ return a_offset < b_offset + ggml_webgpu_tensor_binding_size(b, alignment) &&
+ b_offset < a_offset + ggml_webgpu_tensor_binding_size(a, alignment);
+}
+
+static bool ggml_webgpu_tensor_binding_overlap_range(const webgpu_global_context & global_ctx,
+ ggml_tensor * tensor,
+ ggml_backend_buffer_t buffer,
+ size_t offset,
+ size_t size) {
+ if (tensor->buffer != buffer) {
+ return false;
+ }
+
+ const size_t alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment;
+ const size_t tensor_offset = ggml_webgpu_tensor_align_offset(tensor, alignment);
+ return tensor_offset < offset + size && offset < tensor_offset + ggml_webgpu_tensor_binding_size(tensor, alignment);
}
struct ggml_webgpu_merged_binding_range {
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}
+// Same param/binding layout as conv_2d; the shader differs
+static webgpu_encoded_op ggml_webgpu_conv_2d_dw(webgpu_context & ctx,
+ ggml_tensor * src0,
+ ggml_tensor * src1,
+ ggml_tensor * dst) {
+ const int32_t s0 = ggml_get_op_params_i32(dst, 0);
+ const int32_t s1 = ggml_get_op_params_i32(dst, 1);
+ const int32_t p0 = ggml_get_op_params_i32(dst, 2);
+ const int32_t p1 = ggml_get_op_params_i32(dst, 3);
+ const int32_t d0 = ggml_get_op_params_i32(dst, 4);
+ const int32_t d1 = ggml_get_op_params_i32(dst, 5);
+
+ // Scalar params matching conv2d_dw.wgsl (weight src0 [KW,KH,1,C], input src1, output dst).
+ std::vector<uint32_t> params = {
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+
+ (uint32_t) ggml_nelements(dst),
+ (uint32_t) dst->ne[2],
+ (uint32_t) dst->ne[3],
+ (uint32_t) dst->ne[0],
+ (uint32_t) dst->ne[1],
+ (uint32_t) src1->ne[0],
+ (uint32_t) src1->ne[1],
+ (uint32_t) src0->ne[0],
+ (uint32_t) src0->ne[1],
+
+ (uint32_t) s0,
+ (uint32_t) s1,
+ (uint32_t) p0,
+ (uint32_t) p1,
+ (uint32_t) d0,
+ (uint32_t) d1,
+ };
+
+ std::vector<wgpu::BindGroupEntry> entries = {
+ ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0),
+ ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1),
+ ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst),
+ };
+
+ ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+ shader_lib_ctx.src0 = src0;
+ shader_lib_ctx.src1 = src1;
+ shader_lib_ctx.dst = dst;
+ shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+
+ // Input layout: contiguous -> WHCN, contiguous-channels -> CWHN
+ const bool whcn = ggml_is_contiguous(src1);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_conv2d_dw_pipeline(shader_lib_ctx, whcn);
+ auto * decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+
+ uint32_t wg_x;
+ uint32_t wg_y;
+ uint32_t total_wg = CEIL_DIV((uint32_t) ggml_nelements(dst), decisions->wg_size);
+ compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
+
+ return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
+}
+
static webgpu_encoded_op ggml_webgpu_im2col(webgpu_context & ctx,
ggml_tensor * src0,
ggml_tensor * src1,
ggml_webgpu_shader_lib_context shader_lib_ctx = {};
shader_lib_ctx.src0 = src0;
shader_lib_ctx.src1 = src1;
+ shader_lib_ctx.src2 = src2;
+ shader_lib_ctx.src3 = src3;
shader_lib_ctx.src4 = src4;
shader_lib_ctx.src5 = src5;
shader_lib_ctx.dst = dst;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
shader_lib_ctx.supports_subgroups = ctx->global_ctx->capabilities.supports_subgroups;
+ bool xbc_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src2) ||
+ ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src4) ||
+ ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src5) ||
+ ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src2, src4) ||
+ ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src2, src5) ||
+ ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src4, src5);
+ bool a_overlap = false;
+ bool ids_overlap = false;
+ ggml_webgpu_merged_binding_range xbc_merged_range = {};
+ if (xbc_overlap) {
+ xbc_merged_range = ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src4, src5 });
+ a_overlap = ggml_webgpu_tensor_binding_overlap_range(ctx->global_ctx, src3, src1->buffer,
+ xbc_merged_range.offset, xbc_merged_range.size);
+ if (a_overlap) {
+ xbc_merged_range = ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src3, src4, src5 });
+ }
+ ids_overlap = ggml_webgpu_tensor_binding_overlap_range(ctx->global_ctx, src6, src1->buffer,
+ xbc_merged_range.offset, xbc_merged_range.size);
+ if (ids_overlap) {
+ xbc_merged_range =
+ a_overlap ? ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src3, src4, src5, src6 }) :
+ ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src4, src5, src6 });
+ }
+ }
- webgpu_pipeline pipeline = ctx->shader_lib->get_ssm_scan_pipeline(shader_lib_ctx);
- auto * decisions = static_cast<ggml_webgpu_ssm_scan_shader_decisions *>(pipeline.context.get());
- const bool xbc_overlap = decisions->xbc_overlap;
+ webgpu_pipeline pipeline =
+ ctx->shader_lib->get_ssm_scan_pipeline(shader_lib_ctx, xbc_overlap, a_overlap, ids_overlap);
+ auto * decisions = static_cast<ggml_webgpu_ssm_scan_shader_decisions *>(pipeline.context.get());
+ xbc_overlap = decisions->xbc_overlap;
+ a_overlap = decisions->a_overlap;
+ ids_overlap = decisions->ids_overlap;
uint32_t offset_x = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type));
+ uint32_t offset_dt = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type));
+ uint32_t offset_A = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src3) / ggml_type_size(src3->type));
uint32_t offset_B = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src4) / ggml_type_size(src4->type));
uint32_t offset_C = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src5) / ggml_type_size(src5->type));
+ uint32_t offset_ids = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src6) / ggml_type_size(src6->type));
size_t xbc_bind_offset = 0;
size_t xbc_bind_size = 0;
if (xbc_overlap) {
- const ggml_webgpu_merged_binding_range merged_range =
- ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src4, src5 });
- xbc_bind_offset = merged_range.offset;
- xbc_bind_size = merged_range.size;
- offset_x = ggml_webgpu_tensor_merged_element_offset(src1, merged_range);
- offset_B = ggml_webgpu_tensor_merged_element_offset(src4, merged_range);
- offset_C = ggml_webgpu_tensor_merged_element_offset(src5, merged_range);
+ xbc_bind_offset = xbc_merged_range.offset;
+ xbc_bind_size = xbc_merged_range.size;
+ offset_x = ggml_webgpu_tensor_merged_element_offset(src1, xbc_merged_range);
+ offset_dt = ggml_webgpu_tensor_merged_element_offset(src2, xbc_merged_range);
+ if (a_overlap) {
+ offset_A = ggml_webgpu_tensor_merged_element_offset(src3, xbc_merged_range);
+ }
+ offset_B = ggml_webgpu_tensor_merged_element_offset(src4, xbc_merged_range);
+ offset_C = ggml_webgpu_tensor_merged_element_offset(src5, xbc_merged_range);
+ if (ids_overlap) {
+ offset_ids = ggml_webgpu_tensor_merged_element_offset(src6, xbc_merged_range);
+ }
}
std::vector<uint32_t> params = {
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
offset_x,
- (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)),
- (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src3) / ggml_type_size(src3->type)),
+ offset_dt,
+ offset_A,
offset_B,
offset_C,
- (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src6) / ggml_type_size(src6->type)),
+ offset_ids,
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
(uint32_t) (src0->nb[1] / ggml_type_size(src0->type)),
if (xbc_overlap) {
entries.push_back(
ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(src1), xbc_bind_offset, xbc_bind_size));
- entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src2));
- entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, src3));
- entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 4, src6));
- entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 5, dst));
+ if (ids_overlap) {
+ if (!a_overlap) {
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src3));
+ }
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, a_overlap ? 2 : 3, dst));
+ } else if (a_overlap) {
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src6));
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, dst));
+ } else {
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src3));
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, src6));
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 4, dst));
+ }
} else {
entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1));
entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src2));
(uint32_t) (idx->ne[1]), (uint32_t) (idx->ne[2])
};
- std::vector<wgpu::BindGroupEntry> entries = {
- ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src),
- ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx),
- ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst),
- };
+ std::vector<wgpu::BindGroupEntry> entries;
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src));
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx));
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst));
if (decisions->i64_idx) {
entries.push_back(ggml_webgpu_make_bind_group_entry(3, ctx->set_rows_dev_error_buf, 0,
op.has_mask = mask != nullptr;
op.has_sinks = sinks != nullptr;
- op.kv_overlap = ggml_webgpu_tensor_overlap(K, V);
+ op.kv_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, K, V);
uint32_t offset_k = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, K) / ggml_type_size(K->type));
uint32_t offset_v = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, V) / ggml_type_size(V->type));
}
static webgpu_encoded_op ggml_webgpu_flash_attn_direct(webgpu_context & ctx, const ggml_webgpu_flash_attn_op & op) {
- webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_pipeline(op.shader_lib_ctx);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_pipeline(op.shader_lib_ctx, op.kv_overlap);
auto * decisions = static_cast<ggml_webgpu_flash_attn_decisions *>(pipeline.context.get());
uint32_t wg_per_head = CEIL_DIV(op.shader_lib_ctx.src0->ne[1], decisions->q_tile);
uint32_t wg_x = wg_per_head * op.shader_lib_ctx.src0->ne[2] * op.shader_lib_ctx.src0->ne[3];
ggml_tensor * sinks,
ggml_tensor * dst,
ggml_webgpu_flash_attn_op op) {
- webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_vec_pipeline(op.shader_lib_ctx);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_vec_pipeline(op.shader_lib_ctx, op.kv_overlap);
auto * decisions = static_cast<ggml_webgpu_flash_attn_vec_decisions *>(pipeline.context.get());
wgpu::Buffer blk_buf = {};
shader_lib_ctx.dst = dst;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
- webgpu_pipeline pipeline = ctx->shader_lib->get_binary_pipeline(shader_lib_ctx);
- auto * decisions = static_cast<ggml_webgpu_binary_shader_decisions *>(pipeline.context.get());
+ const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_binary_pipeline(shader_lib_ctx, src_overlap);
+ auto * decisions = static_cast<ggml_webgpu_binary_shader_decisions *>(pipeline.context.get());
uint32_t ne = (uint32_t) ggml_nelements(dst);
ggml_tensor * dst) {
uint32_t ne = (uint32_t) ggml_nelements(dst);
uint32_t dim = (uint32_t) dst->op_params[0];
+ if (ggml_nbytes(src0) == 0 && ggml_nbytes(src1) == 0) {
+ return {};
+ }
ggml_webgpu_shader_lib_context shader_lib_ctx = {};
shader_lib_ctx.src0 = src0;
shader_lib_ctx.dst = dst;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
- webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx);
- auto * decisions = static_cast<ggml_webgpu_binary_shader_decisions *>(pipeline.context.get());
+ const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1) ||
+ ggml_nbytes(src0) == 0 || ggml_nbytes(src1) == 0;
+ webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx, src_overlap);
+ auto * decisions = static_cast<ggml_webgpu_binary_shader_decisions *>(pipeline.context.get());
uint32_t offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type));
uint32_t offset_src1 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type));
size_t merged_offset = 0;
size_t merged_size = 0;
if (decisions->src_overlap) {
- const ggml_webgpu_merged_binding_range merged_range =
- ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 });
- merged_offset = merged_range.offset;
- merged_size = merged_range.size;
- offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range);
- offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range);
+ if (ggml_nbytes(src0) == 0) {
+ merged_offset = ggml_webgpu_tensor_align_offset(ctx, src1);
+ merged_size = ggml_webgpu_tensor_binding_size(ctx, src1);
+ offset_src0 = 0;
+ offset_src1 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type));
+ } else if (ggml_nbytes(src1) == 0) {
+ merged_offset = ggml_webgpu_tensor_align_offset(ctx, src0);
+ merged_size = ggml_webgpu_tensor_binding_size(ctx, src0);
+ offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type));
+ offset_src1 = 0;
+ } else {
+ const ggml_webgpu_merged_binding_range merged_range =
+ ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 });
+ merged_offset = merged_range.offset;
+ merged_size = merged_range.size;
+ offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range);
+ offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range);
+ }
}
std::vector<uint32_t> params = { ne,
shader_lib_ctx.dst = dst;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
- webgpu_pipeline pipeline = ctx->shader_lib->get_rms_norm_mul_pipeline(shader_lib_ctx);
- auto * decisions = static_cast<ggml_webgpu_rms_norm_mul_shader_decisions *>(pipeline.context.get());
+ const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, rn_src, mul_src);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_rms_norm_mul_pipeline(shader_lib_ctx, src_overlap);
+ auto * decisions = static_cast<ggml_webgpu_rms_norm_mul_shader_decisions *>(pipeline.context.get());
if (decisions->src_overlap) {
const ggml_webgpu_merged_binding_range merged_range =
shader_lib_ctx.dst = dst;
shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
- webgpu_pipeline pipeline = ctx->shader_lib->get_glu_pipeline(shader_lib_ctx);
+ const bool src_overlap = src1 != nullptr && ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1);
+ webgpu_pipeline pipeline = ctx->shader_lib->get_glu_pipeline(shader_lib_ctx, src_overlap);
- auto * decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+ auto * decisions = static_cast<ggml_webgpu_glu_shader_decisions *>(pipeline.context.get());
const int split = (src1 != nullptr);
+ uint32_t offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type));
+ uint32_t offset_src1 =
+ src1 != nullptr ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0;
+ size_t merged_offset = 0;
+ size_t merged_size = 0;
+ if (decisions->src_overlap) {
+ const ggml_webgpu_merged_binding_range merged_range =
+ ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 });
+ merged_offset = merged_range.offset;
+ merged_size = merged_range.size;
+ offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range);
+ offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range);
+ }
+
std::vector<uint32_t> params = {
- (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
- src1 != nullptr ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0,
+ offset_src0,
+ offset_src1,
(uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
(uint32_t) (src0->nb[1] / ggml_type_size(src0->type)),
(uint32_t) (src0->nb[2] / ggml_type_size(src0->type)),
ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 3)), // limit, for swiglu_oai
};
- std::vector<wgpu::BindGroupEntry> entries = {
- ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0),
- };
- uint32_t dst_binding = 1;
- if (split) {
+ std::vector<wgpu::BindGroupEntry> entries;
+ uint32_t dst_binding = 1;
+ if (decisions->src_overlap) {
+ entries.push_back(
+ ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src0), merged_offset, merged_size));
+ } else {
+ entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0));
+ }
+ if (split && !decisions->src_overlap) {
dst_binding = 2;
entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1));
}
return ggml_webgpu_sum_rows(ctx, src0, node);
case GGML_OP_CONV_2D:
return ggml_webgpu_conv_2d(ctx, src0, src1, node);
+ case GGML_OP_CONV_2D_DW:
+ return ggml_webgpu_conv_2d_dw(ctx, src0, src1, node);
case GGML_OP_IM2COL:
return ggml_webgpu_im2col(ctx, src0, src1, node);
case GGML_OP_UPSCALE:
if (!supports_op) {
break;
}
- if (ggml_webgpu_tensor_overlap(src1, src2) && src1->type != src2->type &&
- !ggml_is_quantized(src1->type) && !ggml_is_quantized(src2->type)) {
+ if (ggml_webgpu_tensor_binding_overlap(ctx->webgpu_global_ctx, src1, src2) &&
+ src1->type != src2->type && !ggml_is_quantized(src1->type) && !ggml_is_quantized(src2->type)) {
supports_op = false;
break;
}
(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) &&
(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
break;
+ case GGML_OP_CONV_2D_DW:
+ supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
+ (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) &&
+ (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16) &&
+ (ggml_is_contiguous(src1) || ggml_is_contiguous_channels(src1));
+ break;
case GGML_OP_IM2COL:
supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
--- /dev/null
+#include "common_decls.tmpl"
+enable f16;
+
+// Ported from the Vulkan backend's conv2d_dw.comp. Two variants (based on WHCN)
+// selected by the input (src1) layout: contiguous -> WHCN, else CWHN.
+// weight (src0) is [KW,KH,1,C]; output matches the input layout.
+
+@group(0) @binding(0)
+#if defined(WEIGHT_F32)
+var<storage, read_write> weights: array<f32>;
+#elif defined(WEIGHT_F16)
+var<storage, read_write> weights: array<f16>;
+#endif
+
+@group(0) @binding(1)
+#if defined(INPUT_F32)
+var<storage, read_write> input: array<f32>;
+#elif defined(INPUT_F16)
+var<storage, read_write> input: array<f16>;
+#endif
+
+@group(0) @binding(2)
+#if defined(OUTPUT_F32)
+var<storage, read_write> output: array<f32>;
+#elif defined(OUTPUT_F16)
+var<storage, read_write> output: array<f16>;
+#endif
+
+struct Params {
+ offset_w: u32,
+ offset_i: u32,
+ offset_o: u32,
+
+ ne: u32,
+ channels: u32,
+ batches: u32,
+ dst_w: u32, dst_h: u32,
+ src_w: u32, src_h: u32,
+ knl_w: u32, knl_h: u32,
+
+ stride_x: i32, stride_y: i32,
+ pad_x: i32, pad_y: i32,
+ dilation_x: i32, dilation_y: i32,
+};
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+fn load_weight(idx: u32) -> f32 {
+ #if defined(WEIGHT_F32)
+ return weights[idx];
+ #elif defined(WEIGHT_F16)
+ return f32(weights[idx]);
+ #endif
+}
+fn load_input(idx: u32) -> f32 {
+ #if defined(INPUT_F32)
+ return input[idx];
+ #elif defined(INPUT_F16)
+ return f32(input[idx]);
+ #endif
+}
+fn store_output(idx: u32, val: f32) {
+ #if defined(OUTPUT_F32)
+ output[idx] = val;
+ #elif defined(OUTPUT_F16)
+ output[idx] = f16(val);
+ #endif
+}
+
+#if defined(WHCN)
+// Input/output/kernel contiguous in [W, H, C, N] order (kernel [KW,KH,C]).
+fn conv_2d_dw(idx: u32) -> f32 {
+ let i0 = idx / params.dst_w;
+ let dst_x = idx - i0 * params.dst_w;
+ let i1 = i0 / params.dst_h;
+ let dst_y = i0 - i1 * params.dst_h;
+ let n = i1 / params.channels;
+ let c = i1 - n * params.channels;
+
+ let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w
+ + c * params.src_h * params.src_w;
+ let knl_i = params.offset_w + c * params.knl_h * params.knl_w;
+
+ var sum: f32 = 0.0;
+ for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
+ let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
+ if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
+ for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
+ let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
+ if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
+ let v = load_input(src_i + u32(src_y) * params.src_w + u32(src_x));
+ let k = load_weight(knl_i + ky * params.knl_w + kx);
+ sum += v * k;
+ }
+ }
+ return sum;
+}
+#else
+// Channels contiguous (CWHN): channel is the innermost axis.
+fn conv_2d_dw(idx: u32) -> f32 {
+ let i0 = idx / params.channels;
+ let c = idx - i0 * params.channels;
+ let i1 = i0 / params.dst_w;
+ let dst_x = i0 - i1 * params.dst_w;
+ let n = i1 / params.dst_h;
+ let dst_y = i1 - n * params.dst_h;
+
+ let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w;
+ let src_row = params.src_w * params.channels;
+ let knl_row = params.knl_w * params.channels;
+
+ var sum: f32 = 0.0;
+ for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
+ let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
+ if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
+ for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
+ let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
+ if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
+ let v = load_input(src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c);
+ let k = load_weight(params.offset_w + ky * knl_row + kx * params.channels + c);
+ sum += v * k;
+ }
+ }
+ return sum;
+}
+#endif
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_wg: vec3<u32>
+) {
+ let idx = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
+ if (idx >= params.ne) { return; }
+ store_output(params.offset_o + idx, conv_2d_dw(idx));
+}
@group(0) @binding(0)
var<storage, read_write> src0: array<DataType>;
-#ifdef NO_SPLIT
+#ifdef SRC_OVERLAP
+@group(0) @binding(1)
+var<storage, read_write> dst: array<DataType>;
+
+@group(0) @binding(2)
+var<uniform> params: Params;
+
+fn a_value(base: u32) -> DataType {
+ return src0[base];
+}
+
+fn b_value(base: u32) -> DataType {
+ return src0[base];
+}
+
+#elif defined(NO_SPLIT)
@group(0) @binding(1)
var<storage, read_write> dst: array<DataType>;
@group(0) @binding(0) var<storage, read_write> s_in: array<f32>;
#ifdef XBC_OVERLAP
-@group(0) @binding(1) var<storage, read_write> x_B_C_merged: array<f32>;
-@group(0) @binding(2) var<storage, read_write> dt: array<f32>;
-@group(0) @binding(3) var<storage, read_write> A: array<f32>;
-@group(0) @binding(4) var<storage, read_write> ids: array<i32>;
-@group(0) @binding(5) var<storage, read_write> dst: array<f32>;
-@group(0) @binding(6) var<uniform> params: Params;
+#ifdef IDS_OVERLAP
+@group(0) @binding(1) var<storage, read_write> x_dt_B_C_ids_merged: array<u32>;
+#ifdef A_OVERLAP
+@group(0) @binding(2) var<storage, read_write> dst: array<f32>;
+@group(0) @binding(3) var<uniform> params: Params;
+#else
+@group(0) @binding(2) var<storage, read_write> A: array<f32>;
+@group(0) @binding(3) var<storage, read_write> dst: array<f32>;
+@group(0) @binding(4) var<uniform> params: Params;
+#endif
+#else
+@group(0) @binding(1) var<storage, read_write> x_dt_B_C_merged: array<f32>;
+#ifdef A_OVERLAP
+@group(0) @binding(2) var<storage, read_write> ids: array<i32>;
+@group(0) @binding(3) var<storage, read_write> dst: array<f32>;
+@group(0) @binding(4) var<uniform> params: Params;
+#else
+@group(0) @binding(2) var<storage, read_write> A: array<f32>;
+@group(0) @binding(3) var<storage, read_write> ids: array<i32>;
+@group(0) @binding(4) var<storage, read_write> dst: array<f32>;
+@group(0) @binding(5) var<uniform> params: Params;
+#endif
+#endif
#else
@group(0) @binding(1) var<storage, read_write> x: array<f32>;
@group(0) @binding(2) var<storage, read_write> dt: array<f32>;
return token_in_tile * WG_SIZE;
}
+#ifdef XBC_OVERLAP
+fn read_merged_f32(idx: u32) -> f32 {
+#ifdef IDS_OVERLAP
+ return bitcast<f32>(x_dt_B_C_ids_merged[idx]);
+#else
+ return x_dt_B_C_merged[idx];
+#endif
+}
+#endif
+
+fn read_state_slot(i3: u32) -> u32 {
+#ifdef IDS_OVERLAP
+ return x_dt_B_C_ids_merged[params.offset_ids + i3];
+#else
+ return u32(ids[params.offset_ids + i3]);
+#endif
+}
+
@compute @workgroup_size(WG_SIZE)
fn main(
@builtin(local_invocation_id) local_id: vec3<u32>,
let ir = head_seq % params.n_head;
let i3 = head_seq / params.n_head;
- let state_slot = u32(ids[params.offset_ids + i3]);
+ let state_slot = read_state_slot(i3);
let g = ir / (params.n_head / params.n_group);
let s_idx = params.offset_s + tid + i1 * params.stride_s1 + ir * params.stride_s2 + state_slot * params.stride_s3;
var s_prev = s_in[s_idx];
- let A0 = A[params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1];
+ let a_idx = params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1;
+#ifdef A_OVERLAP
+ let A0 = read_merged_f32(a_idx);
+#else
+ let A0 = A[a_idx];
+#endif
for (var token_base = 0u; token_base < params.n_seq_tokens; token_base += TOKENS_PER_TILE) {
if (tid < TOKENS_PER_TILE) {
if (token < params.n_seq_tokens) {
let x_idx = params.offset_x + i1 + ir * params.stride_x1 + token * params.stride_x2 + i3 * params.stride_x3;
let dt_idx = params.offset_dt + ir + token * params.stride_dt1 + i3 * params.stride_dt2;
+#ifdef XBC_OVERLAP
+ let dt0 = read_merged_f32(dt_idx);
+#else
let dt0 = dt[dt_idx];
+#endif
let dtsp = select(log(1.0 + exp(dt0)), dt0, dt0 > 20.0);
shared_dtsp[tid] = dtsp;
#ifdef XBC_OVERLAP
- shared_x_dt[tid] = x_B_C_merged[x_idx] * dtsp;
+ shared_x_dt[tid] = read_merged_f32(x_idx) * dtsp;
#else
shared_x_dt[tid] = x[x_idx] * dtsp;
#endif
let b_idx = params.offset_B + tid + g * params.stride_B1 + token * params.stride_B2 + i3 * params.stride_B3;
let c_idx = params.offset_C + tid + g * params.stride_C1 + token * params.stride_C2 + i3 * params.stride_C3;
#ifdef XBC_OVERLAP
- let s = s_prev * dA + x_B_C_merged[b_idx] * x_dt;
+ let s = s_prev * dA + read_merged_f32(b_idx) * x_dt;
#else
let s = s_prev * dA + B[b_idx] * x_dt;
#endif
#ifdef USE_SUBGROUP_REDUCTION
#ifdef XBC_OVERLAP
- let subgroup_partial = subgroupAdd(s * x_B_C_merged[c_idx]);
+ let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx));
#else
let subgroup_partial = subgroupAdd(s * C[c_idx]);
#endif
}
#else
#ifdef XBC_OVERLAP
- shared_reduce[reduce_idx] = s * x_B_C_merged[c_idx];
+ shared_reduce[reduce_idx] = s * read_merged_f32(c_idx);
#else
shared_reduce[reduce_idx] = s * C[c_idx];
#endif
"RWKV_WKV7",
"SOLVE_TRI",
"GATED_DELTA_NET",
+ "LIGHTNING_INDEXER",
+ "DSV4_HC_COMB",
+ "DSV4_HC_PRE",
+ "DSV4_HC_POST",
"UNARY",
"GLU",
};
-static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
+static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
"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)",
+ "dsv4_hc_comb(mixes, scale, base)",
+ "dsv4_hc_pre(x, weights)",
+ "dsv4_hc_post(x, residual, post, comb)",
"unary(x)",
"glu(x)",
};
-static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
+static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
return tensor->nb[0] > tensor->nb[1];
}
-static bool ggml_is_contiguous_n(const struct ggml_tensor * tensor, int n) {
+static bool ggml_is_contiguous_m_n(const struct ggml_tensor * tensor, int m, int n) {
size_t next_nb = ggml_type_size(tensor->type);
if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) {
return false;
}
next_nb *= tensor->ne[0]/ggml_blck_size(tensor->type);
- for (int i = 1; i < GGML_MAX_DIMS; i++) {
- if (i > n) {
+ for (int i = 1; i < n; i++) {
+ if (i > m) {
if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) {
return false;
}
}
bool ggml_is_contiguous_0(const struct ggml_tensor * tensor) {
- return ggml_is_contiguous_n(tensor, 0);
+ return ggml_is_contiguous_m_n(tensor, 0, GGML_MAX_DIMS);
}
bool ggml_is_contiguous_1(const struct ggml_tensor * tensor) {
- return ggml_is_contiguous_n(tensor, 1);
+ return ggml_is_contiguous_m_n(tensor, 1, GGML_MAX_DIMS);
}
bool ggml_is_contiguous_2(const struct ggml_tensor * tensor) {
- return ggml_is_contiguous_n(tensor, 2);
+ return ggml_is_contiguous_m_n(tensor, 2, GGML_MAX_DIMS);
+}
+
+bool ggml_is_contiguous_to_1(const struct ggml_tensor * tensor) {
+ return ggml_is_contiguous_m_n(tensor, 0, 1);
+}
+
+bool ggml_is_contiguous_to_2(const struct ggml_tensor * tensor) {
+ return ggml_is_contiguous_m_n(tensor, 0, 2);
+}
+
+bool ggml_is_contiguous_to_3(const struct ggml_tensor * tensor) {
+ return ggml_is_contiguous_m_n(tensor, 0, 3);
}
bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) {
int s0,
int p0,
int d0) {
- struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); // [N, OL, IC * K]
+ struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N, OL, IC * K]
struct ggml_tensor * result =
ggml_mul_mat(ctx,
int d0) {
struct ggml_tensor * new_b = ggml_reshape_4d(ctx, b, b->ne[0], 1, b->ne[1], b->ne[2]);
- struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16);
+ struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16);
struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a);
int p1,
int d0,
int d1) {
- struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type); // [N, OH, OW, IC * KH * KW]
+ struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N, OH, OW, IC * KH * KW]
struct ggml_tensor * result =
ggml_mul_mat(ctx,
int d1, // dilation height
int d2 // dilation depth
) {
- struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type); // [N*OD, OH, OW, IC * KD * KH * KW]
+ struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N*OD, OH, OW, IC * KD * KH * KW]
int64_t OC = a->ne[3] / IC;
int64_t N = b->ne[3] / IC;
struct ggml_tensor * new_a = ggml_reshape_4d(ctx, a, a->ne[0], a->ne[1], 1, a->ne[2] * a->ne[3]);
struct ggml_tensor * im2col = ggml_im2col(ctx, new_a,
ggml_reshape_4d(ctx, b, b->ne[0], b->ne[1], 1, b->ne[2] * b->ne[3]),
- s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW]
+ s0, s1, p0, p1, d0, d1, true, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW]
struct ggml_tensor * new_b = ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], b->ne[2], b->ne[3]); // [N * IC, OH, OW, KH * KW] => [N, IC, OH * OW, KH * KW]
new_a = ggml_reshape_4d(ctx, new_a, (new_a->ne[0] * new_a->ne[1]), new_a->ne[2], new_a->ne[3], 1); // [OC，1, KH, KW] => [1, OC, 1, KH * KW]
return result;
}
+
void ggml_flash_attn_ext_set_prec(
struct ggml_tensor * a,
enum ggml_prec prec) {
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;
+}
+
+// ggml_dsv4_hc_comb
+
+struct ggml_tensor * ggml_dsv4_hc_comb(
+ struct ggml_context * ctx,
+ struct ggml_tensor * mixes,
+ struct ggml_tensor * scale,
+ struct ggml_tensor * base,
+ float eps,
+ int32_t n_iter) {
+ GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+ GGML_ASSERT(scale->type == GGML_TYPE_F32);
+ GGML_ASSERT(base->type == GGML_TYPE_F32);
+ GGML_ASSERT(n_iter > 0);
+
+ const int64_t hc_mix_dim = mixes->ne[0];
+ const int64_t n_tokens = mixes->ne[1];
+
+ int64_t hc = 0;
+ for (int64_t i = 1; i*i + 2*i <= hc_mix_dim; ++i) {
+ if ((2 + i)*i == hc_mix_dim) {
+ hc = i;
+ break;
+ }
+ }
+
+ GGML_ASSERT(hc > 0);
+ GGML_ASSERT(hc == 4);
+ GGML_ASSERT(mixes->ne[2] == 1);
+ GGML_ASSERT(mixes->ne[3] == 1);
+ GGML_ASSERT(scale->ne[0] >= 3);
+ GGML_ASSERT(scale->ne[1] == 1);
+ GGML_ASSERT(scale->ne[2] == 1);
+ GGML_ASSERT(scale->ne[3] == 1);
+ GGML_ASSERT(base->ne[0] == hc_mix_dim);
+ GGML_ASSERT(base->ne[1] == 1);
+ GGML_ASSERT(base->ne[2] == 1);
+ GGML_ASSERT(base->ne[3] == 1);
+
+ struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hc, hc, n_tokens);
+
+ ggml_set_op_params_f32(result, 0, eps);
+ ggml_set_op_params_i32(result, 1, n_iter);
+
+ result->op = GGML_OP_DSV4_HC_COMB;
+ result->src[0] = mixes;
+ result->src[1] = scale;
+ result->src[2] = base;
+
+ return result;
+}
+
+// ggml_dsv4_hc_pre
+
+struct ggml_tensor * ggml_dsv4_hc_pre(
+ struct ggml_context * ctx,
+ struct ggml_tensor * x,
+ struct ggml_tensor * weights) {
+ GGML_ASSERT(x->type == GGML_TYPE_F32);
+ GGML_ASSERT(weights->type == GGML_TYPE_F32);
+
+ const int64_t n_embd = x->ne[0];
+ const int64_t hc = x->ne[1];
+ const int64_t n_tokens = x->ne[2];
+
+ GGML_ASSERT(hc > 0);
+ GGML_ASSERT(x->ne[3] == 1);
+ GGML_ASSERT(weights->ne[0] == hc);
+ GGML_ASSERT(weights->ne[1] == n_tokens);
+ GGML_ASSERT(weights->ne[2] == 1);
+ GGML_ASSERT(weights->ne[3] == 1);
+
+ struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens);
+
+ result->op = GGML_OP_DSV4_HC_PRE;
+ result->src[0] = x;
+ result->src[1] = weights;
+
+ return result;
+}
+
+// ggml_dsv4_hc_post
+
+struct ggml_tensor * ggml_dsv4_hc_post(
+ struct ggml_context * ctx,
+ struct ggml_tensor * x,
+ struct ggml_tensor * residual,
+ struct ggml_tensor * post,
+ struct ggml_tensor * comb) {
+ GGML_ASSERT(x->type == GGML_TYPE_F32);
+ GGML_ASSERT(residual->type == GGML_TYPE_F32);
+ GGML_ASSERT(post->type == GGML_TYPE_F32);
+ GGML_ASSERT(comb->type == GGML_TYPE_F32);
+
+ const int64_t n_embd = x->ne[0];
+ const int64_t n_tokens = x->ne[1];
+ const int64_t hc = residual->ne[1];
+
+ GGML_ASSERT(hc > 0);
+ GGML_ASSERT(x->ne[2] == 1);
+ GGML_ASSERT(x->ne[3] == 1);
+
+ GGML_ASSERT(residual->ne[0] == n_embd);
+ GGML_ASSERT(residual->ne[2] == n_tokens);
+ GGML_ASSERT(residual->ne[3] == 1);
+
+ GGML_ASSERT(post->ne[0] == hc);
+ GGML_ASSERT(post->ne[1] == n_tokens);
+ GGML_ASSERT(post->ne[2] == 1);
+ GGML_ASSERT(post->ne[3] == 1);
+
+ GGML_ASSERT(comb->ne[0] == hc);
+ GGML_ASSERT(comb->ne[1] == hc);
+ GGML_ASSERT(comb->ne[2] == n_tokens);
+ GGML_ASSERT(comb->ne[3] == 1);
+
+ struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, hc, n_tokens);
+
+ result->op = GGML_OP_DSV4_HC_POST;
+ result->src[0] = x;
+ result->src[1] = residual;
+ result->src[2] = post;
+ result->src[3] = comb;
+
+ return result;
+}
+
////////////////////////////////////////////////////////////////////////////////
struct ggml_hash_set ggml_hash_set_new(size_t size) {
}
void ggml_set_output(struct ggml_tensor * tensor) {
- tensor->flags |= GGML_TENSOR_FLAG_OUTPUT;
+ for (struct ggml_tensor * cur = tensor; cur != NULL; cur = cur->view_src) {
+ cur->flags |= GGML_TENSOR_FLAG_OUTPUT;
+ }
}
void ggml_set_param(struct ggml_tensor * tensor) {
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);
return ctx->info[tensor_id].t.name;
}
+const int64_t * gguf_get_tensor_ne(const struct gguf_context * ctx, int64_t tensor_id) {
+ GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
+ return ctx->info[tensor_id].t.ne;
+}
+
enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) {
GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx));
return ctx->info[tensor_id].t.type;
struct gguf_writer_base {
size_t written_bytes {0u};
- ~gguf_writer_base(void) = default;
+ virtual ~gguf_writer_base(void) = default;
// we bet on devirtualization
virtual void write(int8_t val) = 0;
fout.write(struct.pack("i", n_mels))
fout.write(struct.pack("i", n_freqs))
- f32(n_mels, n_freqs).tofile(fout)
+ # Mel filter must be non-negative.
+ np.abs(f32(n_mels, n_freqs)).tofile(fout)
fout.write(struct.pack("i", n_fft))
f32(n_fft).tofile(fout)
-524f974bb21a1013408f76d71c15732482c0c3fe
+06ca97616793248fadb410ea8d69c7511b2005e4
overview / pkg / ggml / sources / whisper.cpp / commitdiff