include(ExternalProject)
option(GGML_HEXAGON_HTP_DEBUG "ggml-hexagon: enable HTP debug output" OFF)
-option(GGML_HEXAGON_FA_EXP2_HF "ggml-hexagon: use FP16 exp2 polynomial in FA softmax instead of F32 exp round-trip" OFF)
set(GGML_HEXAGON_HTP_CERT "$ENV{HEXAGON_HTP_CERT}" CACHE PATH "ggml-hexagon: enable HTP library signing using certificate")
add_library(htp_iface OBJECT
#include "htp-opnode.h"
#include "htp-ops.h"
#include "htp/matmul-ops.h"
+#include "htp/flash-attn-ops.h"
#include "htp_iface.h"
#include "htp-drv.h"
static int opt_hostbuf = 1; // hostbuf ON by default
static int opt_mm_select = 3; // 3 = HMX -> Tiled -> Flat -> CPU, 2 = Tiled -> Flat -> CPU, 1 = Flat -> CPU
+static int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
// Default PMU events, if profiling with PMU (mode=2) is enabled
// See https://docs.qualcomm.com/doc/80-N2040-60/topic/pmu-events.html
case HTP_TRACE_EVT_HVX_W_DEQUANT: return "HVX_W_DEQUANT";
case HTP_TRACE_EVT_HVX_W_PREP: return "HVX_W_PREP";
case HTP_TRACE_EVT_HVX_O_PROC: return "HVX_O_PROC";
+ case HTP_TRACE_EVT_HVX_FA_QK: return "HVX_QK_FA";
+ case HTP_TRACE_EVT_HVX_FA_SFM: return "HVX_SFM_FA";
+ case HTP_TRACE_EVT_HVX_FA_Q_PREP: return "HVX_Q_PREP";
+ 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";
default: return "UNKNOWN";
}
// ** backend interface
+static bool ggml_hexagon_flash_attn_is_hmx_eligible(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * q,
+ const struct ggml_tensor * k,
+ const struct ggml_tensor * v,
+ const struct ggml_tensor * sinks
+) {
+ if (sess->n_hmx == 0) {
+ return false;
+ }
+
+ if (opt_fa_select < 2) {
+ return false;
+ }
+
+ if (k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) {
+ return false;
+ }
+
+ const uint32_t DK = q->ne[0];
+ const uint32_t DV = v->ne[0];
+
+ if (DK % 64 != 0 || DV % 64 != 0) {
+ return false;
+ }
+
+ // Fall back to HVX for small token counts if head dimension is small (DK <= 128)
+ const uint32_t neq1 = q->ne[1];
+ if (DK <= 128 && neq1 < 5) {
+ return false;
+ }
+
+ return true;
+}
+
+static bool ggml_hexagon_precompute_flash_attn_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * op,
+ struct htp_fa_kernel_params * kparams
+) {
+ if (opt_fa_select < 1) {
+ return false;
+ }
+
+ memset(kparams, 0, sizeof(*kparams));
+
+ const struct ggml_tensor * q = op->src[0];
+ const struct ggml_tensor * k = op->src[1];
+ const struct ggml_tensor * v = op->src[2];
+ const struct ggml_tensor * mask = op->src[3];
+ const struct ggml_tensor * dst = op;
+
+ const uint32_t neq0 = q->ne[0]; // head_dim (DK)
+ const uint32_t neq1 = q->ne[1]; // n_tokens
+ const uint32_t neq2 = q->ne[2]; // n_heads
+
+ const uint32_t nek1 = k->ne[1]; // kv_len
+
+ const uint32_t nev0 = v->ne[0]; // head_dim (DV)
+
+ const uint32_t DK = neq0;
+ const uint32_t DV = nev0;
+
+ const uint32_t n_kv_heads = k->ne[2];
+ const uint32_t G = neq2 / n_kv_heads;
+
+ float scale = 1.0f;
+ float max_bias = 0.0f;
+ float logit_softcap = 0.0f;
+ memcpy(&scale, &op->op_params[0], sizeof(float));
+ memcpy(&max_bias, &op->op_params[1], sizeof(float));
+ memcpy(&logit_softcap, &op->op_params[2], sizeof(float));
+
+ if (logit_softcap != 0.0f) {
+ scale /= logit_softcap;
+ }
+
+ kparams->scale = scale;
+ kparams->max_bias = max_bias;
+ kparams->logit_softcap = logit_softcap;
+
+ kparams->is_q_fp32 = (q->type == GGML_TYPE_F32) ? 1 : 0;
+ kparams->is_dst_fp32 = (dst->type == GGML_TYPE_F32) ? 1 : 0;
+ kparams->G = G;
+
+ const uint32_t n_head = q->ne[2];
+ kparams->n_head_log2 = 1u << (uint32_t) std::floor(std::log2(n_head));
+ kparams->m0 = std::pow(2.0f, -(max_bias) / kparams->n_head_log2);
+ kparams->m1 = std::pow(2.0f, -(max_bias / 2.0f) / kparams->n_head_log2);
+
+ // Check HMX eligibility
+ 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);
+ if (ret == 0) {
+ kparams->kernel_type = HTP_FA_KERNEL_HMX;
+ kparams->Br = Br;
+ kparams->Bc = Bc;
+ kparams->n_kv_blocks = (nek1 + Bc - 1) / Bc;
+ kparams->n_threads = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? sess->n_threads : 1;
+
+ 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);
+
+ 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
+
+ const size_t m_line_bytes = hex_align_up(Bc * sizeof(uint16_t), 128);
+ kparams->u.hmx.mask_buf_row_stride = m_line_bytes / sizeof(uint16_t);
+ kparams->u.hmx.mask_broadcast = (mask != nullptr && mask->ne[2] == 1) ? 1 : 0;
+ kparams->u.hmx.div_G = init_fastdiv_values(G);
+ if (mask) {
+ kparams->src3_div2 = init_fastdiv_values(mask->ne[2]);
+ kparams->src3_div3 = init_fastdiv_values(mask->ne[3]);
+ }
+
+ kparams->qrows = 0;
+ kparams->qrows_per_thread = 0;
+ return true;
+ }
+ }
+
+ // Fallback to HVX
+ kparams->kernel_type = HTP_FA_KERNEL_HVX;
+ kparams->Br = 1;
+ kparams->Bc = 64; // FLASH_ATTN_BLOCK_SIZE
+ kparams->n_kv_blocks = (k->ne[1] + 64 - 1) / 64;
+ kparams->n_threads = sess->n_threads;
+
+ const size_t size_q_row_padded = hex_round_up(q->ne[0] * (kparams->is_q_fp32 ? 4 : 2), 128);
+ const size_t size_k_row_padded = hex_round_up(k->ne[0] * 2, 128);
+ const size_t size_v_row_padded = hex_round_up(v->ne[0] * 2, 128);
+
+ kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, kparams->is_q_fp32 != 0, mask != nullptr, sess->n_threads);
+
+ kparams->u.hvx.size_q_row_padded = size_q_row_padded;
+ kparams->u.hvx.size_k_row_padded = size_k_row_padded;
+ kparams->u.hvx.size_v_row_padded = size_v_row_padded;
+ kparams->u.hvx.src0_div21 = init_fastdiv_values(q->ne[2] * q->ne[1]);
+ kparams->u.hvx.src0_div1 = init_fastdiv_values(q->ne[1]);
+ kparams->u.hvx.broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]);
+ kparams->u.hvx.broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]);
+ kparams->u.hvx.broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]);
+ kparams->u.hvx.broadcast_rv3 = init_fastdiv_values(q->ne[3]/v->ne[3]);
+ if (mask) {
+ kparams->src3_div2 = init_fastdiv_values(mask->ne[2]);
+ kparams->src3_div3 = init_fastdiv_values(mask->ne[3]);
+ }
+
+ kparams->qrows = q->ne[1] * q->ne[2] * q->ne[3];
+ kparams->qrows_per_thread = (kparams->qrows + sess->n_threads - 1) / sess->n_threads;
+
+ return true;
+}
static bool ggml_hexagon_supported_flash_attn_ext(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const struct ggml_tensor * src0 = op->src[0];
return false;
}
+ struct htp_fa_kernel_params kparams;
+ if (!ggml_hexagon_precompute_flash_attn_params(sess, op, &kparams)) {
+ return false;
+ }
+
+ if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
+ HEX_VERBOSE("ggml-hex: skip flash_attn_ext because VTCM needed (%d) > budget (%zu)\n",
+ kparams.vtcm_size, sess->vtcm_size);
+ return false;
+ }
+
return true;
}
kparams->kernel_type = (src1_nrows < (int) sess->n_threads) ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW;
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);
- size_t vtcm_src0_size = 0, vtcm_src1_size = 0;
+ size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
uint32_t best_n_prefetch = 2;
size_t total_size = 0;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
total_size = htp_mm_hvx_id_get_vtcm_sizes(
wtype, ne10, src1_nrows, sess->n_threads, src0->nb[1], d,
- &vtcm_src0_size, &vtcm_src1_size
+ &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
);
if (total_size <= vtcm_budget) {
best_n_prefetch = d;
if (best_n_prefetch == 2 && total_size > vtcm_budget) {
total_size = htp_mm_hvx_id_get_vtcm_sizes(
wtype, ne10, src1_nrows, sess->n_threads, src0->nb[1], 2,
- &vtcm_src0_size, &vtcm_src1_size
+ &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
);
}
kparams->n_prefetch = best_n_prefetch;
kparams->vtcm_size = total_size;
kparams->vtcm_src0_size = vtcm_src0_size;
kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = 0;
+ kparams->vtcm_dst_size = vtcm_dst_size;
} else {
bool try_tiled = (k_align && opt_mm_select >= 2);
if (try_tiled) {
size_t src3_sz_per_thread = 0;
uint32_t best_n_prefetch = 16;
+ size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * sess->n_threads;
+
if (is_repack) {
uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t src1_row_size_padded = hex_round_up(src1_row_size, QK_Q8_0_TILED * sizeof(float));
size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
size_t src1_sz = src1_sz_per_thread;
best_n_prefetch = 2;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
size_t repacked_vtcm_size = hex_round_up(d * tile_row_size, 128);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
size_t src0_sz = repacked_vtcm_size * sess->n_threads;
size_t src2_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
size_t src3_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz;
+ size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
if (tiled_vtcm_size <= sess->vtcm_size) {
best_n_prefetch = d;
}
if (best_n_prefetch == 2 && src0_sz_per_thread == 0) {
size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
src0_sz_per_thread = repacked_vtcm_size;
src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
src3_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
size_t src2_sz = src2_sz_per_thread * sess->n_threads;
size_t src3_sz = src3_sz_per_thread * sess->n_threads;
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz;
+ size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
bool try_tiled = (opt_mm_select >= 2);
if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
kparams->vtcm_src1_size = src1_sz;
kparams->vtcm_src2_size = src2_sz;
kparams->vtcm_src3_size = src3_sz;
+ kparams->vtcm_dst_size = quant_scratch_size;
kparams->vtcm_size = tiled_vtcm_size;
kparams->n_prefetch = best_n_prefetch;
} else {
kparams->vtcm_src1_size = flat_src1_sz;
kparams->vtcm_src2_size = src2_sz;
kparams->vtcm_src3_size = src3_sz;
- kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + src3_sz;
+ kparams->vtcm_dst_size = quant_scratch_size;
+ kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + src3_sz + quant_scratch_size;
kparams->n_prefetch = best_n_prefetch;
}
}
size_t src2_sz_per_thread = 0;
uint32_t best_n_prefetch = 16;
+ size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * sess->n_threads;
+
if (is_repack) {
uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t src1_row_size_padded = hex_round_up(src1_row_size, QK_Q8_0_TILED * sizeof(float));
size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
size_t src1_sz = src1_sz_per_thread;
best_n_prefetch = 2;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
size_t repacked_vtcm_size = hex_round_up(d * tile_row_size, 128);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
size_t src0_sz = repacked_vtcm_size * sess->n_threads;
size_t src2_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz;
+ size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
if (tiled_vtcm_size <= sess->vtcm_size) {
best_n_prefetch = d;
}
if (best_n_prefetch == 2 && src0_sz_per_thread == 0) {
size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
src0_sz_per_thread = repacked_vtcm_size;
src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
}
size_t src1_sz = src1_sz_per_thread;
size_t src2_sz = src2_sz_per_thread * sess->n_threads;
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz;
+ size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
bool try_tiled = (opt_mm_select >= 2);
if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
kparams->vtcm_src0_size = src0_sz;
kparams->vtcm_src1_size = src1_sz;
kparams->vtcm_src2_size = src2_sz;
+ kparams->vtcm_dst_size = quant_scratch_size;
kparams->vtcm_size = tiled_vtcm_size;
kparams->n_prefetch = best_n_prefetch;
} else {
kparams->vtcm_src0_size = src0_sz;
kparams->vtcm_src1_size = flat_src1_sz;
kparams->vtcm_src2_size = src2_sz;
- kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz;
+ kparams->vtcm_dst_size = quant_scratch_size;
+ kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + quant_scratch_size;
kparams->n_prefetch = best_n_prefetch;
}
}
return !ggml_op_is_empty(node->op) && !ggml_is_empty(node) && (node->flags & GGML_TENSOR_FLAG_COMPUTE);
}
-static bool is_hmx_eligible(const ggml_tensor * t) {
+static bool mm_is_hmx_eligible(const ggml_tensor * t) {
if (opt_nhmx == 0) { return false; }
const ggml_tensor * src0 = t->src[0];
static bool is_mergeable_mul_mat(const ggml_tensor * t) {
if (!t || t->op != GGML_OP_MUL_MAT) return false;
if (t->src[1]->type != GGML_TYPE_F32) return false;
- return ggml_is_quantized(t->src[0]->type) && !is_hmx_eligible(t);
+ return ggml_is_quantized(t->src[0]->type) && !mm_is_hmx_eligible(t);
}
static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) {
}
}
+ 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) {
+ 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) {
+ 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 {
+ HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n",
+ kparams.vtcm_size, sess->vtcm_size);
+ }
+ }
+ }
+ }
+
return false;
}
node.node->src[0], node.node->src[1], node.node,
(struct htp_mm_kernel_params *)node.kernel_params
);
+ } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+ ggml_hexagon_precompute_flash_attn_params(sess,
+ node.node,
+ (struct htp_fa_kernel_params *)node.kernel_params
+ );
}
computed_nodes.push_back(std::move(node));
}
const char * str_use_hmx = getenv("GGML_HEXAGON_USE_HMX");
const char * str_nhmx = getenv("GGML_HEXAGON_NHMX");
const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
+ const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
const char * str_ndev = getenv("GGML_HEXAGON_NDEV");
const char * str_arch = getenv("GGML_HEXAGON_ARCH");
const char * str_vmem = getenv("GGML_HEXAGON_VMEM");
opt_nhvx = str_nhvx ? strtoul(str_nhvx, NULL, 0) : opt_nhvx;
opt_nhmx = str_nhmx ? atoi(str_nhmx) : (str_use_hmx ? atoi(str_use_hmx) : opt_nhmx);
opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select;
+ opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select;
opt_ndev = str_ndev ? strtoul(str_ndev, NULL, 0) : opt_ndev;
opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) : opt_hostbuf;
opt_mbuf = str_mbuf ? strtoul(str_mbuf, NULL, 0) * MiB : opt_mbuf;
#include <stdio.h>
#include "htp-ops.h"
#include "htp/matmul-ops.h"
+#include "htp/flash-attn-ops.h"
struct htp_opnode {
ggml_tensor * node = nullptr;
}
void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
- node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN) {
+ node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN ||
+ node.opcode == HTP_OP_MUL_MAT_ADD) {
const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
const char * path = "unknown";
int32_t type = kparams->kernel_type;
path = "hvx-flat";
}
snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+ } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+ const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params;
+ const char * path = "unknown";
+ int32_t type = kparams->kernel_type;
+ if (type == HTP_FA_KERNEL_HMX) {
+ path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
+ } else if (type == HTP_FA_KERNEL_HVX) {
+ path = "hvx";
+ }
+ snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
} else {
snprintf(str, max_size, "----");
}
worker-pool.c
hex-dma.c
hmx-queue.c
- flash-attn-ops.c
- hmx-flash-attn-ops.c
- matmul-ops.c
binary-ops.c
unary-ops.c
sum-rows-ops.c
solve-tri-ops.c
gated-delta-net-ops.c
pad-ops.c
+ matmul-ops.c
+ flash-attn-ops.c
)
target_compile_definitions(${HTP_LIB} PRIVATE
$<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,HTP_DEBUG=1,NDEBUG=1>
$<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,FARF_HIGH=1,>)
-if (GGML_HEXAGON_FA_EXP2_HF)
- message(STATUS "ggml-htp: HMX_FA_USE_EXP2_HF=1 (use FP16 exp2 polynomial in FA softmax)")
- target_compile_definitions(${HTP_LIB} PRIVATE HMX_FA_USE_EXP2_HF=1)
-endif()
build_idl(htp_iface.idl ${HTP_LIB})
#pragma clang diagnostic ignored "-Wunused-but-set-variable"
#include <assert.h>
+#include <HAP_compute_res.h>
#include <HAP_farf.h>
#include <HAP_perf.h>
#include <math.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
#include <string.h>
#include "hex-dma.h"
+#include "hex-fastdiv.h"
+#include "hex-profile.h"
+#include "hmx-queue.h"
+#include "hmx-utils.h"
#include "hvx-utils.h"
#include "hvx-dump.h"
+#include "hvx-copy.h"
+#include "hvx-reduce.h"
#include "hvx-flash-attn.h"
+#include "vtcm-utils.h"
+#include "worker-pool.h"
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "htp-ctx.h"
#include "htp-ops.h"
-#include "htp-ops.h"
-int hmx_flash_attn_ext(struct htp_ops_context * octx);
+#include "flash-attn-ops.h"
+#include "hvx-fa-kernels.h"
+#include "hmx-fa-kernels.h"
// Must be multiple of 32
#define FLASH_ATTN_BLOCK_SIZE (32 * 2)
-#if __HVX_ARCH__ < 79
-#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
-#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b))
-#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
-#else
-#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
-#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b)
-#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
-#endif
-
-// This is a bit of a hack because the compiler is strugling to properly inline
-// the default hvx_vec_f32_to_f16 with output into the local array.
-static __attribute__((noinline)) void hvx_vec_f32_to_f16_a(void *ptr, HVX_Vector v0, HVX_Vector v1)
-{
- *(HVX_Vector *) ptr = hvx_vec_f32_to_f16(v0, v1);
-}
-
-// Dot product of two F16 vectors, accumulating to float
-static inline void hvx_dot_f16_f16_aa(float * restrict r, const void * restrict x, const void * restrict y, unsigned int n, float s) {
- const HVX_Vector * restrict vx = (const HVX_Vector * restrict) x; // fp16
- const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
-
- uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
- uint32_t nloe = n % VLEN_FP16; // leftover elements
-
- HVX_VectorPair rsum_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
-
- uint32_t i = 0;
-
- #pragma unroll(4)
- for (i = 0; i < nvec; i++) {
- rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, vx[i], vy[i]);
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
- HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
- HVX_Vector x_hf = Q6_V_vand_QV(bmask, vx[i]);
-
- rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
- }
-
- HVX_Vector rsum = HVX_OP_ADD_F32(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p));
- rsum = HVX_OP_MUL_F32(hvx_vec_splat_f32(s), hvx_vec_reduce_sum_f32(rsum));
- hvx_vec_store_u(r, 4, rsum);
-}
-
-static inline HVX_Vector hvx_dot_f16_f16_aa_rx4(const void * restrict y,
- const uint8_t * restrict x,
- const size_t stride_x,
- const size_t nvec,
- const size_t nloe) {
- const HVX_Vector * restrict vx0 = (const HVX_Vector * restrict) x; // fp16
- const HVX_Vector * restrict vx1 = (const HVX_Vector * restrict) (x + stride_x); // fp16
- const HVX_Vector * restrict vx2 = (const HVX_Vector * restrict) (x + stride_x * 2); // fp16
- const HVX_Vector * restrict vx3 = (const HVX_Vector * restrict) (x + stride_x * 3); // fp16
- const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
-
- HVX_VectorPair rsum0_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
- HVX_VectorPair rsum1_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
- HVX_VectorPair rsum2_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
- HVX_VectorPair rsum3_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
-
- uint32_t i = 0;
-
- for (i = 0; i < nvec; i++) {
- HVX_Vector y_hf = vy[i];
- HVX_Vector x0_hf = vx0[i];
- HVX_Vector x1_hf = vx1[i];
- HVX_Vector x2_hf = vx2[i];
- HVX_Vector x3_hf = vx3[i];
-
- rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
- rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
- rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
- rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
- }
-
- if (nloe) {
- // Load x (fp16) and zero-out unused elements
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
- HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
- HVX_Vector x0_hf = Q6_V_vand_QV(bmask, vx0[i]);
- HVX_Vector x1_hf = Q6_V_vand_QV(bmask, vx1[i]);
- HVX_Vector x2_hf = Q6_V_vand_QV(bmask, vx2[i]);
- HVX_Vector x3_hf = Q6_V_vand_QV(bmask, vx3[i]);
-
- rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
- rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
- rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
- rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
- }
-
- HVX_Vector rsum0 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p));
- HVX_Vector rsum1 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p));
- HVX_Vector rsum2 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum2_p), Q6_V_hi_W(rsum2_p));
- HVX_Vector rsum3 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum3_p), Q6_V_hi_W(rsum3_p));
-
- HVX_Vector_x4 rsum0123 = { .v = { rsum0, rsum1, rsum2, rsum3 } };
- return hvx_vec_reduce_sum_f32x4(rsum0123);
-}
-
-static inline HVX_Vector hvx_dot_f16_f16_aa_rx32(const void * restrict y,
- const uint8_t * restrict x,
- const size_t stride_x,
- const size_t n,
- float s) {
-
- const size_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
- const size_t nloe = n % VLEN_FP16; // leftover elements
-
- HVX_Vector sums = Q6_V_vzero();
- const size_t stride_x_4 = stride_x * 4;
- for (uint32_t j = 0; j < VLEN_FP32; j += 4) {
- HVX_Vector sums_x4 = hvx_dot_f16_f16_aa_rx4(y, x, stride_x, nvec, nloe);
- HVX_VectorPred pred = Q6_Q_vsetq_R(j * SIZEOF_FP32);
- sums = Q6_V_vmux_QVV(pred, sums, sums_x4);
- x += stride_x_4;
- }
-
- return HVX_OP_MUL_F32(hvx_vec_splat_f32(s), sums);
-}
-
-// MAD: y (F32) += x (F16) * s (F16)
-static inline void hvx_mad_f32_f16_aa(float * restrict y, const void * restrict x, const __fp16 * restrict s, int 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
-
- HVX_Vector S0 = hvx_vec_splat_f16(*s);
-
- 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);
- }
- }
-}
-
-// MAD: y (F32) += x0 (F16) * s0 (F16) + x1 (F16) * s1 (F16)
-static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restrict x0, const void * restrict x1,
- const __fp16 * restrict s0, const __fp16 * restrict s1, int 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
-
- HVX_Vector S0 = hvx_vec_splat_f16(*s0);
- HVX_Vector S1 = hvx_vec_splat_f16(*s1);
-
- 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);
- }
- }
-}
-
struct htp_fa_context {
const struct htp_ops_context * octx;
float scale;
float max_bias;
- float logit_softcap;
+ __fp16 logit_softcap;
uint32_t n_head_log2;
float m0;
float m1;
- float slopes[512];
+ __fp16 slopes[512];
uint32_t n_blocks;
bool is_q_fp32;
- uint64_t t_start;
-};
-
-static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, HVX_Vector vs) {
- assert((size_t) dst % 128 == 0);
- assert((size_t) src % 128 == 0);
+ size_t size_q_block;
+ size_t size_vkq_acc;
- const HVX_Vector * restrict vsrc = (const HVX_Vector * restrict) src;
- HVX_Vector * restrict vdst = (HVX_Vector * restrict) dst;
+ uint8_t * spad_q;
+ uint8_t * spad_k;
+ uint8_t * spad_v;
+ uint8_t * spad_m;
+ uint8_t * spad_a;
- const uint32_t nvec = n / VLEN_FP32;
- const uint32_t nloe = n % VLEN_FP32;
+ uint64_t t_start;
+};
- uint32_t i = 0;
- #pragma unroll(4)
- for (; i < nvec; ++i) {
- vdst[i] = HVX_OP_MUL_F32(vsrc[i], vs);
- }
- if (nloe) {
- hvx_vec_store_a(&vdst[i], nloe * sizeof(float), HVX_OP_MUL_F32(vsrc[i], vs));
- }
-}
+struct hmx_fa_context {
+ const struct htp_ops_context * octx;
+ const struct htp_tensor * sinks; // attention sinks (src[4]), NULL if absent
+ bool pipeline; // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2
+ uint32_t n_threads;
+
+ // Op parameters
+ __fp16 scale;
+ float max_bias;
+ __fp16 logit_softcap;
+ uint32_t n_head_log2;
+ float m0, m1;
+
+ // Dimensions
+ uint32_t DK, DV;
+ uint32_t n_kv; // kv_len
+ uint32_t n_kv_heads; // number of KV heads
+ uint32_t n_heads; // number of Q heads
+ uint32_t G; // GQA factor = n_heads / n_kv_heads
+ struct fastdiv_values div_G;
+ struct fastdiv_values src3_div2;
+ struct fastdiv_values src3_div3;
+ uint32_t n_kv_blocks;
+ uint32_t neq1; // Q token count
+
+ // Types
+ bool is_q_fp32;
+ bool is_dst_fp32;
+
+ // Dynamic block sizes
+ uint32_t Br; // Q tokens per block (before GQA expansion)
+ uint32_t Bc;
+ uint32_t g_br; // hex_align_up(G * Br, 32) - actual tile row dim
+
+ // VTCM buffers (allocated by vtcm_seq_alloc)
+ __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_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_d_tiles; // Diagonal rescale [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]
+ HVX_Vector * vtcm_p_rowsum; // Softmax intermediate [g_br]
+ HVX_Vector * vtcm_row_bufs; // Per-thread softmax row scratch [n_threads][2][Bc/64]
+ uint8_t * vtcm_hmx_scales_id; // HMX output scales (identity)
+ uint8_t * vtcm_hmx_scales_qk; // HMX output scales (qk_scale)
+ __fp16 * vtcm_mask_buf; // VTCM mask buffer [Br * m_line], DMA'd per KV block
+ __fp16 * vtcm_slopes; // ALiBi slopes [g_br]
+ size_t row_buf_stride; // HVX vectors per row buffer (Bc/64)
+ size_t mask_buf_row_stride; // elements (__fp16) per row in mask buffer
+ bool mask_broadcast; // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
+ dma_cache m_cache;
+};
static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * data) {
struct htp_fa_context * factx = (struct htp_fa_context *) data;
if (ir0 >= ir1) return;
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0);
dma_queue * dma = octx->ctx->dma[ith];
const size_t size_v_row = DV * sizeof(__fp16);
// Scratchpad buffers for Q, K, V, Mask, and VKQ32 accumulator
- uint8_t * spad_q = octx->src0_spad.data + octx->src0_spad.size_per_thread * ith;
- uint8_t * spad_k = octx->src1_spad.data + octx->src1_spad.size_per_thread * ith;
- uint8_t * spad_v = octx->src2_spad.data + octx->src2_spad.size_per_thread * ith;
- uint8_t * spad_m = octx->src3_spad.data + octx->src3_spad.size_per_thread * ith;
- uint8_t * spad_a = octx->dst_spad.data + octx->dst_spad.size_per_thread * ith;
-
- const HVX_Vector logit_cap = hvx_vec_splat_f32(factx->logit_softcap);
+ uint8_t * spad_q = factx->spad_q + factx->size_q_block * ith;
+ uint8_t * spad_k = factx->spad_k + factx->size_k_block * 2 * ith;
+ uint8_t * spad_v = factx->spad_v + factx->size_v_block * 2 * ith;
+ uint8_t * spad_m = factx->spad_m + (mask ? factx->size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * ith;
+ uint8_t * spad_a = factx->spad_a + factx->size_vkq_acc * ith;
dma_cache m_cache;
- dma_cache_init(&m_cache, spad_m, factx->size_m_block, DMA_CACHE_MAX_SIZE);
+ dma_cache_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
for (uint32_t ir = ir0; ir < ir1; ++ir) {
const uint32_t iq3 = fastdiv(ir, &factx->src0_div21);
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);
- // FARF(HIGH, "fa %u: prefetch Q: ir %u iq1 %u iq2 %u iq3 %u q_row_ptr %p size %u : usec %u", ith, ir, iq1, iq2, iq3, q_row_ptr, size_q_row,
- // (unsigned)HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - factx->t_start));
-
const __fp16 * mp_base = NULL;
if (mask) {
const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
// 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);
}
-
- // FARF(HIGH, "fa %u: prefetch KVM: ir %u ib %u iq1 %u iq2 %u iq3 %u : size_k_row %u size_v_row %u bs %u: usec %u",
- // ith, ir, ib, iq1, iq2, iq3,
- // size_k_row, size_v_row, current_block_size,
- // (unsigned)HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - factx->t_start));
}
const uint32_t h = iq2; // head index
- const float slope = factx->slopes[h];
+ const __fp16 slope = factx->slopes[h];
HVX_Vector S_vec = hvx_vec_splat_f32(0.0f);
- HVX_Vector M_vec = hvx_vec_splat_f32(-INFINITY);
+ HVX_Vector M_vec = hvx_vec_splat_f32(HTP_FA_M_INITIAL_VAL);
// Clear accumulator
hvx_splat_f32_a(spad_a, 0, DV);
}
const HVX_Vector slope_vec = hvx_vec_splat_f16(slope);
+ const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
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);
uint8_t * v_base = dma_queue_pop(dma).dst; // V
__fp16 * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M
- // FARF(HIGH, "fa %u: process: ir %u ib %u : iq1 %u iq2 %u iq3 %u q_ptr_vtcm %p : usec %u",
- // ith, ir, ib, iq1, iq2, iq3, q_ptr_vtcm,
- // (unsigned)HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - factx->t_start));
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);
// Inner loop processing the block from VTCM
- uint32_t ic = 0;
-
- // Process in sub-blocks of 32 (VLEN_FP32)
- HVX_Vector sb_scores[FLASH_ATTN_BLOCK_SIZE / VLEN_FP32];
- HVX_Vector v_max = hvx_vec_splat_f32(-INFINITY);
- for (uint32_t iv = 0; ic < current_block_size; ic += VLEN_FP32, ++iv) {
- // 1. Compute scores
- HVX_Vector scores = hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base + ic * factx->size_k_row_padded, factx->size_k_row_padded, DK, factx->scale);
-
- // 2. Softcap
- if (factx->logit_softcap != 0.0f) {
- scores = hvx_vec_tanh_f32(scores);
- scores = HVX_OP_MUL_F32(scores, logit_cap);
- }
+ // 1. Compute scores (64 elements FP16)
+ HVX_Vector scores_f16 = Q6_V_vzero();
+ if (current_block_size > 0) {
+ HVX_Vector scores0 = hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base, factx->size_k_row_padded, DK, factx->scale);
+ HVX_Vector scores1 = (current_block_size > 32) ? hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base + 32 * factx->size_k_row_padded, factx->size_k_row_padded, DK, factx->scale) : Q6_V_vzero();
+ scores_f16 = hvx_vec_f32_to_f16(scores0, scores1);
+ }
- // 3. Mask
- if (mask) {
- const __fp16 * mp = m_base + ic;
- HVX_Vector m_vals_f16 = *(const HVX_UVector *) mp;
-
- // Multiplying -INFINITY (0xFC00) by a slope in VhfVhf instructions can incorrectly produce NaN on v79.
- // Clamp -INFINITY to the max negative fp16 finite value (-65504.0f).
- 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);
-
- #if __HVX_ARCH__ >= 79
- HVX_VectorPair m_vals_f32_pair = Q6_Wsf_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(m_vals_f16), slope_vec);
- HVX_Vector add_val = Q6_V_lo_W(m_vals_f32_pair);
- scores = Q6_Vsf_vadd_VsfVsf(add_val, scores);
- #else
- HVX_VectorPair m_vals_f32_pair = Q6_Wqf32_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(m_vals_f16), slope_vec);
- HVX_Vector add_val = Q6_V_lo_W(m_vals_f32_pair);
- scores = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(add_val, scores));
- #endif
- }
+ // 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);
+ }
- // Mask out invalid lanes for leftover handling
- uint32_t valid_lanes = current_block_size - ic;
- if (valid_lanes < VLEN_FP32) {
- HVX_VectorPred valid_pred = Q6_Q_vsetq_R(valid_lanes * 4); // 4 bytes per fp32 lane
- scores = Q6_V_vmux_QVV(valid_pred, scores, hvx_vec_splat_f32(-INFINITY));
- }
+ HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(current_block_size * sizeof(__fp16));
- sb_scores[iv] = scores;
- v_max = hvx_vec_reduce_max2_f32(scores, v_max); // All lanes have block max
+ // 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 m_scaled = hvx_vec_mul_f16_f16(m_vals_f16, slope_vec);
+ scores_f16 = Q6_V_vmux_QVV(q_tail_keep, hvx_vec_add_f16_f16(scores_f16, m_scaled), v_neg_inf);
+ } else {
+ scores_f16 = Q6_V_vmux_QVV(q_tail_keep, scores_f16, v_neg_inf);
}
+ // Compute block max in FP16
+ HVX_Vector v_max_f16 = hvx_vec_reduce_max_f16(scores_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);
+
+ 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);
- HVX_Vector ms_vec = hvx_vec_exp_f32(diff_vec);
+
+ HVX_Vector diff_f16 = hvx_vec_f32_to_f16(diff_vec, diff_vec);
+ HVX_Vector diff_base2 = hvx_vec_mul_f16_f16(diff_f16, v_log2e);
+ HVX_Vector ms_f16 = hvx_vec_exp2_f16(diff_base2);
+ HVX_Vector ms_vec = Q6_V_lo_W(hvx_vec_f16_to_f32(ms_f16));
+
M_vec = M_new_vec;
hvx_scale_vec_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms_vec);
- HVX_Vector p_sum_vec = hvx_vec_splat_f32(0.0f);
- for (uint32_t ic2 = 0, iv = 0; ic2 < current_block_size; ic2 += VLEN_FP32, ++iv) {
- HVX_Vector scores = sb_scores[iv];
- HVX_Vector scores_shifted = HVX_OP_SUB_F32(scores, M_vec);
- HVX_Vector P = hvx_vec_exp_f32(scores_shifted);
+ // Compute P = exp2((S - M) * log2(e)) in FP16
+ HVX_Vector v_m_vec_f16 = hvx_vec_f32_to_f16(M_vec, M_vec);
+ HVX_Vector v_s_minus_m = Q6_Vqf16_vsub_VhfVhf(scores_f16, v_m_vec_f16);
- p_sum_vec = HVX_OP_ADD_F32(p_sum_vec, P);
+ HVX_Vector v_s_minus_m_base2 = hvx_vec_mul_f16_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m), v_log2e);
- // 5. Accumulate V
- __fp16 __attribute__((aligned(VLEN))) p_arr[VLEN_FP16];
- hvx_vec_f32_to_f16_a(p_arr, P, hvx_vec_splat_f32(0));
+ HVX_Vector P = hvx_vec_exp2_f16(v_s_minus_m_base2);
+ P = Q6_V_vmux_QVV(q_tail_keep, P, Q6_V_vzero());
- float __attribute__((aligned(128))) P_arr[VLEN_FP32];
- hvx_vec_store_a(P_arr, 128, P);
+ // Convert P to FP32 to update the running sum S_vec
+ HVX_VectorPair P_pair = hvx_vec_f16_to_f32(P);
+ HVX_Vector P0 = Q6_V_lo_W(P_pair);
+ HVX_Vector P1 = Q6_V_hi_W(P_pair);
+ HVX_Vector p_sum_vec = hvx_vec_reduce_sum_f32(HVX_OP_ADD_F32(P0, P1));
- for (uint32_t j = 0; j < VLEN_FP32; j += 2) {
- const uint32_t cur_ic = ic2 + j;
- if (cur_ic >= current_block_size) {
- break;
- }
+ S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec);
- if (cur_ic + 1 == current_block_size) {
- // Odd leftover, process single row
- if (P_arr[j] != 0.0f) {
- const uint8_t * v_ptr = v_base + cur_ic * factx->size_v_row_padded;
- hvx_mad_f32_f16_aa(VKQ32, v_ptr, (p_arr + j), DV);
- }
- break;
- }
+ // 5. Accumulate V (F16 * F16 -> F32 accumulator)
+ __fp16 __attribute__((aligned(128))) p_arr[VLEN_FP16];
+ hvx_vec_store_a(p_arr, 128, P);
- // Avoid NaN * 0.0 = NaN for uninitialized V cache rows.
- // Check the f32 values to safely avoid strict aliasing violations.
- if (P_arr[j] == 0.0f && P_arr[j + 1] == 0.0f) {
- continue;
+ 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);
}
+ break;
+ }
- const uint8_t * v_ptr = v_base + cur_ic * 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);
+ if (p_arr[j] == 0.0f && p_arr[j + 1] == 0.0f) {
+ continue;
}
- }
- p_sum_vec = hvx_vec_reduce_sum_f32(p_sum_vec);
- S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec);
+ 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);
+ }
}
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
// Issue DMA for next+1 block (if exists)
if (ib + 2 < factx->n_blocks) {
const uint8_t * m_src = (const uint8_t *) (mp_base + next_ic_start);
dma_cache_push(dma, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
}
-
- // FARF(HIGH, "fa %u: prefetch KVM: ir %u ib %u : iq1 %u iq2 %u iq3 %u : size_k_row %u size_v_row %u bs %u: usec %u",
- // ith, ir, next_ib, iq1, iq2, iq3,
- // size_k_row, size_v_row, next_block_size,
- // (unsigned)HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - factx->t_start));
}
}
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir);
// sinks
float M = hvx_vec_get_f32(M_vec);
float S = hvx_vec_get_f32(S_vec);
// Store result
// dst indices
- const int i1 = iq1;
- const int i2 = iq2;
- const int i3 = iq3;
+ const uint32_t i1 = iq1;
+ const uint32_t i2 = iq2;
+ const uint32_t i3 = iq3;
// dst is permuted: [DV, n_heads, n_tokens, n_seq]
// head stride is nb[1], token stride is nb[2], batch stride is nb[3]
} else if (dst->type == HTP_TYPE_F16) {
hvx_copy_f16_f32_ua(dst_ptr, (uint8_t *) VKQ32, DV);
}
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, ir);
+ }
+}
+
+// ============================================================================
+// HMX Phase args and thread logic
+// ============================================================================
+
+typedef struct {
+ struct hmx_fa_context * factx;
+ uint32_t kv_rows;
+ size_t src_stride;
+ size_t buf_idx;
+ uint32_t kv_start;
+ uint32_t rows_per_t;
+} fa_k_int_args_t;
+
+static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
+ fa_k_int_args_t * args = (fa_k_int_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const uint32_t total_rows = args->kv_rows;
+ const uint32_t rows_per_t = args->rows_per_t;
+ const uint32_t start = i * rows_per_t;
+ const uint32_t end = (uint32_t) hex_smin(start + rows_per_t, total_rows);
+
+ if (start >= total_rows) {
+ return;
+ }
+
+ struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ 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, factx->vtcm_k_fp16[args->buf_idx], 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, size_t buf_idx, uint32_t kv_start) {
+ worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ 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, buf_idx, kv_start, rows_per_t };
+ if (n > 1) {
+ worker_pool_run_func(wp, fa_k_interleave_thread, &args, n);
+ } else {
+ fa_k_interleave_thread(1, 0, &args);
+ }
+}
+
+typedef struct {
+ struct hmx_fa_context * factx;
+ uint32_t kv_rows;
+ size_t src_stride;
+ size_t buf_idx;
+ size_t n_col_tiles;
+ uint32_t kv_start;
+ uint32_t rows_per_t;
+} fa_v_int_args_t;
+
+static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) {
+ fa_v_int_args_t * args = (fa_v_int_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const uint32_t total_rows = args->kv_rows;
+ const uint32_t rows_per_t = args->rows_per_t;
+ const uint32_t start = i * rows_per_t;
+ const uint32_t end = (uint32_t) hex_smin(start + rows_per_t, total_rows);
+
+ if (start >= total_rows) {
+ return;
+ }
+
+ __fp16 * v_tiles_dest = factx->pipeline ? factx->vtcm_v_tiles[args->buf_idx] : factx->vtcm_v_tiles[0];
+
+ struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ 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_dest, factx->vtcm_v_fp16[args->buf_idx], total_rows, factx->DV,
+ args->src_stride, (uint32_t) args->n_col_tiles, start, end);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
+}
+
+static void fa_phase_v_interleave(struct hmx_fa_context * factx,
+ uint32_t kv_rows,
+ size_t src_stride,
+ size_t buf_idx,
+ size_t n_col_tiles,
+ uint32_t kv_start) {
+ worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ 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, buf_idx, n_col_tiles, kv_start, rows_per_t };
+ if (n > 1) {
+ worker_pool_run_func(wp, fa_v_interleave_thread, &args, n);
+ } else {
+ fa_v_interleave_thread(1, 0, &args);
+ }
+}
+
+typedef struct {
+ struct hmx_fa_context * factx;
+ const struct htp_tensor * q;
+ uint32_t q_start;
+ uint32_t kv_head;
+ uint32_t ib3;
+ size_t n_rows_g;
+ size_t rows_per_t;
+} fa_q_load_args_t;
+
+static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
+ fa_q_load_args_t * args = (fa_q_load_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const size_t n_rows_g = args->n_rows_g;
+ const size_t G = factx->G;
+ const size_t DK = factx->DK;
+
+ // Partition the padded Q rows (g_br) across threads.
+ // Keep start/end even so r and r+1 are always in the same thread's range.
+ const size_t rows_per_t = args->rows_per_t;
+ 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;
+ 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
+ {
+ const uint32_t g_br = factx->g_br;
+ const uint32_t DV = factx->DV;
+
+ const size_t col_vec_bytes = hex_align_up(g_br * sizeof(float), 256);
+ const size_t d_tile_bytes = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
+ const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
+
+ // Initialize vtcm_l_vec & vtcm_m_vec
+ const size_t l_bytes_per_t = hex_align_up(col_vec_bytes / n, 128);
+ const size_t l_start = i * l_bytes_per_t;
+ const size_t l_end = hex_smin(l_start + l_bytes_per_t, col_vec_bytes);
+
+ const size_t m_bytes_per_t = hex_align_up(col_vec_bytes / n, 128);
+ const size_t m_start = i * m_bytes_per_t;
+ const size_t m_end = hex_smin(m_start + m_bytes_per_t, col_vec_bytes);
+
+ if (factx->sinks) {
+ const float * sinks_data = (const float *) (uintptr_t) factx->sinks->data;
+ float * m_vec = (float *) factx->vtcm_m_vec;
+ const size_t r_start = l_start / sizeof(float);
+ const size_t r_end = l_end / sizeof(float);
+ const float scale_factor = EXP_LOG2E_F;
+
+ const HVX_Vector v_scale = hvx_vec_splat_f32(scale_factor);
+
+ for (size_t r = r_start; r < r_end; r += 32) {
+ HVX_VectorAlias local_m;
+ for (size_t j = 0; j < 32; ++j) {
+ size_t curr_r = r + j;
+ if (curr_r < n_rows_g) {
+ const size_t h_idx = fastmodulo(curr_r, G, &factx->div_G);
+ const size_t head = args->kv_head * G + h_idx;
+ local_m.fp32[j] = sinks_data[head];
+ } else {
+ local_m.fp32[j] = HTP_FA_M_INITIAL_VAL;
+ }
+ }
+ HVX_Vector v_scaled = HVX_OP_MUL_F32(local_m.v, v_scale);
+ *(HVX_Vector *) (m_vec + r) = v_scaled;
+ }
+ if (l_start < col_vec_bytes) {
+ hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start);
+ }
+ } else {
+ if (l_start < col_vec_bytes) {
+ hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start);
+ }
+ if (m_start < col_vec_bytes) {
+ hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float));
+ }
+ }
+
+ // Initialize vtcm_d_tiles 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);
+ }
+
+ // Initialize vtcm_o_tiles[0] to 0
+ __fp16 * o_tile_prev = factx->vtcm_o_tiles[0];
+ const size_t o_bytes_per_t = hex_align_up(o_tile_bytes / n, 128);
+ const size_t o_start = i * o_bytes_per_t;
+ const size_t o_end = hex_smin(o_start + o_bytes_per_t, o_tile_bytes);
+ if (o_start < o_tile_bytes) {
+ hvx_splat_u8_a((char *) o_tile_prev + o_start, 0, o_end - o_start);
+ }
+ }
+
+ if (start >= factx->g_br) {
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
+ return;
+ }
+
+ const struct htp_tensor * q = args->q;
+ const uint32_t q_start = args->q_start;
+ const uint32_t kv_head = args->kv_head;
+ const uint32_t ib3 = args->ib3;
+
+ for (size_t r = start; r < end; r += 2) {
+ const size_t q_idx0 = fastdiv(r + 0, &factx->div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, &factx->div_G);
+ const size_t q_idx1 = fastdiv(r + 1, &factx->div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, &factx->div_G);
+
+ const uint8_t * q_ptr0 = (r + 0 < n_rows_g) ? ((const uint8_t *) q->data + (q_start + q_idx0) * q->nb[1] +
+ (kv_head * G + h_idx0) * q->nb[2] + ib3 * q->nb[3]) :
+ NULL;
+ const uint8_t * q_ptr1 = (r + 1 < n_rows_g) ? ((const uint8_t *) q->data + (q_start + q_idx1) * q->nb[1] +
+ (kv_head * G + h_idx1) * q->nb[2] + ib3 * q->nb[3]) :
+ NULL;
+
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = factx->vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (factx->is_q_fp32) {
+ const HVX_Vector * pv_in0 = q_ptr0 ? (const HVX_Vector *) q_ptr0 : NULL;
+ const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
+
+ for (uint32_t d = 0; d < DK / 32; ++d) {
+ HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+ HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+ out_tile[r1 / 2] = v_hf;
+ }
+ } else {
+ const HVX_Vector * pv_in0 = q_ptr0 ? (const HVX_Vector *) q_ptr0 : NULL;
+ const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
+
+ for (uint32_t d = 0; d < DK / 64; ++d) {
+ HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+ __fp16 * out_dual_tile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dual_tile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+
+ *pv_out0 = Q6_V_lo_W(vp);
+ *pv_out1 = Q6_V_hi_W(vp);
+ }
+ }
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
+}
+
+static void fa_phase_q_load(struct hmx_fa_context * factx,
+ const struct htp_tensor * q,
+ uint32_t q_start,
+ uint32_t kv_head,
+ uint32_t ib3,
+ size_t n_rows_g) {
+ worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ uint32_t n = 1;
+ if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
+ n = factx->n_threads;
+ }
+ size_t rows_per_t = hex_align_up(hmx_ceil_div(factx->g_br, n), 2);
+ fa_q_load_args_t args = { factx, q, q_start, kv_head, ib3, n_rows_g, rows_per_t };
+ if (n > 1) {
+ worker_pool_run_func(wp, fa_q_load_thread, &args, n);
+ } else {
+ fa_q_load_thread(1, 0, &args);
}
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0);
+}
+
+typedef struct {
+ struct hmx_fa_context * factx;
+ const struct htp_tensor * dst;
+ const __fp16 * o_tile_src;
+ uint32_t q_start;
+ uint32_t kv_head;
+ uint32_t ib3;
+ size_t n_rows_g;
+ size_t rows_per_t;
+} fa_o_store_args_t;
+
+static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
+ fa_o_store_args_t * args = (fa_o_store_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const size_t n_rows_g = args->n_rows_g;
+ const size_t G = factx->G;
+ const size_t DV = factx->DV;
+
+ const size_t rows_per_t = args->rows_per_t;
+ const size_t start = (size_t) i * rows_per_t;
+ const size_t end = hex_smin(start + rows_per_t, n_rows_g);
+
+ if (start >= n_rows_g) {
+ return;
+ }
+
+ struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ 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 __fp16 * o_tile_src = args->o_tile_src;
+ const uint32_t q_start = args->q_start;
+ 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);
+
+ 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]);
+
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
+
+ for (uint32_t d = 0; d < DV / 32; ++d) {
+ const HVX_Vector * in_tile = (const HVX_Vector *) (tile_row_base + d * HMX_FP16_TILE_N_ELMS);
+ HVX_VectorPair vp = hvx_vec_f16_to_f32_shuff(in_tile[r1 / 2]);
+ if (r1 % 2 == 0) {
+ *(HVX_UVector *) (out + d * 32) = Q6_V_lo_W(vp);
+ } else {
+ *(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
+ }
+ }
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
+}
+
+static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
+ fa_o_store_args_t * args = (fa_o_store_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const size_t n_rows_g = args->n_rows_g;
+ const size_t G = factx->G;
+ const size_t DV = factx->DV;
+
+ const size_t rows_per_t = args->rows_per_t;
+ const size_t start = (size_t) i * rows_per_t;
+ const size_t end = hex_smin(start + rows_per_t, n_rows_g);
+
+ if (start >= n_rows_g) {
+ return;
+ }
+
+ struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ 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 __fp16 * o_tile_src = args->o_tile_src;
+ const uint32_t q_start = args->q_start;
+ 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);
+
+ __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]);
+
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
+
+ for (uint32_t d = 0; d < DV / 64; ++d) {
+ const __fp16 * in_dual_tile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
+ const HVX_Vector * pv_in1 = pv_in0 + 16;
+ HVX_VectorPair vp = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
+ if (r1 % 2 == 0) {
+ *(HVX_UVector *) (out + d * 64) = Q6_V_lo_W(vp);
+ } else {
+ *(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
+ }
+ }
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
+}
+
+static void fa_phase_o_store(struct hmx_fa_context * factx,
+ const struct htp_tensor * dst,
+ const __fp16 * o_tile_src,
+ uint32_t q_start,
+ uint32_t kv_head,
+ uint32_t ib3,
+ size_t n_rows_g) {
+ worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+ uint32_t n = 1;
+ if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
+ n = factx->n_threads;
+ }
+ size_t rows_per_t = hmx_ceil_div(n_rows_g, n);
+ 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);
+ } else {
+ store_fn(1, 0, &args);
+ }
+}
+
+typedef struct {
+ struct hmx_fa_context * factx;
+ size_t kv_rows;
+ size_t n_rows_g;
+ size_t n_col_tiles;
+ size_t n_tiles_per_bc;
+ size_t n_row_tiles;
+ size_t n_row_tiles_g_br;
+ uint32_t Bc;
+ uint32_t G;
+ uint32_t kv_head;
+ uint32_t kv_start;
+ uint32_t q_start;
+ uint32_t ib3;
+ bool has_alibi; // true when max_bias != 0 (need slope * mask + add)
+ __fp16 * slopes;
+ const struct htp_tensor * mask;
+ const __fp16 * mask_vtcm; // VTCM mask buffer base (NULL = DDR fallback)
+ size_t mask_vtcm_row_stride; // elements (__fp16) per row in VTCM mask buffer
+ struct fastdiv_values thread_div;
+} fa_softmax_args_t;
+
+static inline void fa_softmax_impl(
+ unsigned int n, unsigned int i, void * data,
+ const bool has_mask,
+ const bool mask_broadcast,
+ const bool is_g1,
+ const bool has_alibi,
+ const bool has_softcap
+) {
+ fa_softmax_args_t * args = (fa_softmax_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const size_t n_rows_g = args->n_rows_g;
+ const size_t kv_rows = args->kv_rows;
+ const size_t Bc = args->Bc;
+ const size_t G = args->G;
+ const size_t n_tiles_per_bc = args->n_tiles_per_bc;
+ const size_t n_row_vec_cnt = hmx_ceil_div(n_rows_g, 64);
+ const uint32_t im3 = has_mask ? fastmodulo(args->ib3, args->mask->ne[3], &factx->src3_div3) : 0;
+
+ size_t vec_start = 0;
+ size_t vec_end = n_row_vec_cnt;
+ if (n > 1) {
+ const size_t vecs_per_t = fastdiv(n_row_vec_cnt + n - 1, &args->thread_div);
+ vec_start = i * vecs_per_t;
+ vec_end = hex_smin(vec_start + vecs_per_t, n_row_vec_cnt);
+ }
+
+ if (vec_start >= n_row_vec_cnt) {
+ return;
+ }
+
+ struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+ 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
+ const size_t row_buf_stride = factx->row_buf_stride;
+ HVX_Vector * my_row_buf0 = factx->vtcm_row_bufs + i * 2 * row_buf_stride;
+ HVX_Vector * my_row_buf1 = my_row_buf0 + row_buf_stride;
+
+ const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
+
+ for (size_t r_vec_idx = vec_start; r_vec_idx < vec_end; ++r_vec_idx) {
+ HVX_Vector rowmax_acc_v = v_neg_inf;
+ HVX_Vector rowsum_acc_v = Q6_V_vzero();
+ HVX_Vector m_prev_v0 = factx->vtcm_m_vec[r_vec_idx * 2 + 0];
+ HVX_Vector m_prev_v1 = factx->vtcm_m_vec[r_vec_idx * 2 + 1];
+
+ HVX_Vector v_slopes = Q6_V_vzero();
+ if (has_alibi) {
+ v_slopes = hvx_vmem(args->slopes + r_vec_idx * 64);
+ }
+
+ for (uint32_t r_vec_off = 0; r_vec_off < 64; r_vec_off += 2) {
+ uint32_t r = r_vec_idx * 64 + r_vec_off;
+ if (r >= hex_align_up(n_rows_g, 2)) {
+ break;
+ }
+
+ 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;
+
+ // Decode 2 rows from S tiles into per-thread row buffers
+ if (has_softcap) {
+ const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
+ for (size_t c = 0; c < kv_rows; c += 64) {
+ size_t ci = c / 64;
+ const __fp16 * in_dual_tile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
+ const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
+
+ HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
+ HVX_Vector v_s_row0 = Q6_V_lo_W(vp_s_dual_row);
+ HVX_Vector v_s_row1 = Q6_V_hi_W(vp_s_dual_row);
+
+ HVX_Vector t0 = hvx_vec_tanh_f16(v_s_row0);
+ my_row_buf0[ci] = hvx_vec_mul_f16_f16(t0, v_cap);
+
+ HVX_Vector t1 = hvx_vec_tanh_f16(v_s_row1);
+ 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 ci = c / 64;
+ const __fp16 * in_dual_tile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
+ const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
+
+ HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
+ my_row_buf0[ci] = Q6_V_lo_W(vp_s_dual_row);
+ my_row_buf1[ci] = Q6_V_hi_W(vp_s_dual_row);
+ }
+ }
+
+ // Apply mask & compute rowmax(S)
+ HVX_Vector v_slope0 = Q6_V_vzero();
+ HVX_Vector v_slope1 = Q6_V_vzero();
+ if (has_alibi) {
+ v_slope0 = hvx_vec_repl_f16(Q6_V_vror_VR(v_slopes, r_vec_off * 2));
+ v_slope1 = (r + 1 < n_rows_g) ? hvx_vec_repl_f16(Q6_V_vror_VR(v_slopes, (r_vec_off + 1) * 2)) : Q6_V_vzero();
+ }
+
+ const HVX_Vector v_threshold = Q6_Vh_vsplat_R(0xcc00); // fp16 -16.0
+
+ 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) {
+ HVX_Vector v_mask0, v_mask1;
+
+ if (mask_broadcast) {
+ if (is_g1) {
+ const size_t qi0 = r + 0;
+ v_mask0 = *(const HVX_Vector *) (args->mask_vtcm + qi0 * args->mask_vtcm_row_stride + c);
+ v_mask1 = v_neg_inf;
+ if (r + 1 < n_rows_g) {
+ const size_t qi1 = r + 1;
+ v_mask1 = *(const HVX_Vector *) (args->mask_vtcm + qi1 * args->mask_vtcm_row_stride + c);
+ }
+ } else {
+ const size_t qi0 = fastdiv(r + 0, &factx->div_G);
+ v_mask0 = *(const HVX_Vector *) (args->mask_vtcm + qi0 * args->mask_vtcm_row_stride + c);
+ v_mask1 = v_neg_inf;
+ if (r + 1 < n_rows_g) {
+ const size_t qi1 = fastdiv(r + 1, &factx->div_G);
+ if (qi1 == qi0) {
+ v_mask1 = v_mask0;
+ } else {
+ v_mask1 = *(const HVX_Vector *) (args->mask_vtcm + qi1 * args->mask_vtcm_row_stride + c);
+ }
+ }
+ }
+ } else {
+ // Head-dependent mask: pre-interleaved per row r.
+ const size_t r0 = r + 0;
+ v_mask0 = *(const HVX_Vector *) (args->mask_vtcm + r0 * args->mask_vtcm_row_stride + c);
+ v_mask1 = v_neg_inf;
+ if (r + 1 < n_rows_g) {
+ const size_t r1 = r + 1;
+ v_mask1 = *(const HVX_Vector *) (args->mask_vtcm + r1 * args->mask_vtcm_row_stride + c);
+ }
+ }
+
+ // Threshold: mask values below -16.0 are treated as -inf (causal mask).
+ HVX_VectorPred q_keep0 = Q6_Q_and_QQ(Q6_Q_vcmp_gt_VhfVhf(v_mask0, v_threshold), q_tail_keep);
+ HVX_VectorPred q_keep1 = Q6_Q_and_QQ(Q6_Q_vcmp_gt_VhfVhf(v_mask1, v_threshold), q_tail_keep);
+
+ // Scale mask values by log2(e) for base-2 calculations
+ const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
+ HVX_Vector v_mask0_scaled = hvx_vec_mul_f16_f16(v_mask0, v_log2e);
+ HVX_Vector v_mask1_scaled = hvx_vec_mul_f16_f16(v_mask1, v_log2e);
+
+ if (has_alibi) {
+ HVX_Vector v_sm0 = hvx_vec_mul_f16_f16(v_mask0_scaled, v_slope0);
+ HVX_Vector v_sm1 = hvx_vec_mul_f16_f16(v_mask1_scaled, v_slope1);
+ my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_sm0), v_neg_inf);
+ my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_sm1), v_neg_inf);
+ } else {
+ 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 {
+ 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 = hvx_vec_reduce_max_f16(v_s_rowmax0);
+ v_s_rowmax1 = hvx_vec_reduce_max_f16(v_s_rowmax1);
+
+ // Splat m_prev[r], m_prev[r+1] from the float per-row accumulators and convert to fp16 vectors
+ HVX_Vector v_m_prev0, v_m_prev1;
+ if (r_vec_off < 32) {
+ HVX_Vector v0 = hvx_vec_repl_f32(Q6_V_vror_VR(m_prev_v0, r_vec_off * 4));
+ v_m_prev0 = hvx_vec_f32_to_f16(v0, v0);
+ if (r + 1 < n_rows_g) {
+ HVX_Vector v1 = hvx_vec_repl_f32(Q6_V_vror_VR(m_prev_v0, (r_vec_off + 1) * 4));
+ v_m_prev1 = hvx_vec_f32_to_f16(v1, v1);
+ } else {
+ v_m_prev1 = Q6_V_vzero();
+ }
+ } else {
+ HVX_Vector v0 = hvx_vec_repl_f32(Q6_V_vror_VR(m_prev_v1, (r_vec_off - 32) * 4));
+ v_m_prev0 = hvx_vec_f32_to_f16(v0, v0);
+ if (r + 1 < n_rows_g) {
+ HVX_Vector v1 = hvx_vec_repl_f32(Q6_V_vror_VR(m_prev_v1, (r_vec_off + 1 - 32) * 4));
+ v_m_prev1 = hvx_vec_f32_to_f16(v1, v1);
+ } else {
+ v_m_prev1 = Q6_V_vzero();
+ }
+ }
+
+ HVX_Vector v_dup_m0 = Q6_Vhf_vmax_VhfVhf(v_m_prev0, v_s_rowmax0);
+ HVX_Vector v_dup_m1 = Q6_Vhf_vmax_VhfVhf(v_m_prev1, v_s_rowmax1);
+
+ // Insert row r, r+1 rowmax into rowmax_acc_v
+ {
+ HVX_VectorPred p_start = Q6_Q_vsetq_R(r_vec_off * 2);
+ HVX_VectorPred p_mid = Q6_Q_vsetq_R((r_vec_off + 1) * 2);
+ HVX_VectorPred p_end = Q6_Q_vsetq2_R((r_vec_off + 2) * 2);
+ HVX_VectorPred p_lane0 = Q6_Q_and_QQn(p_mid, p_start);
+ HVX_VectorPred p_lane1 = Q6_Q_and_QQn(p_end, p_mid);
+ rowmax_acc_v = Q6_V_vmux_QVV(p_lane0, v_dup_m0, rowmax_acc_v);
+ rowmax_acc_v = Q6_V_vmux_QVV(p_lane1, v_dup_m1, rowmax_acc_v);
+ }
+
+ // Compute P = exp(S - m_new)
+ const HVX_Vector v_zero = Q6_V_vzero();
+ 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;
+ 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);
+
+ HVX_Vector v_p_row0_hf = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0));
+ HVX_Vector v_p_row1_hf = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1));
+ __fp16 * out_dual_tile = p_st_base + (c / 64) * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_p_out0 = ((HVX_Vector *) out_dual_tile) + r1 / 2;
+ HVX_Vector * pv_p_out1 = pv_p_out0 + 16;
+
+ HVX_VectorPair vp_p_dual = Q6_W_vshuff_VVR(v_p_row1_hf, v_p_row0_hf, -2);
+ *pv_p_out0 = Q6_V_lo_W(vp_p_dual);
+ *pv_p_out1 = Q6_V_hi_W(vp_p_dual);
+
+ HVX_VectorPair vp_p0 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf);
+ HVX_VectorPair vp_p1 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf);
+
+ v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0), Q6_V_hi_W(vp_p0)));
+ v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1), Q6_V_hi_W(vp_p1)));
+ }
+
+ HVX_Vector rowsum0_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(v_p_rowsum0));
+ HVX_Vector rowsum1_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(v_p_rowsum1));
+ {
+ HVX_Vector rv0_v = hvx_vec_f32_to_f16(rowsum0_sf, rowsum0_sf);
+ HVX_Vector rv1_v = hvx_vec_f32_to_f16(rowsum1_sf, rowsum1_sf);
+
+ HVX_VectorPred p_start = Q6_Q_vsetq_R(r_vec_off * 2);
+ HVX_VectorPred p_mid = Q6_Q_vsetq_R((r_vec_off + 1) * 2);
+ HVX_VectorPred p_end = Q6_Q_vsetq2_R((r_vec_off + 2) * 2);
+ HVX_VectorPred p_lane0 = Q6_Q_and_QQn(p_mid, p_start);
+ HVX_VectorPred p_lane1 = Q6_Q_and_QQn(p_end, p_mid);
+ rowsum_acc_v = Q6_V_vmux_QVV(p_lane0, rv0_v, rowsum_acc_v);
+ rowsum_acc_v = Q6_V_vmux_QVV(p_lane1, rv1_v, rowsum_acc_v);
+ }
+ }
+
+ // Inline fa_ml_update_and_build_d for this vector (lock-free and in parallel)
+ HVX_VectorPair rowmax_acc_pair = hvx_vec_f16_to_f32(rowmax_acc_v);
+ HVX_Vector v_rowmax_acc_f32_0 = Q6_V_lo_W(rowmax_acc_pair);
+ HVX_Vector v_rowmax_acc_f32_1 = Q6_V_hi_W(rowmax_acc_pair);
+
+ HVX_Vector v_m_curr0 = Q6_Vsf_vmax_VsfVsf(m_prev_v0, v_rowmax_acc_f32_0);
+ HVX_Vector v_m_curr1 = Q6_Vsf_vmax_VsfVsf(m_prev_v1, v_rowmax_acc_f32_1);
+
+ HVX_Vector v_m_diff0 = HVX_OP_SUB_F32(m_prev_v0, v_m_curr0);
+ HVX_Vector v_m_diff1 = HVX_OP_SUB_F32(m_prev_v1, v_m_curr1);
+
+ HVX_Vector v_m_diff_f16 = hvx_vec_f32_to_f16(v_m_diff0, v_m_diff1);
+ HVX_Vector exp_m_diff_f16 = hvx_vec_exp2_f16(v_m_diff_f16);
+
+ HVX_VectorPair exp_m_diff_pair = hvx_vec_f16_to_f32(exp_m_diff_f16);
+ HVX_Vector exp_m_diff0 = Q6_V_lo_W(exp_m_diff_pair);
+ HVX_Vector exp_m_diff1 = Q6_V_hi_W(exp_m_diff_pair);
+
+ HVX_VectorPair rowsum_acc_pair = hvx_vec_f16_to_f32(rowsum_acc_v);
+ HVX_Vector v_rowsum_acc_f32_0 = Q6_V_lo_W(rowsum_acc_pair);
+ HVX_Vector v_rowsum_acc_f32_1 = Q6_V_hi_W(rowsum_acc_pair);
+
+ HVX_Vector v_l_curr0;
+ HVX_Vector v_l_curr1;
+ if (args->kv_start == 0 && factx->sinks != NULL) {
+ // First KV block with sinks: m_prev holds the seeded sink value (not -inf),
+ // so exp_m_diff = exp2(sink - m_curr) is the sink's contribution to the
+ // denominator. l_prev is 0 here, so add exp_m_diff directly instead of
+ // multiplying the (uninitialized) l_prev term.
+ v_l_curr0 = HVX_OP_ADD_F32(exp_m_diff0, v_rowsum_acc_f32_0);
+ v_l_curr1 = HVX_OP_ADD_F32(exp_m_diff1, v_rowsum_acc_f32_1);
+ } else {
+ HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0];
+ HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1];
+ v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0);
+ v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1);
+ }
+
+ factx->vtcm_m_vec[r_vec_idx * 2 + 0] = v_m_curr0;
+ factx->vtcm_m_vec[r_vec_idx * 2 + 1] = v_m_curr1;
+ factx->vtcm_l_vec[r_vec_idx * 2 + 0] = v_l_curr0;
+ factx->vtcm_l_vec[r_vec_idx * 2 + 1] = v_l_curr1;
+
+ // Build diagonal tile D = diag(exp(m_diff))
+ const HVX_Vector v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets;
+ const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16));
+ HVX_Vector v_exp_m_diff = exp_m_diff_f16;
+
+ size_t t0 = r_vec_idx * 2;
+ if (t0 < args->n_row_tiles) {
+ const HVX_Vector v_content = v_exp_m_diff;
+ __fp16 * out_base = factx->vtcm_d_tiles + t0 * (args->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);
+ }
+
+ size_t t1 = r_vec_idx * 2 + 1;
+ if (t1 < args->n_row_tiles) {
+ const HVX_Vector v_content = Q6_V_vror_VR(v_exp_m_diff, 64);
+ __fp16 * out_base = factx->vtcm_d_tiles + t1 * (args->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);
+ }
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, (uint16_t) (args->q_start * G + vec_start * 64));
+}
+
+static void fa_softmax_thread_nomask(unsigned int n, unsigned int i, void * data) {
+ fa_softmax_impl(n, i, data,
+ /*has_mask=*/false,
+ /*mask_broadcast=*/false,
+ /*is_g1=*/false,
+ /*has_alibi=*/false,
+ /*has_softcap=*/false);
+}
+
+static void fa_softmax_thread_mask_broadcast_g1(unsigned int n, unsigned int i, void * data) {
+ fa_softmax_impl(n, i, data,
+ /*has_mask=*/true,
+ /*mask_broadcast=*/true,
+ /*is_g1=*/true,
+ /*has_alibi=*/false,
+ /*has_softcap=*/false);
+}
+
+static void fa_softmax_thread_mask_broadcast_gn(unsigned int n, unsigned int i, void * data) {
+ fa_softmax_impl(n, i, data,
+ /*has_mask=*/true,
+ /*mask_broadcast=*/true,
+ /*is_g1=*/false,
+ /*has_alibi=*/false,
+ /*has_softcap=*/false);
+}
+
+static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) {
+ fa_softmax_args_t * args = (fa_softmax_args_t *) data;
+ struct hmx_fa_context * factx = args->factx;
+
+ const bool has_mask = (args->mask != NULL);
+ const bool mask_broadcast = factx->mask_broadcast;
+ const bool is_g1 = (args->G == 1);
+ const bool has_alibi = args->has_alibi;
+ const bool has_softcap = (factx->logit_softcap != 0.0f);
+
+ fa_softmax_impl(n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap);
+}
+
+static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_context * factx,
+ size_t n_row_tiles,
+ size_t n_row_tiles_g_br) {
+ const HVX_Vector v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets;
+ const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16));
+ const HVX_Vector one = hvx_vec_splat_f32(1.0f);
+
+ HVX_Vector v_content = Q6_V_vzero();
+ for (size_t i = 0; i < n_row_tiles; ++i) {
+ if ((i % 2) == 0) {
+ HVX_Vector inv_lo = HVX_OP_MUL_F32(one, hvx_vec_inverse_f32(factx->vtcm_l_vec[i]));
+ HVX_Vector inv_hi = (i + 1 < n_row_tiles) ? HVX_OP_MUL_F32(one, hvx_vec_inverse_f32(factx->vtcm_l_vec[i + 1])) : Q6_V_vzero();
+ v_content = hvx_vec_f32_to_f16(inv_lo, inv_hi);
+ } else {
+ 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;
+ Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
+ }
+}
+
+static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx,
+ 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;
+ const size_t n_row_vec_cnt = hmx_ceil_div(sargs->n_rows_g, 64);
+
+ worker_callback_t softmax_fn = fa_softmax_thread;
+ if (sargs->mask == NULL && factx->logit_softcap == 0.0f && !sargs->has_alibi) {
+ softmax_fn = fa_softmax_thread_nomask;
+ } else if (sargs->mask != NULL && factx->mask_broadcast && factx->logit_softcap == 0.0f && !sargs->has_alibi) {
+ if (sargs->G == 1) {
+ softmax_fn = fa_softmax_thread_mask_broadcast_g1;
+ } else {
+ softmax_fn = fa_softmax_thread_mask_broadcast_gn;
+ }
+ }
+
+ 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);
+ } else {
+ softmax_fn(1, 0, sargs);
+ }
+}
+
+// ============================================================================
+// HMX job structs and worker functions
+// ============================================================================
+
+typedef struct {
+ const __fp16 * q_tiles;
+ const __fp16 * k_tiles;
+ __fp16 * s_tiles;
+ size_t n_row_tiles;
+ size_t n_col_tiles;
+ size_t n_dot_tiles; // DK / 32
+ size_t n_tiles_per_bc;
+ uint8_t * hmx_scales;
+} hmx_fa_qk_job_t;
+
+static void hmx_fa_qk_dot_worker(void * data) {
+ hmx_fa_qk_job_t * job = (hmx_fa_qk_job_t *) data;
+ const size_t n_row_tiles = job->n_row_tiles;
+ const size_t n_col_tiles = job->n_col_tiles;
+ const size_t n_dot_tiles = job->n_dot_tiles;
+ const size_t n_tiles_per_bc = job->n_tiles_per_bc;
+ const __fp16 * restrict q_tiles = job->q_tiles;
+ const __fp16 * restrict k_tiles = job->k_tiles;
+ __fp16 * restrict s_tiles = job->s_tiles;
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(n_col_tiles > 0);
+ __builtin_assume(n_dot_tiles > 0);
+
+ Q6_bias_mxmem2_A((void *) job->hmx_scales);
+ for (size_t r = 0; r < n_row_tiles; ++r) {
+ for (size_t c = 0; c < n_col_tiles; ++c) {
+ const __fp16 * row_tiles = q_tiles + r * HMX_FP16_TILE_N_ROWS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
+ const __fp16 * col_tiles = k_tiles + c * HMX_FP16_TILE_N_COLS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
+ __fp16 * out_tile = s_tiles + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
+
+ hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles);
+ }
+ }
+}
+
+typedef struct {
+ __fp16 * o_curr;
+ const __fp16 * o_prev;
+ const __fp16 * p_tiles;
+ const __fp16 * v_tiles;
+ const __fp16 * d_tiles;
+ uint8_t * hmx_scales;
+ size_t n_row_tiles;
+ size_t n_col_tiles;
+ size_t n_row_tiles_g_br;
+ size_t n_tiles_per_bc;
+ size_t DV;
+} hmx_fa_o_update_job_t;
+
+static void hmx_fa_o_update_worker(void * data) {
+ hmx_fa_o_update_job_t * job = (hmx_fa_o_update_job_t *) data;
+ const size_t n_row_tiles = job->n_row_tiles;
+ const size_t n_col_tiles = job->n_col_tiles;
+ const size_t n_row_tiles_g_br = job->n_row_tiles_g_br;
+ const size_t n_tiles_per_bc = job->n_tiles_per_bc;
+ const size_t DV_tiles = job->DV / 32;
+ const __fp16 * restrict d_tiles = job->d_tiles;
+ const __fp16 * restrict p_tiles = job->p_tiles;
+ const __fp16 * restrict v_tiles = job->v_tiles;
+ const __fp16 * restrict o_prev = job->o_prev;
+ __fp16 * restrict o_curr = job->o_curr;
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(n_col_tiles > 0);
+ __builtin_assume(DV_tiles > 0);
+
+ Q6_bias_mxmem2_A((void *) job->hmx_scales);
+ for (size_t r = 0; r < n_row_tiles; ++r) {
+ for (size_t c = 0; c < DV_tiles; ++c) {
+ // D[r,r] @ O_prev[r,c] — only the diagonal tile
+ const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * o_rc = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * p_tile_in = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * v_tile_in = v_tiles + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
+ __fp16 * o_tile_out = o_curr + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+
+ hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles);
+ }
+ }
+}
+
+typedef struct {
+ __fp16 * o_curr; // output (row-major tile layout)
+ const __fp16 * o_prev; // input (column-major tile layout)
+ const __fp16 * d_tiles; // diag(1/l) tiles
+ uint8_t * hmx_scales;
+ size_t n_row_tiles;
+ size_t n_row_tiles_g_br;
+ size_t DV;
+} hmx_fa_o_norm_job_t;
+
+static void hmx_fa_o_norm_worker(void * data) {
+ hmx_fa_o_norm_job_t * job = (hmx_fa_o_norm_job_t *) data;
+ const size_t n_row_tiles = job->n_row_tiles;
+ const size_t n_row_tiles_g_br = job->n_row_tiles_g_br;
+ const size_t DV_tiles = job->DV / 32;
+ const __fp16 * restrict d_tiles = job->d_tiles;
+ const __fp16 * restrict o_prev = job->o_prev;
+ __fp16 * restrict o_curr = job->o_curr;
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(DV_tiles > 0);
+
+ Q6_bias_mxmem2_A((void *) job->hmx_scales);
+ for (size_t r = 0; r < n_row_tiles; ++r) {
+ for (size_t c = 0; c < DV_tiles; ++c) {
+ const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * o_rc = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+ __fp16 * o_out = o_curr + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
+
+ hmx_fa_o_norm_tile(d_diag, o_rc, o_out);
+ }
+ }
+}
+
+// Populate per-GQA-row ALiBi slopes for a given KV head.
+static __attribute__((noinline)) void fa_compute_slopes(
+ const struct hmx_fa_context * factx,
+ uint32_t kv_head,
+ size_t n_rows_g) {
+ __fp16 * slopes = factx->vtcm_slopes;
+ if (factx->max_bias == 0.0f) {
+ hvx_splat_f16_a(slopes, 1.0f, n_rows_g);
+ return;
+ }
+
+ const uint32_t G = factx->G;
+ const uint32_t n_head_log2 = factx->n_head_log2;
+ const float m0 = factx->m0;
+ const float m1 = factx->m1;
+
+ __fp16 temp_slopes[512] __attribute__((aligned(128)));
+ if (G <= 32) {
+ // Fast path: Compute G unique slope values in vector registers
+ HVX_Vector v_val = hvx_alibi_slopes(kv_head, G, n_head_log2, m0, m1);
+
+ __fp16 temp_slopes_aligned[64] __attribute__((aligned(128)));
+ hvx_vmem(temp_slopes_aligned) = hvx_vec_f32_to_f16(v_val, Q6_V_vzero());
+
+ for (uint32_t i = 0; i < G; ++i) {
+ temp_slopes[i] = temp_slopes_aligned[i];
+ }
+ } else {
+ // Fallback path: G > 32 (rare configurations)
+ for (uint32_t i = 0; i < G; ++i) {
+ temp_slopes[i] = (__fp16)alibi_slope(kv_head * G + i, n_head_log2, m0, m1);
+ }
+ }
+
+ // Allocate stack buffer to avoid scalar writes to VTCM (which generates L2 misses)
+ __fp16 local_slopes[n_rows_g] __attribute__((aligned(128)));
+ for (size_t r = 0; r < n_rows_g; ++r) {
+ local_slopes[r] = temp_slopes[fastmodulo(r, G, &factx->div_G)];
+ }
+
+ // Copy to VTCM slopes using HVX block copy (both are aligned to 128 bytes)
+ hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)local_slopes, n_rows_g);
+}
+
+static void fa_push_mask_dma_gqa(
+ dma_queue * dma,
+ const struct htp_tensor * mask,
+ uint32_t q_start,
+ uint32_t im3,
+ uint32_t kv_start,
+ uint32_t kv_head,
+ uint32_t G,
+ uint32_t m_line_bytes,
+ uint32_t kv_rows,
+ uint32_t n_q_rows,
+ struct hmx_fa_context * factx
+) {
+ for (uint32_t g = 0; g < G; ++g) {
+ const uint32_t h_idx = kv_head * G + g;
+ const uint32_t im2 = fastmodulo(h_idx, mask->ne[2], &factx->src3_div2);
+ const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] +
+ im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
+ uint8_t * ms_dst = (uint8_t *) factx->vtcm_mask_buf + g * m_line_bytes;
+ dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+ }
+}
+
+static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) {
+ for (uint32_t g = 0; g < G; ++g) {
+ dma_queue_pop(dma);
+ }
+}
+
+// ============================================================================
+// 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;
+ const struct htp_tensor * q = octx->src[0];
+ const struct htp_tensor * k = octx->src[1];
+ const struct htp_tensor * v = octx->src[2];
+ const struct htp_tensor * mask = (octx->src[3] && octx->src[3]->data) ? octx->src[3] : NULL;
+ const struct htp_tensor * dst = octx->dst;
+
+ struct htp_context * const ctx = octx->ctx;
+
+ if (!ctx->hmx_enabled) {
+ return HTP_STATUS_NO_SUPPORT;
+ }
+
+ // Dimensions
+ const uint32_t neq0 = q->ne[0]; // head_dim (DK)
+ const uint32_t neq1 = q->ne[1]; // n_tokens
+ const uint32_t neq2 = q->ne[2]; // n_heads
+ const uint32_t neq3 = q->ne[3]; // n_seqs
+
+ const uint32_t nek0 = k->ne[0]; // head_dim
+ const uint32_t nek1 = k->ne[1]; // kv_len
+
+ const uint32_t nev0 = v->ne[0]; // head_dim (DV)
+
+ const uint32_t DK = neq0;
+ const uint32_t DV = nev0;
+
+ // HMX requires head_dim to be multiple of 32
+ if (DK % 32 != 0 || DV % 32 != 0) {
+ return HTP_STATUS_NO_SUPPORT;
+ }
+
+ const struct htp_fa_kernel_params * kparams = (const struct htp_fa_kernel_params *) octx->kernel_params;
+ const uint32_t n_kv_heads = k->ne[2];
+
+ // ======== Build context ========
+ struct hmx_fa_context factx;
+ memset(&factx, 0, sizeof(factx));
+ factx.octx = octx;
+ factx.sinks = octx->src[4]; // NULL if this op has no attention sinks
+ factx.n_threads = kparams->n_threads;
+ factx.DK = DK;
+ factx.DV = DV;
+ factx.n_kv = nek1;
+ factx.n_kv_heads = n_kv_heads;
+ factx.n_heads = neq2;
+ factx.G = kparams->G;
+ factx.div_G = kparams->u.hmx.div_G;
+ factx.neq1 = neq1;
+ factx.Br = kparams->Br;
+ factx.Bc = kparams->Bc;
+ factx.g_br = kparams->u.hmx.g_br;
+ factx.n_kv_blocks = kparams->n_kv_blocks;
+ factx.is_q_fp32 = (kparams->is_q_fp32 != 0);
+ factx.is_dst_fp32 = (kparams->is_dst_fp32 != 0);
+ factx.pipeline = (kparams->u.hmx.pipeline != 0);
+ factx.mask_broadcast = (kparams->u.hmx.mask_broadcast != 0);
+ if (mask) {
+ factx.src3_div2 = kparams->src3_div2;
+ factx.src3_div3 = kparams->src3_div3;
+ }
+
+ if (kparams->logit_softcap == 0.0f) {
+ factx.scale = (__fp16) (kparams->scale * EXP_LOG2E_F); // log2(e)
+ } else {
+ factx.scale = (__fp16) kparams->scale;
+ }
+ factx.max_bias = kparams->max_bias;
+ factx.logit_softcap = (__fp16) (kparams->logit_softcap * EXP_LOG2E_F);
+
+ factx.n_head_log2 = kparams->n_head_log2;
+ factx.m0 = kparams->m0;
+ factx.m1 = kparams->m1;
+
+ const uint32_t Br = factx.Br;
+ const uint32_t Bc = factx.Bc;
+ const uint32_t g_br = factx.g_br;
+ const bool pipeline = factx.pipeline;
+ const uint32_t n_threads = factx.n_threads;
+ const uint32_t G = factx.G;
+
+ // ======== VTCM allocation (GQA-aware) ========
+ const size_t size_k_row = DK * sizeof(__fp16);
+ const size_t size_v_row = DV * sizeof(__fp16);
+ const size_t size_k_row_padded = hex_round_up(size_k_row, 128);
+ const size_t size_v_row_padded = hex_round_up(size_v_row, 128);
+
+ const size_t q_tile_bytes = hex_align_up(g_br * DK * sizeof(__fp16), 4096);
+ const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
+ const size_t k_dma_bytes = hex_align_up(Bc * size_k_row_padded, 4096);
+ const size_t v_dma_bytes = hex_align_up(Bc * size_v_row_padded, 4096);
+ const size_t k_tile_bytes = hex_align_up(Bc * DK * sizeof(__fp16), 4096);
+ const size_t v_tile_bytes = hex_align_up(Bc * DV * sizeof(__fp16), 4096);
+ const size_t s_tile_bytes = hex_align_up(g_br * Bc * sizeof(__fp16), 4096);
+ const size_t d_tile_bytes = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
+ const size_t col_vec_bytes = hex_align_up(g_br * sizeof(float), 256);
+ const size_t row_vec_bytes = hex_align_up(Bc * sizeof(__fp16), 256);
+ const size_t m_line_bytes = hex_align_up(Bc * sizeof(__fp16), 128);
+ const size_t m_buf_bytes = hex_align_up(Br * m_line_bytes, 4096) * HMX_FA_DMA_CACHE_SIZE;
+ const size_t slopes_bytes = hex_align_up(g_br * sizeof(__fp16), 128);
+
+ uint8_t * vtcm_cur = ctx->vtcm_base;
+
+ factx.vtcm_q_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, q_tile_bytes);
+ factx.vtcm_o_tiles[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
+ factx.vtcm_o_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
+ factx.vtcm_k_fp16[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
+ factx.vtcm_k_fp16[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
+ factx.vtcm_v_fp16[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
+ factx.vtcm_v_fp16[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
+ factx.vtcm_k_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_tile_bytes);
+ factx.vtcm_v_tiles[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
+ if (pipeline) {
+ factx.vtcm_v_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
+ } else {
+ factx.vtcm_v_tiles[1] = NULL;
+ }
+ factx.vtcm_s_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
+ factx.vtcm_p_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
+ factx.vtcm_d_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, d_tile_bytes);
+ factx.vtcm_m_vec = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
+ factx.vtcm_l_vec = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
+ factx.vtcm_s_rowmax = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
+ factx.vtcm_p_rowsum = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
+ factx.vtcm_row_bufs = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, row_vec_bytes * 2 * n_threads);
+ factx.row_buf_stride = row_vec_bytes / sizeof(HVX_Vector);
+ factx.vtcm_hmx_scales_id = vtcm_seq_alloc(&vtcm_cur, 256);
+ factx.vtcm_hmx_scales_qk = vtcm_seq_alloc(&vtcm_cur, 256);
+ factx.vtcm_mask_buf = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, m_buf_bytes);
+ factx.mask_buf_row_stride = m_line_bytes / sizeof(__fp16);
+ factx.vtcm_slopes = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, slopes_bytes);
+
+ dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, hex_align_up(Br * m_line_bytes, 4096), HMX_FA_DMA_CACHE_SIZE);
+
+ if ((size_t) (vtcm_cur - ctx->vtcm_base) > ctx->vtcm_size) {
+ return HTP_STATUS_VTCM_TOO_SMALL;
+ }
+
+ // ======== Initialize HMX output scales ========
+ hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0
+ hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale));
+
+ // ======== Skip compute if profiling ========
+ if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
+ return HTP_STATUS_OK;
+ }
+
+ // ======== DMA setup ========
+ dma_queue * const dma = ctx->dma[0];
+
+ const size_t n_row_tiles_g_br = g_br / HMX_FP16_TILE_N_ROWS;
+ const size_t n_tiles_per_bc = Bc / HMX_FP16_TILE_N_COLS;
+
+ const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
+
+ // ======== HMX lock strategy ========
+ if (!factx.pipeline) {
+ HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+ }
+
+ // ======== Reusable job descriptors for pipeline ========
+ hmx_fa_qk_job_t qk_job;
+ hmx_fa_o_update_job_t ou_job;
+ hmx_fa_o_norm_job_t on_job;
+
+ // ======== Main loop ========
+ for (uint32_t ib3 = 0; ib3 < neq3; ++ib3) {
+ const uint32_t im3 = mask ? fastmodulo(ib3, mask->ne[3], &factx.src3_div3) : 0;
+ for (uint32_t q_start = 0; q_start < neq1; q_start += Br) {
+ const uint32_t n_q_rows = hex_smin(Br, neq1 - q_start);
+ const size_t n_rows_g = n_q_rows * G;
+ const size_t g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
+ const size_t n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
+
+ for (uint32_t kv_head = 0; kv_head < n_kv_heads; ++kv_head) {
+ const uint32_t ik2 = kv_head;
+ const uint32_t ik3 = ib3 / (neq3 / k->ne[3]);
+ const uint32_t iv2 = kv_head;
+ const uint32_t iv3 = ib3 / (neq3 / v->ne[3]);
+
+ // Prefetch first KV block
+ if (factx.n_kv_blocks > 0) {
+ const uint32_t kv_rows0 = hex_smin(Bc, nek1);
+
+ 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);
+ }
+
+ // ---- Load Q block & Initialize per-block state ----
+ fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
+
+ __fp16 * o_tile_prev = factx.vtcm_o_tiles[0];
+ __fp16 * o_tile_curr = factx.vtcm_o_tiles[1];
+
+ // ---- KV block loop with DMA double-buffering ----
+ size_t buf_idx = 0;
+
+ htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) q_start);
+ fa_compute_slopes(&factx, kv_head, n_rows_g);
+ htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) q_start);
+
+ const size_t k_src_stride = size_k_row_padded / sizeof(__fp16);
+ const size_t v_src_stride = size_v_row_padded / sizeof(__fp16);
+
+ if (factx.pipeline) {
+ // ==================================================================
+ // Pipeline path
+ // ==================================================================
+ struct hmx_queue * hmx_q = ctx->hmx_queue;
+
+ 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_q_rows);
+ } else {
+ fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx);
+ }
+ }
+
+ // Wait for current KV DMA
+ dma_queue_pop(dma); // K
+ dma_queue_pop(dma); // V
+
+ // ---- 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));
+ }
+ fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, 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));
+
+ 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);
+ }
+ fa_phase_v_interleave(&factx, kv_rows, v_src_stride, 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 ----
+ __fp16 * current_mask_vtcm = NULL;
+ if (mask) {
+ if (__builtin_expect(factx.mask_broadcast, true)) {
+ current_mask_vtcm = (__fp16 *) dma_queue_pop(dma).dst;
+ } else {
+ fa_pop_mask_dma_gqa(dma, G);
+ current_mask_vtcm = factx.vtcm_mask_buf;
+ }
+ }
+
+ fa_softmax_args_t sargs;
+ memset(&sargs, 0, sizeof(sargs));
+ sargs.factx = &factx;
+ sargs.kv_rows = kv_rows;
+ sargs.n_rows_g = n_rows_g;
+ sargs.n_col_tiles = n_col_tiles;
+ sargs.n_tiles_per_bc = n_tiles_per_bc;
+ sargs.n_row_tiles = n_row_tiles;
+ sargs.n_row_tiles_g_br = n_row_tiles_g_br;
+ sargs.Bc = Bc;
+ sargs.G = G;
+ sargs.kv_head = kv_head;
+ sargs.kv_start = kv_start;
+ sargs.q_start = q_start;
+ sargs.ib3 = ib3;
+ sargs.has_alibi = (factx.max_bias != 0.0f);
+ sargs.mask = mask;
+ sargs.mask_vtcm = current_mask_vtcm;
+ sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
+ sargs.slopes = factx.vtcm_slopes;
+ fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
+
+ buf_idx = 1 - buf_idx;
+ }
+
+ // Epilogue
+ 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));
+ hmx_queue_pop(hmx_q);
+
+ hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
+ }
+
+ } else {
+ // ==================================================================
+ // Fallback path
+ // ==================================================================
+ 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);
+ dma_queue_pop(dma); // K
+ dma_queue_pop(dma); // V
+
+ 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_q_rows);
+ } else {
+ fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx);
+ }
+ }
+ 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);
+ }
+ fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start);
+
+ {
+ const size_t n_dot_tiles = (size_t) (DK / 32);
+ const __fp16 * restrict q_base = factx.vtcm_q_tiles;
+ const __fp16 * restrict k_base = factx.vtcm_k_tiles;
+ __fp16 * restrict s_base = factx.vtcm_s_tiles;
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(n_col_tiles > 0);
+ __builtin_assume(n_dot_tiles > 0);
+
+ htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_qk);
+ for (size_t r = 0; r < n_row_tiles; ++r) {
+ for (size_t c = 0; c < n_col_tiles; ++c) {
+ const __fp16 * row_tiles = q_base + r * HMX_FP16_TILE_N_ROWS * DK;
+ const __fp16 * col_tiles = k_base + c * HMX_FP16_TILE_N_COLS * DK;
+ __fp16 * out_tile = s_base + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
+
+ hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles);
+ }
+ }
+ htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ }
+
+ // ---- Phase 3: softmax + build_D ----
+ __fp16 * current_mask_vtcm = NULL;
+ if (mask) {
+ if (__builtin_expect(factx.mask_broadcast, true)) {
+ current_mask_vtcm = (__fp16 *) dma_queue_pop(dma).dst;
+ } else {
+ fa_pop_mask_dma_gqa(dma, G);
+ current_mask_vtcm = factx.vtcm_mask_buf;
+ }
+ }
+
+ fa_softmax_args_t sargs;
+ memset(&sargs, 0, sizeof(sargs));
+ sargs.factx = &factx;
+ sargs.kv_rows = kv_rows;
+ sargs.n_rows_g = n_rows_g;
+ sargs.n_col_tiles = n_col_tiles;
+ sargs.n_tiles_per_bc = n_tiles_per_bc;
+ sargs.n_row_tiles = n_row_tiles;
+ sargs.n_row_tiles_g_br = n_row_tiles_g_br;
+ sargs.Bc = Bc;
+ sargs.G = G;
+ sargs.kv_head = kv_head;
+ sargs.kv_start = kv_start;
+ sargs.q_start = q_start;
+ sargs.ib3 = ib3;
+ sargs.has_alibi = (factx.max_bias != 0.0f);
+ sargs.mask = mask;
+ sargs.mask_vtcm = current_mask_vtcm;
+ sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
+ sargs.slopes = factx.vtcm_slopes;
+ fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
+ fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc, kv_start);
+
+ {
+ const size_t DV_tiles = (size_t) (DV / 32);
+ const __fp16 * restrict d_base = factx.vtcm_d_tiles;
+ const __fp16 * restrict p_base = factx.vtcm_p_tiles;
+ const __fp16 * restrict v_base = factx.vtcm_v_tiles[0];
+ const __fp16 * restrict op_base = o_tile_prev;
+ __fp16 * restrict oc_base = o_tile_curr;
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(n_col_tiles > 0);
+ __builtin_assume(DV_tiles > 0);
+
+ htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
+ for (size_t r = 0; r < n_row_tiles; ++r) {
+ for (size_t c = 0; c < DV_tiles; ++c) {
+ const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * o_rc = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * p_tile_in = p_base + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * v_tile_in = v_base + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
+ __fp16 * o_tile_out = oc_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+
+ hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles);
+ }
+ }
+ htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
+ }
+
+ buf_idx = 1 - buf_idx;
+ }
+ }
+
+ // ---- 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);
+
+ if (factx.pipeline) {
+ on_job.o_curr = o_tile_curr;
+ on_job.o_prev = o_tile_prev;
+ on_job.d_tiles = factx.vtcm_d_tiles;
+ 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;
+ on_job.DV = DV;
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
+ hmx_queue_pop(ctx->hmx_queue);
+ } else {
+ const size_t DV_tiles = (size_t) (DV / 32);
+ const __fp16 * restrict d_base = factx.vtcm_d_tiles;
+ const __fp16 * restrict op_base = o_tile_prev;
+ __fp16 * restrict oc_base = o_tile_curr;
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(DV_tiles > 0);
+
+ htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
+ for (size_t r = 0; r < n_row_tiles; ++r) {
+ for (size_t c = 0; c < DV_tiles; ++c) {
+ const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * o_rc = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+ __fp16 * o_out = oc_base + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
+
+ hmx_fa_o_norm_tile(d_diag, o_rc, o_out);
+ }
+ }
+ htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ }
+ }
+
+ // ---- Store O block ----
+ fa_phase_o_store(&factx, dst, o_tile_curr, q_start, kv_head, ib3, n_rows_g);
+ }
+ }
+ }
+
+ if (factx.pipeline) {
+ hmx_queue_suspend(ctx->hmx_queue);
+ } else {
+ HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
+ }
+
+ return HTP_STATUS_OK;
}
int op_flash_attn_ext(struct htp_ops_context * octx) {
return HTP_STATUS_NO_SUPPORT;
}
- // HMX path: head_dim multiple of 64, F16 KV, and no sinks
- if (k->type == HTP_TYPE_F16 && v->type == HTP_TYPE_F16 && k->ne[0] % 64 == 0 && v->ne[0] % 64 == 0 && octx->src[4] == NULL) {
- int ret = hmx_flash_attn_ext(octx);
- if (ret == HTP_STATUS_OK) {
- return ret;
- }
- // VTCM too small or other failure -> fall through to HVX path
+ const struct htp_fa_kernel_params * kparams = (const struct htp_fa_kernel_params *) octx->kernel_params;
+
+ if (kparams->kernel_type == HTP_FA_KERNEL_UNSUPPORTED) {
+ return HTP_STATUS_NO_SUPPORT;
+ }
+
+ if (kparams->kernel_type == HTP_FA_KERNEL_HMX) {
+ return hmx_flash_attn_ext(octx);
}
struct htp_fa_context factx;
factx.t_start = HAP_perf_get_qtimer_count();
- factx.src0_div21 = init_fastdiv_values(q->ne[2] * q->ne[1]);
- factx.src0_div1 = init_fastdiv_values(q->ne[1]);
+ factx.src0_div21 = kparams->u.hvx.src0_div21;
+ factx.src0_div1 = kparams->u.hvx.src0_div1;
- factx.broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]);
- factx.broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]);
- factx.broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]);
- factx.broadcast_rv3 = init_fastdiv_values(q->ne[3]/v->ne[3]);
+ factx.broadcast_rk2 = kparams->u.hvx.broadcast_rk2;
+ factx.broadcast_rk3 = kparams->u.hvx.broadcast_rk3;
+ factx.broadcast_rv2 = kparams->u.hvx.broadcast_rv2;
+ factx.broadcast_rv3 = kparams->u.hvx.broadcast_rv3;
if (mask) {
- factx.src3_div2 = init_fastdiv_values(mask->ne[2]);
- factx.src3_div3 = init_fastdiv_values(mask->ne[3]);
+ factx.src3_div2 = kparams->src3_div2;
+ factx.src3_div3 = kparams->src3_div3;
}
- factx.is_q_fp32 = (q->type == HTP_TYPE_F32);
- factx.size_q_row_padded = hex_round_up(q->ne[0] * (factx.is_q_fp32 ? 4 : 2), 128);
- factx.size_k_row_padded = hex_round_up(k->ne[0] * sizeof(__fp16), 128);
- factx.size_v_row_padded = hex_round_up(v->ne[0] * sizeof(__fp16), 128);
+ factx.is_q_fp32 = (kparams->is_q_fp32 != 0);
+ factx.size_q_row_padded = kparams->u.hvx.size_q_row_padded;
+ factx.size_k_row_padded = kparams->u.hvx.size_k_row_padded;
+ factx.size_v_row_padded = kparams->u.hvx.size_v_row_padded;
size_t size_q_block = factx.size_q_row_padded * 1; // single row for now
factx.size_k_block = factx.size_k_row_padded * FLASH_ATTN_BLOCK_SIZE;
factx.size_v_block = factx.size_v_row_padded * FLASH_ATTN_BLOCK_SIZE;
factx.size_m_block = hex_round_up(FLASH_ATTN_BLOCK_SIZE * sizeof(__fp16), 128);
- factx.n_blocks = (k->ne[1] + FLASH_ATTN_BLOCK_SIZE - 1) / FLASH_ATTN_BLOCK_SIZE;
-
- float scale = 1.0f;
- float max_bias = 0.0f;
- float logit_softcap = 0.0f;
-
- memcpy(&scale, (float *) octx->op_params + 0, sizeof(float));
- memcpy(&max_bias, (float *) octx->op_params + 1, sizeof(float));
- memcpy(&logit_softcap, (float *) octx->op_params + 2, sizeof(float));
-
- if (logit_softcap != 0.0f) {
- scale /= logit_softcap;
- }
+ factx.n_blocks = kparams->n_kv_blocks;
- factx.scale = scale;
- factx.max_bias = max_bias;
- factx.logit_softcap = logit_softcap;
+ factx.scale = kparams->scale;
+ factx.max_bias = kparams->max_bias;
+ factx.logit_softcap = (__fp16) kparams->logit_softcap;
- uint32_t n_head = q->ne[2];
- factx.n_head_log2 = 1u << (uint32_t) floor(log2(n_head));
- factx.m0 = powf(2.0f, -(max_bias ) / factx.n_head_log2);
- factx.m1 = powf(2.0f, -(max_bias / 2.0f) / factx.n_head_log2);
+ factx.n_head_log2 = kparams->n_head_log2;
+ factx.m0 = kparams->m0;
+ factx.m1 = kparams->m1;
+ const uint32_t n_head = q->ne[2];
if (n_head > 512) {
return HTP_STATUS_NO_SUPPORT;
}
for (uint32_t h = 0; h < n_head; ++h) {
- factx.slopes[h] = (max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f;
+ factx.slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
}
// total rows in q
- const uint32_t neq0 = q->ne[0];
- const uint32_t neq1 = q->ne[1];
- const uint32_t neq2 = q->ne[2];
- const uint32_t neq3 = q->ne[3];
-
- factx.qrows = neq1*neq2*neq3;
- factx.qrows_per_thread = (factx.qrows + octx->n_threads - 1) / octx->n_threads;
+ factx.qrows = kparams->qrows;
+ factx.qrows_per_thread = kparams->qrows_per_thread;
size_t size_vkq_acc = hex_round_up(v->ne[0] * sizeof(float), 128); // VKQ32
- octx->src0_spad.size_per_thread = size_q_block * 1;
- octx->src1_spad.size_per_thread = factx.size_k_block * 2;
- octx->src2_spad.size_per_thread = factx.size_v_block * 2;
- octx->src3_spad.size_per_thread = mask ? factx.size_m_block * DMA_CACHE_MAX_SIZE : 0;
- octx->dst_spad.size_per_thread = size_vkq_acc;
+ factx.size_q_block = size_q_block;
+ factx.size_vkq_acc = size_vkq_acc;
- octx->src0_spad.size = octx->src0_spad.size_per_thread * octx->n_threads;
- octx->src1_spad.size = octx->src1_spad.size_per_thread * octx->n_threads;
- octx->src2_spad.size = octx->src2_spad.size_per_thread * octx->n_threads;
- octx->src3_spad.size = octx->src3_spad.size_per_thread * octx->n_threads;
- octx->dst_spad.size = octx->dst_spad.size_per_thread * octx->n_threads;
+ uint8_t * vtcm_cur = octx->ctx->vtcm_base;
- size_t total_spad = octx->src0_spad.size + octx->src1_spad.size + octx->src2_spad.size + octx->src3_spad.size + octx->dst_spad.size;
+ factx.spad_q = vtcm_seq_alloc(&vtcm_cur, size_q_block * octx->n_threads);
+ factx.spad_k = vtcm_seq_alloc(&vtcm_cur, factx.size_k_block * 2 * octx->n_threads);
+ factx.spad_v = vtcm_seq_alloc(&vtcm_cur, factx.size_v_block * 2 * octx->n_threads);
+ factx.spad_m = vtcm_seq_alloc(&vtcm_cur, (mask ? factx.size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * octx->n_threads);
+ factx.spad_a = vtcm_seq_alloc(&vtcm_cur, size_vkq_acc * octx->n_threads);
- if (octx->ctx->vtcm_size < total_spad) {
+ if ((size_t) (vtcm_cur - octx->ctx->vtcm_base) > octx->ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL;
}
- octx->src0_spad.data = octx->ctx->vtcm_base; octx->src0_spad.src = NULL;
- octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL;
- octx->src2_spad.data = octx->src1_spad.data + octx->src1_spad.size; octx->src2_spad.src = NULL;
- octx->src3_spad.data = octx->src2_spad.data + octx->src2_spad.size; octx->src3_spad.src = NULL;
- octx->dst_spad.data = octx->src3_spad.data + octx->src3_spad.size; octx->dst_spad.src = NULL;
-
if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
worker_pool_run_func(octx->ctx->worker_pool, flash_attn_ext_f16_thread, &factx, octx->n_threads);
}
--- /dev/null
+#ifndef HTP_FLASH_ATTN_OPS_H
+#define HTP_FLASH_ATTN_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+
+#include "hex-fastdiv.h"
+#include "hex-common.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Tile constants (mirrored from hmx-utils.h for use on host side if needed)
+#define HMX_FP16_TILE_N_ROWS 32
+#define HMX_FP16_TILE_N_COLS 32
+#define HMX_FP16_TILE_N_ELMS 1024
+#define HMX_FP16_TILE_SIZE 2048
+#define HVX_FA_DMA_CACHE_SIZE 128
+#define HMX_FA_DMA_CACHE_SIZE 4
+
+#define HTP_FA_M_INITIAL_VAL -10000.0f
+
+enum htp_fa_kernel_type {
+ HTP_FA_KERNEL_UNSUPPORTED = 0,
+ HTP_FA_KERNEL_HVX,
+ HTP_FA_KERNEL_HMX
+};
+
+struct htp_fa_kernel_params {
+ uint8_t kernel_type; // enum htp_fa_kernel_type
+ uint8_t is_q_fp32; // 1 = Q type is F32, 0 = F16
+ uint8_t is_dst_fp32; // 1 = dst type is F32, 0 = F16
+ uint8_t n_threads; // Number of threads to run
+
+ // Common parameters
+ uint16_t Br;
+ uint16_t Bc;
+ uint16_t n_kv_blocks; // also HVX's n_blocks
+ uint16_t G; // GQA factor (n_heads / n_kv_heads)
+
+ float scale;
+ float max_bias;
+ float logit_softcap;
+ uint32_t vtcm_size;
+
+ uint32_t qrows;
+ uint32_t qrows_per_thread;
+ float m0;
+ float m1;
+ uint32_t n_head_log2;
+
+ struct fastdiv_values src3_div2;
+ struct fastdiv_values src3_div3;
+
+ union {
+ struct {
+ uint32_t g_br;
+ uint32_t row_buf_stride;
+ uint32_t mask_buf_row_stride;
+ int32_t mask_broadcast;
+ int32_t pipeline;
+ struct fastdiv_values div_G;
+ } hmx;
+ struct {
+ uint32_t size_q_row_padded;
+ uint32_t size_k_row_padded;
+ uint32_t size_v_row_padded;
+ struct fastdiv_values src0_div21;
+ struct fastdiv_values src0_div1;
+ struct fastdiv_values broadcast_rk2;
+ struct fastdiv_values broadcast_rk3;
+ struct fastdiv_values broadcast_rv2;
+ struct fastdiv_values broadcast_rv3;
+ } hvx;
+ } u;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
+#endif
+
+// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
+// g_br = hex_align_up(gqa_factor * Br, 32) replaces Br for all Q/O/S/P/D dimensions.
+// Layout: Q + O_ping + O_pong + K_dma*2 + V_dma*2 + K_tile + V_tile + S + P + D + vectors + scales
+// Mask is DMA'd into a VTCM buffer (Br rows per KV block) to avoid DDR reads in softmax.
+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) {
+ 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), 4096); // Q: [g_br, DK]
+ const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), 4096); // O: [g_br, DV] x2 ping-pong
+ const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 4096); // K DMA: [Bc, DK] x2 double-buf
+ const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 4096); // V DMA: [Bc, DV] x2 double-buf
+ const size_t k_tile_size = hex_align_up(Bc * DK * sizeof(__fp16), 4096); // K tiles: [Bc, DK] interleaved
+ const size_t v_tile_size = hex_align_up(Bc * DV * sizeof(__fp16), 4096); // V tiles: [Bc, DV] interleaved
+ const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), 4096); // S/P:[g_br, Bc]
+ const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), 4096); // D: [g_br, g_br]
+ const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256); // m, l, etc.
+ const size_t row_vec_size = hex_align_up(Bc * sizeof(__fp16), 256);
+ const size_t m_line_size = hex_align_up(Bc * sizeof(__fp16), 128);
+ const size_t m_buf_size = hex_align_up(Br * m_line_size, 4096) * HMX_FA_DMA_CACHE_SIZE;
+ const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128);
+
+ return q_tile_size * 1 // Q tiles
+ + o_tile_size * 2 // O ping-pong
+ + k_dma_size * 2 // K DMA x2
+ + v_dma_size * 2 // V DMA x2
+ + k_tile_size * 1 // K tiles
+ + v_tile_size * (pipeline ? 2 : 1) // V tiles (double-buffered if pipelining)
+ + s_tile_size * 2 // S + P
+ + d_tile_size * 1 // D (diagonal matrix)
+ + col_vec_size * 4 // m_vec, l_vec, s_rowmax, p_rowsum
+ + row_vec_size * 2 * n_threads // per-thread softmax row scratch
+ + m_buf_size * 1 // mask VTCM buffer [Br rows]
+ + slopes_size // Slopes
+ + 256 * 2; // HMX scales (id + qk)
+}
+
+#define FA_HVX_BLOCK_SIZE 64
+
+static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, size_t n_threads) {
+ const size_t size_q_row_padded = hex_round_up(DK * (is_q_fp32 ? 4 : 2), 128);
+ const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
+ const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
+
+ const size_t size_q_block = size_q_row_padded * 1;
+ const size_t size_k_block = size_k_row_padded * FA_HVX_BLOCK_SIZE;
+ const size_t size_v_block = size_v_row_padded * FA_HVX_BLOCK_SIZE;
+ const size_t size_m_block = hex_round_up(FA_HVX_BLOCK_SIZE * sizeof(__fp16), 128);
+ const size_t size_vkq_acc = hex_round_up(DV * sizeof(float), 128);
+
+ const size_t size_per_thread = size_q_block * 1
+ + size_k_block * 2
+ + size_v_block * 2
+ + (has_mask ? size_m_block * HVX_FA_DMA_CACHE_SIZE : 0)
+ + size_vkq_acc;
+
+ return size_per_thread * n_threads;
+}
+
+#define FA_MIN_KV_BLOCKS 3
+
+// Cost-based (Br, Bc) search for flash attention with pipeline constraint.
+static inline int hmx_fa_find_chunk_size(size_t * Br_out,
+ size_t * Bc_out,
+ size_t gqa_factor,
+ size_t DK,
+ size_t DV,
+ size_t qo_len,
+ size_t kv_len,
+ size_t vtcm_budget,
+ size_t n_threads) {
+ 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 fp16 = sizeof(__fp16);
+ const bool can_pipeline = (kv_len >= FA_MIN_KV_BLOCKS * bc_unit && n_threads >= 2);
+
+ // Approximate per-unit VTCM costs (without per-buffer alignment padding).
+ const size_t per_gbr = (DK + 2 * DV) * fp16 + 4 * sizeof(float); // Q + O*2 + 4 col vectors
+ const size_t per_gbr2 = fp16; // D diagonal matrix
+ const size_t per_bc =
+ 3 * DK * fp16 + (can_pipeline ? 4 : 3) * DV * fp16 + 2 * n_threads * fp16; // K/V DMA x2 + tiles + row bufs
+ const size_t per_gbr_bc = 2 * fp16; // S + P
+
+ const size_t overhead = 256 * 2 + 13 * 4096;
+
+ if (vtcm_budget <= overhead) {
+ return -1;
+ }
+ const size_t usable = vtcm_budget - overhead;
+
+ // Br_max: largest Br aligned to br_unit that does not exceed qo_len.
+ const size_t Br_max = qo_len >= br_unit ? hex_align_down(qo_len, br_unit) : br_unit;
+
+ // Pipeline constraint: cap Bc so n_kv_blocks >= FA_MIN_KV_BLOCKS.
+ // Only relax when kv_len is too short to form enough blocks.
+ 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
+
+ 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) {
+ const size_t g_br = hex_align_up(gqa_factor * Br, T);
+
+ // g_br-dependent VTCM cost: g_br * per_gbr + g_br*g_br * per_gbr2
+ const size_t gbr_cost = g_br * per_gbr + g_br * g_br * per_gbr2;
+ if (gbr_cost >= usable) {
+ if (Br == br_unit) {
+ break;
+ }
+ continue;
+ }
+
+ // Analytically solve for max Bc:
+ // remain >= Bc * (per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE)
+ // The Br * fp16 term accounts for the VTCM mask buffer [Br * Bc].
+ const size_t remain = usable - gbr_cost;
+ const size_t bc_denom = per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE;
+ size_t Bc = hex_smin(hex_align_down(remain / bc_denom, bc_unit), Bc_limit);
+ if (Bc < bc_unit) {
+ if (Br == br_unit) {
+ break;
+ }
+ continue;
+ }
+
+ // Exact VTCM verification (alignment padding may push over budget)
+ while (Bc >= bc_unit && hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline) > vtcm_budget) {
+ Bc -= bc_unit;
+ }
+ if (Bc < bc_unit) {
+ if (Br == br_unit) {
+ break;
+ }
+ continue;
+ }
+
+ 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;
+
+ if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
+ best_cost = cost;
+ best_mn = mn;
+ best_Br = Br;
+ best_Bc = Bc;
+ }
+
+ if (Br == br_unit) {
+ break;
+ }
+ }
+
+ if (best_Br == 0) {
+ return -1;
+ }
+
+ *Br_out = best_Br;
+ *Bc_out = best_Bc;
+ return 0;
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* HTP_FLASH_ATTN_OPS_H */
}
dma_descriptor_1d * desc = (dma_descriptor_1d *) &q->desc[q->push_idx];
- 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;
- desc->src = (void *) dptr.src;
- desc->dst = (void *) dptr.dst;
- desc->size = size;
+ 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->done = 1;
+ desc->desc_size = 0;
+ 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;
}
return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
}
-#define DMA_CACHE_MAX_SIZE 64U
+#define DMA_CACHE_MAX_SIZE 256U
typedef struct {
uint8_t *base;
if (c->src[i] == (uint32_t) src) {
c->age[i] = 0;
dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
- // FARF(ERROR, "dma-cache: found %p", src);
} else {
c->age[i]++;
if (c->age[i] > o_age) { o_age = c->age[i]; o_idx = i; }
}
}
if (!dst) {
- // FARF(ERROR, "dma-cache: replacing #%u : age %u %p -> %p", o_idx, c->age[o_idx], (void *) c->src[o_idx], src);
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(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
--- /dev/null
+#ifndef HMX_FA_KERNELS_H
+#define HMX_FA_KERNELS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include "hvx-utils.h"
+#include "hmx-utils.h"
+
+// HMX-specific parameters, offsets and inner kernels for Flash Attention
+
+// Scatter offsets for diagonal tile: entry[2i] = i*136, entry[2i+1] = i*136+6
+// 136 = 4 * 32 + 8 = byte offset to diagonal in a 32x32 fp16 interleaved tile
+static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) = {
+ 0 * 136, 0 * 136 + 6,
+ 1 * 136, 1 * 136 + 6,
+ 2 * 136, 2 * 136 + 6,
+ 3 * 136, 3 * 136 + 6,
+ 4 * 136, 4 * 136 + 6,
+ 5 * 136, 5 * 136 + 6,
+ 6 * 136, 6 * 136 + 6,
+ 7 * 136, 7 * 136 + 6,
+ 8 * 136, 8 * 136 + 6,
+ 9 * 136, 9 * 136 + 6,
+ 10 * 136, 10 * 136 + 6,
+ 11 * 136, 11 * 136 + 6,
+ 12 * 136, 12 * 136 + 6,
+ 13 * 136, 13 * 136 + 6,
+ 14 * 136, 14 * 136 + 6,
+ 15 * 136, 15 * 136 + 6,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+ 0, 0,
+};
+// Inner HMX tile computation kernels
+
+static inline void hmx_fa_qk_dot_tile(
+ const __fp16 * row_tiles,
+ const __fp16 * col_tiles,
+ __fp16 * out_tile,
+ size_t n_dot_tiles
+) {
+ for (size_t k = 0; k < n_dot_tiles; ++k) {
+ Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047);
+ Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047);
+ row_tiles += HMX_FP16_TILE_N_ELMS;
+ col_tiles += HMX_FP16_TILE_N_ELMS;
+ }
+ Q6_mxmem_AR_after_hf(out_tile, 0);
+}
+
+static inline void hmx_fa_o_update_tile(
+ const __fp16 * d_diag,
+ const __fp16 * o_rc,
+ const __fp16 * p_tile_in,
+ const __fp16 * v_tile_in,
+ __fp16 * o_tile_out,
+ size_t n_col_tiles
+) {
+ Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
+ Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
+
+ for (size_t k = 0; k < n_col_tiles; ++k) {
+ Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047);
+ Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047);
+ p_tile_in += HMX_FP16_TILE_N_ELMS;
+ v_tile_in += HMX_FP16_TILE_N_ELMS;
+ }
+
+ Q6_mxmem_AR_after_hf(o_tile_out, 0);
+}
+
+static inline void hmx_fa_o_norm_tile(
+ const __fp16 * d_diag,
+ const __fp16 * o_rc,
+ __fp16 * o_out
+) {
+ Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
+ Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
+ Q6_mxmem_AR_after_hf(o_out, 0);
+}
+
+#endif /* HMX_FA_KERNELS_H */
+++ /dev/null
-// HMX-accelerated Flash Attention for prefill (neq1 >= 32).
-// Ported from htp-ops-lib/src/dsp/ops/flash_attn.c, adapted to the htp/ codebase.
-
-#pragma clang diagnostic ignored "-Wunused-variable"
-#pragma clang diagnostic ignored "-Wunused-function"
-#pragma clang diagnostic ignored "-Wunused-but-set-variable"
-
-#include <assert.h>
-#include <HAP_compute_res.h>
-#include <HAP_farf.h>
-#include <math.h>
-#include <stdbool.h>
-#include <stddef.h>
-#include <stdint.h>
-#include <string.h>
-
-#define GGML_COMMON_DECL_C
-#include "ggml-common.h"
-#include "hex-dma.h"
-#include "hex-fastdiv.h"
-#include "hex-profile.h"
-#include "hmx-queue.h"
-#include "hmx-utils.h"
-#include "htp-ctx.h"
-#include "htp-ops.h"
-#include "hvx-dump.h"
-#include "hvx-copy.h"
-#include "hvx-reduce.h"
-#include "hvx-utils.h"
-#include "hvx-flash-attn.h"
-#include "vtcm-utils.h"
-#include "worker-pool.h"
-
-// ============================================================================
-// Constants
-// ============================================================================
-
-// Tile constants from hmx-utils.h
-// HMX_FP16_TILE_N_ROWS = 32
-// HMX_FP16_TILE_N_COLS = 32
-// HMX_FP16_TILE_N_ELMS = 1024
-// HMX_FP16_TILE_SIZE = 2048
-
-// ============================================================================
-// Dynamic block size computation (GQA-aware)
-// ============================================================================
-
-// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-// g_br = hex_align_up(gqa_factor * Br, 32) replaces Br for all Q/O/S/P/D dimensions.
-// Layout: Q + O_ping + O_pong + K_dma*2 + V_dma*2 + K_tile + V_tile + S + P + D + vectors + scales
-// Mask is DMA'd into a VTCM buffer (Br rows per KV block) to avoid DDR reads in softmax.
-static 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) {
- 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), 4096); // Q: [g_br, DK]
- const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), 4096); // O: [g_br, DV] x2 ping-pong
- const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 4096); // K DMA: [Bc, DK] x2 double-buf
- const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 4096); // V DMA: [Bc, DV] x2 double-buf
- const size_t k_tile_size = hex_align_up(Bc * DK * sizeof(__fp16), 4096); // K tiles: [Bc, DK] interleaved
- const size_t v_tile_size = hex_align_up(Bc * DV * sizeof(__fp16), 4096); // V tiles: [Bc, DV] interleaved
- const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), 4096); // S/P:[g_br, Bc]
- const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), 4096); // D: [g_br, g_br]
- const size_t col_vec_size = hex_align_up(g_br * sizeof(__fp16), 256); // m, l, etc.
- const size_t row_vec_size = hex_align_up(Bc * sizeof(__fp16), 256);
- const size_t m_line_size = hex_align_up(Bc * sizeof(__fp16), 128);
- const size_t m_buf_size = hex_align_up(Br * m_line_size, 4096);
- const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128);
-
- return q_tile_size * 1 // Q tiles
- + o_tile_size * 2 // O ping-pong
- + k_dma_size * 2 // K DMA x2
- + v_dma_size * 2 // V DMA x2
- + k_tile_size * 1 // K tiles
- + v_tile_size * (pipeline ? 2 : 1) // V tiles (double-buffered if pipelining)
- + s_tile_size * 2 // S + P
- + d_tile_size * 1 // D (diagonal matrix)
- + col_vec_size * 4 // m_vec, l_vec, s_rowmax, p_rowsum
- + row_vec_size * 2 * n_threads // per-thread softmax row scratch
- + m_buf_size * 1 // mask VTCM buffer [Br rows]
- + slopes_size // Slopes
- + 256 * 2; // HMX scales (id + qk)
-}
-
-// ============================================================================
-// FP16 exp2 polynomial (ported from htp-ops-lib/include/dsp/hvx_math.h)
-// ============================================================================
-// 5th-order Horner polynomial for exp2(x) in qf16/hf16 domain. Input must be
-// ≤ 0 (safe softmax invariant — overflow handling omitted). ~18 ALU ops per
-// 64 fp16 lanes, fully parallel across HVX threads (no scatter/gather engine).
-// Replaces the F32 round-trip (qf16→f32→exp→f32→f16, ~44 ops for 2×32 lanes).
-static inline HVX_Vector hvx_exp2_hf(HVX_Vector x_v) {
- const HVX_Vector zero_v = Q6_V_vzero();
- const HVX_Vector half_hf_v = Q6_Vh_vsplat_R(0x3800); // fp16 0.5
-
- // k = round_toward_neg_inf(x); f = (float)k; frac = x - f
- HVX_Vector x_minus_half = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(x_v, half_hf_v));
- HVX_Vector k_v = Q6_Vh_equals_Vhf(x_minus_half); // truncate to int16
- HVX_Vector f_v = Q6_Vhf_equals_Vh(k_v); // back to fp16
-
- HVX_Vector x_qf16 = Q6_Vqf16_vsub_VhfVhf(x_v, f_v); // fractional part in qf16
-
- // Horner: y = ((((E5*x + E4)*x + E3)*x + E2)*x + E1)*x + E0
- HVX_Vector y = Q6_Vqf16_vmpy_Vqf16Vqf16(Q6_Vh_vsplat_R(0x5082), x_qf16); // E5*x
- y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x157d)); // + E4
- y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
- y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x20ed)); // + E3
- y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
- y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x2b1b)); // + E2
- y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
- y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x33b0)); // + E1
- y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
- y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x398c)); // + E0
- y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16); // y = y * x
- y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x3c00)); // + 1.0
-
- // Combine polynomial (mantissa) with integer part (exponent): result = y * 2^k
- y = Q6_Vhf_equals_Vqf16(y);
- HVX_Vector y_exp = Q6_Vuh_vlsr_VuhR(Q6_Vh_vasl_VhR(y, 1), 11);
- y_exp = Q6_Vh_vadd_VhVh(k_v, y_exp);
- HVX_VectorPred q_underflow = Q6_Q_vcmp_gt_VhVh(zero_v, y_exp);
- y = Q6_Vh_vaslacc_VhVhR(y, k_v, 10);
- return Q6_V_vmux_QVV(q_underflow, zero_v, y);
-}
-
-#define FA_MIN_KV_BLOCKS 3
-
-// Cost-based (Br, Bc) search for flash attention with pipeline constraint.
-//
-// VTCM model (same as before):
-// overhead + g_br * per_gbr + g_br² * per_gbr2 + Bc * per_bc + g_br * Bc * per_gbr_bc
-//
-// Cost model (minimization objective):
-// Q * (c_q_fixed + K * c_iter_fixed), where Q = ceil(qo/Br), K = ceil(kv/Bc)
-static int hmx_fa_find_chunk_size(size_t * Br_out,
- size_t * Bc_out,
- size_t gqa_factor,
- size_t DK,
- size_t DV,
- size_t qo_len,
- size_t kv_len,
- size_t vtcm_budget,
- size_t n_threads) {
- const size_t T = HMX_FP16_TILE_N_ROWS; // 32
- const size_t br_unit = hmx_ceil_div(T, gqa_factor);
- // Bc must be a multiple of 64 so that n_tiles_per_bc is even. The softmax
- // P-tile write uses a dual-tile pattern (vshuff + two stores 16 slots apart)
- // that would race across r0 blocks if the last dual-tile is half-occupied.
- // See .cursor/todos/hmx-flash-attn-bc-search-space.md for the perf trade-off.
- const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64
- const size_t fp16 = sizeof(__fp16);
- const bool can_pipeline = (kv_len >= FA_MIN_KV_BLOCKS * bc_unit && n_threads >= 2);
-
- // Approximate per-unit VTCM costs (without per-buffer alignment padding).
- const size_t per_gbr = (DK + 2 * DV) * fp16 + 4 * fp16; // Q + O×2 + 4 col vectors
- const size_t per_gbr2 = fp16; // D diagonal matrix
- const size_t per_bc =
- 3 * DK * fp16 + (can_pipeline ? 4 : 3) * DV * fp16 + 2 * n_threads * fp16; // K/V DMA x2 + tiles + row bufs
- const size_t per_gbr_bc = 2 * fp16; // S + P
-
- const size_t overhead = 256 * 2 + 13 * 4096;
-
- if (vtcm_budget <= overhead) {
- return -1;
- }
- const size_t usable = vtcm_budget - overhead;
-
- // Br_max: largest Br aligned to br_unit that does not exceed qo_len.
- const size_t Br_max = qo_len >= br_unit ? hex_align_down(qo_len, br_unit) : br_unit;
-
- // Pipeline constraint: cap Bc so n_kv_blocks >= FA_MIN_KV_BLOCKS.
- // Only relax when kv_len is too short to form enough blocks.
- 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
-
- 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) {
- const size_t g_br = hex_align_up(gqa_factor * Br, T);
-
- // g_br-dependent VTCM cost: g_br * per_gbr + g_br² * per_gbr2
- const size_t gbr_cost = g_br * per_gbr + g_br * g_br * per_gbr2;
- if (gbr_cost >= usable) {
- if (Br == br_unit) {
- break;
- }
- continue;
- }
-
- // Analytically solve for max Bc:
- // remain >= Bc * (per_bc + g_br * per_gbr_bc + Br * fp16_mask)
- // The Br * fp16 term accounts for the VTCM mask buffer [Br × Bc].
- const size_t remain = usable - gbr_cost;
- const size_t bc_denom = per_bc + g_br * per_gbr_bc + Br * fp16;
- size_t Bc = hex_smin(hex_align_down(remain / bc_denom, bc_unit), Bc_limit);
- if (Bc < bc_unit) {
- if (Br == br_unit) {
- break;
- }
- continue;
- }
-
- // Exact VTCM verification (alignment padding may push over budget)
- while (Bc >= bc_unit && hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline) > vtcm_budget) {
- Bc -= bc_unit;
- }
- if (Bc < bc_unit) {
- if (Br == br_unit) {
- break;
- }
- continue;
- }
-
- 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;
-
- if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
- best_cost = cost;
- best_mn = mn;
- best_Br = Br;
- best_Bc = Bc;
- }
-
- if (Br == br_unit) {
- break;
- }
- }
-
- if (best_Br == 0) {
- return -1;
- }
-
- *Br_out = best_Br;
- *Bc_out = best_Bc;
- return 0;
-}
-
-// ============================================================================
-// Tile interleave / extract helpers
-// ============================================================================
-
-// transpose scatter offsets moved to hmx-utils.h as hmx_transpose_scatter_offsets
-
-// Scatter offsets for diagonal tile: entry[2i] = i*136, entry[2i+1] = i*136+6
-// 136 = 4 * 32 + 8 = byte offset to diagonal in a 32x32 fp16 interleaved tile
-static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) = {
- 0 * 136, 0 * 136 + 6,
- 1 * 136, 1 * 136 + 6,
- 2 * 136, 2 * 136 + 6,
- 3 * 136, 3 * 136 + 6,
- 4 * 136, 4 * 136 + 6,
- 5 * 136, 5 * 136 + 6,
- 6 * 136, 6 * 136 + 6,
- 7 * 136, 7 * 136 + 6,
- 8 * 136, 8 * 136 + 6,
- 9 * 136, 9 * 136 + 6,
- 10 * 136, 10 * 136 + 6,
- 11 * 136, 11 * 136 + 6,
- 12 * 136, 12 * 136 + 6,
- 13 * 136, 13 * 136 + 6,
- 14 * 136, 14 * 136 + 6,
- 15 * 136, 15 * 136 + 6,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
- 0, 0,
-};
-
-// hmx_interleave_rows_to_tiles and hmx_interleave_cols_to_tiles are in hmx-utils.h
-
-// ============================================================================
-// HMX Flash Attention context (GQA-merged)
-// ============================================================================
-
-struct hmx_fa_context {
- const struct htp_ops_context * octx;
- bool pipeline; // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2
- uint32_t n_threads;
-
- // Op parameters
- float scale;
- float max_bias;
- float logit_softcap;
- uint32_t n_head_log2;
- float m0, m1;
-
- // Dimensions
- uint32_t DK, DV;
- uint32_t n_kv; // kv_len
- uint32_t n_kv_heads; // number of KV heads
- uint32_t n_heads; // number of Q heads
- uint32_t G; // GQA factor = n_heads / n_kv_heads
- struct fastdiv_values div_G;
- uint32_t n_kv_blocks;
- uint32_t neq1; // Q token count
-
- // Types
- bool is_q_fp32;
- bool is_dst_fp32;
-
- // Dynamic block sizes
- uint32_t Br; // Q tokens per block (before GQA expansion)
- uint32_t Bc;
- uint32_t g_br; // hex_align_up(G * Br, 32) - actual tile row dim
-
- // VTCM buffers (allocated by vtcm_seq_alloc)
- __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_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_d_tiles; // Diagonal rescale [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]
- HVX_Vector * vtcm_p_rowsum; // Softmax intermediate [g_br]
- HVX_Vector * vtcm_row_bufs; // Per-thread softmax row scratch [n_threads][2][Bc/64]
- uint8_t * vtcm_hmx_scales_id; // HMX output scales (identity)
- uint8_t * vtcm_hmx_scales_qk; // HMX output scales (qk_scale)
- __fp16 * vtcm_mask_buf; // VTCM mask buffer [Br × m_line], DMA'd per KV block
- __fp16 * vtcm_slopes; // ALiBi slopes [g_br]
- size_t row_buf_stride; // HVX vectors per row buffer (Bc/64)
- size_t mask_buf_row_stride; // elements (__fp16) per row in mask buffer
- bool mask_broadcast; // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
-};
-
-// ============================================================================
-// Multi-thread K interleave phase
-// ============================================================================
-
-typedef struct {
- struct hmx_fa_context * factx;
- int kv_rows;
- size_t src_stride;
- size_t buf_idx;
-} fa_k_int_args_t;
-
-static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
- fa_k_int_args_t * args = (fa_k_int_args_t *) data;
- struct hmx_fa_context * factx = args->factx;
-
- const int total_rows = args->kv_rows;
- const int rows_per_t = hex_align_up(hmx_ceil_div(total_rows, n), 2); // ensure even (row pairs)
- const int start = i * rows_per_t;
- const int end = hex_smin(start + rows_per_t, total_rows);
-
- if (start >= total_rows) {
- return;
- }
-
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start);
- hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, factx->vtcm_k_fp16[args->buf_idx], total_rows, (int) factx->DK,
- (int) args->src_stride, start, end);
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start);
-}
-
-static void fa_phase_k_interleave(struct hmx_fa_context * factx, int kv_rows, size_t src_stride, size_t buf_idx) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
- fa_k_int_args_t args = { factx, kv_rows, src_stride, buf_idx };
- if (factx->n_threads > 1 && kv_rows >= (int) (factx->n_threads * 2)) {
- worker_pool_run_func(wp, fa_k_interleave_thread, &args, factx->n_threads);
- } else {
- fa_k_interleave_thread(1, 0, &args);
- }
-}
-
-// ============================================================================
-// Multi-thread V interleave phase
-// ============================================================================
-
-typedef struct {
- struct hmx_fa_context * factx;
- int kv_rows;
- size_t src_stride;
- size_t buf_idx;
- size_t n_col_tiles;
-} fa_v_int_args_t;
-
-static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) {
- fa_v_int_args_t * args = (fa_v_int_args_t *) data;
- struct hmx_fa_context * factx = args->factx;
-
- const int total_rows = args->kv_rows;
- const int rows_per_t = hex_align_up(hmx_ceil_div(total_rows, n), 2);
- const int start = i * rows_per_t;
- const int end = hex_smin(start + rows_per_t, total_rows);
-
- if (start >= total_rows) {
- return;
- }
-
- __fp16 * v_tiles_dest = factx->pipeline ? factx->vtcm_v_tiles[args->buf_idx] : factx->vtcm_v_tiles[0];
-
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start);
- hmx_interleave_cols_to_tiles(v_tiles_dest, factx->vtcm_v_fp16[args->buf_idx], total_rows, (int) factx->DV,
- (int) args->src_stride, (int) args->n_col_tiles, start, end);
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start);
-}
-
-static void fa_phase_v_interleave(struct hmx_fa_context * factx,
- int kv_rows,
- size_t src_stride,
- size_t buf_idx,
- size_t n_col_tiles) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
- fa_v_int_args_t args = { factx, kv_rows, src_stride, buf_idx, n_col_tiles };
- if (factx->n_threads > 1 && kv_rows >= (int) (factx->n_threads * 2)) {
- worker_pool_run_func(wp, fa_v_interleave_thread, &args, factx->n_threads);
- } else {
- fa_v_interleave_thread(1, 0, &args);
- }
-}
-
-// ============================================================================
-// Multi-thread Q load phase: read Q[G × neq1, DK] from DDR, convert F32→F16
-// (or deal F16 pairs), and write interleaved into vtcm_q_tiles.
-// Each thread owns a disjoint range of row pairs; writes target distinct tile
-// slots (r0 selects tile row, r1 selects intra-tile slot), so there is no
-// write conflict. Padding fill (when n_rows_g < g_br) is done single-threaded
-// by the caller before dispatching.
-// ============================================================================
-
-typedef struct {
- struct hmx_fa_context * factx;
- const struct htp_tensor * q;
- uint32_t q_start;
- uint32_t kv_head;
- uint32_t ib3;
- size_t n_rows_g;
-} fa_q_load_args_t;
-
-static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
- fa_q_load_args_t * args = (fa_q_load_args_t *) data;
- struct hmx_fa_context * factx = args->factx;
-
- const size_t n_rows_g = args->n_rows_g;
- const size_t G = factx->G;
- const size_t DK = factx->DK;
-
- // Partition row pairs across threads. Keep each thread's start even so r/r+1
- // are always in the same thread's range.
- const size_t rows_per_t = hex_align_up(hmx_ceil_div(n_rows_g, n), 2);
- const size_t start = (size_t) i * rows_per_t;
- const size_t end = hex_smin(start + rows_per_t, n_rows_g);
-
- if (start >= n_rows_g) {
- return;
- }
-
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start);
-
- const struct htp_tensor * q = args->q;
- const uint32_t q_start = args->q_start;
- const uint32_t kv_head = args->kv_head;
- const uint32_t ib3 = args->ib3;
-
- for (size_t r = start; r < end; r += 2) {
- const bool next_row_valid = (r + 1) < n_rows_g;
-
- const size_t q_idx0 = fastdiv(r + 0, &factx->div_G);
- const size_t h_idx0 = fastmodulo(r + 0, G, &factx->div_G);
- const size_t q_idx1 = fastdiv(r + 1, &factx->div_G);
- const size_t h_idx1 = fastmodulo(r + 1, G, &factx->div_G);
-
- const uint8_t * q_ptr0 = (const uint8_t *) q->data + (q_start + q_idx0) * q->nb[1] +
- (kv_head * G + h_idx0) * q->nb[2] + ib3 * q->nb[3];
- const uint8_t * q_ptr1 = next_row_valid ? ((const uint8_t *) q->data + (q_start + q_idx1) * q->nb[1] +
- (kv_head * G + h_idx1) * q->nb[2] + ib3 * q->nb[3]) :
- NULL;
-
- size_t r0 = r / HMX_FP16_TILE_N_ROWS;
- size_t r1 = r % HMX_FP16_TILE_N_ROWS;
- __fp16 * out_base = factx->vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
-
- if (factx->is_q_fp32) {
- const HVX_Vector * pv_in0 = (const HVX_Vector *) q_ptr0;
- const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
-
- for (uint32_t d = 0; d < DK / 32; ++d) {
- HVX_Vector v0 = pv_in0[d];
- HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
- HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
-
- HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
- out_tile[r1 / 2] = v_hf;
- }
- } else {
- const HVX_Vector * pv_in0 = (const HVX_Vector *) q_ptr0;
- const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
-
- for (uint32_t d = 0; d < DK / 64; ++d) {
- HVX_Vector v0 = pv_in0[d];
- HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
- HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
-
- __fp16 * out_dual_tile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
- HVX_Vector * pv_out0 = ((HVX_Vector *) out_dual_tile) + r1 / 2;
- HVX_Vector * pv_out1 = pv_out0 + 16;
-
- *pv_out0 = Q6_V_lo_W(vp);
- *pv_out1 = Q6_V_hi_W(vp);
- }
- }
- }
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start);
-}
-
-static void fa_phase_q_load(struct hmx_fa_context * factx,
- const struct htp_tensor * q,
- uint32_t q_start,
- uint32_t kv_head,
- uint32_t ib3,
- size_t n_rows_g) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
- fa_q_load_args_t args = { factx, q, q_start, kv_head, ib3, n_rows_g };
- // Require >= 2 row pairs per thread so partitioning is worthwhile.
- if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
- worker_pool_run_func(wp, fa_q_load_thread, &args, factx->n_threads);
- } else {
- fa_q_load_thread(1, 0, &args);
- }
-}
-
-// ============================================================================
-// Multi-thread O store phase: read O tiles from VTCM, convert F16->F32 (or
-// deal F16 pairs), and write to strided DDR dst tensor. Each thread owns a
-// disjoint row range; writes target distinct dst rows (different q_idx/h_idx
-// pairs produced by r/G and r%G), so there is no write conflict.
-// ============================================================================
-
-typedef struct {
- struct hmx_fa_context * factx;
- const struct htp_tensor * dst;
- const __fp16 * o_tile_src;
- uint32_t q_start;
- uint32_t kv_head;
- uint32_t ib3;
- size_t n_rows_g;
-} fa_o_store_args_t;
-
-static void fa_o_store_thread(unsigned int n, unsigned int i, void * data) {
- fa_o_store_args_t * args = (fa_o_store_args_t *) data;
- struct hmx_fa_context * factx = args->factx;
-
- const size_t n_rows_g = args->n_rows_g;
- const size_t G = factx->G;
- const size_t DV = factx->DV;
-
- const size_t rows_per_t = hmx_ceil_div(n_rows_g, n);
- const size_t start = (size_t) i * rows_per_t;
- const size_t end = hex_smin(start + rows_per_t, n_rows_g);
-
- if (start >= n_rows_g) {
- return;
- }
-
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start);
-
- const struct htp_tensor * dst = args->dst;
- const __fp16 * o_tile_src = args->o_tile_src;
- const uint32_t q_start = args->q_start;
- 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);
-
- // FIX(dst-indexing): ggml_flash_attn_ext() creates dst as permute(0,2,1,3) ->
- // [DV, n_heads, n_tokens, n_seq], so head stride is nb[1] and token stride is nb[2].
- uint8_t * dst_row = (uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
- (q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3];
-
- size_t r0 = r / HMX_FP16_TILE_N_ROWS;
- size_t r1 = r % HMX_FP16_TILE_N_ROWS;
- const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
-
- if (factx->is_dst_fp32) {
- float * out = (float *) dst_row;
- for (uint32_t d = 0; d < DV / 32; ++d) {
- const HVX_Vector * in_tile = (const HVX_Vector *) (tile_row_base + d * HMX_FP16_TILE_N_ELMS);
- HVX_VectorPair vp = hvx_vec_f16_to_f32_shuff(in_tile[r1 / 2]);
- if (r1 % 2 == 0) {
- *(HVX_UVector *) (out + d * 32) = Q6_V_lo_W(vp);
- } else {
- *(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
- }
- }
- } else {
- __fp16 * out = (__fp16 *) dst_row;
- for (uint32_t d = 0; d < DV / 64; ++d) {
- const __fp16 * in_dual_tile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
- const HVX_Vector * pv_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
- const HVX_Vector * pv_in1 = pv_in0 + 16;
- HVX_VectorPair vp = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
- if (r1 % 2 == 0) {
- *(HVX_UVector *) (out + d * 64) = Q6_V_lo_W(vp);
- } else {
- *(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
- }
- }
- }
- }
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start);
-}
-
-static void fa_phase_o_store(struct hmx_fa_context * factx,
- const struct htp_tensor * dst,
- const __fp16 * o_tile_src,
- uint32_t q_start,
- uint32_t kv_head,
- uint32_t ib3,
- size_t n_rows_g) {
- worker_pool_context_t wp = factx->octx->ctx->worker_pool;
- fa_o_store_args_t args = { factx, dst, o_tile_src, q_start, kv_head, ib3, n_rows_g };
- if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
- worker_pool_run_func(wp, fa_o_store_thread, &args, factx->n_threads);
- } else {
- fa_o_store_thread(1, 0, &args);
- }
-}
-
-// ============================================================================
-// Multi-thread softmax phase + serial m/l update + build_D
-// ============================================================================
-
-typedef struct {
- struct hmx_fa_context * factx;
- size_t kv_rows;
- size_t n_rows_g;
- size_t n_col_tiles;
- size_t n_tiles_per_bc;
- size_t n_row_tiles;
- size_t n_row_tiles_g_br;
- uint32_t Bc;
- uint32_t G;
- uint32_t kv_head;
- uint32_t kv_start;
- uint32_t q_start;
- uint32_t ib3;
- bool has_alibi; // true when max_bias != 0 (need slope * mask + add)
-
- // ALiBi per-head slopes (indexed by GQA-merged row: slope[r] for r in [0, n_rows_g))
- // slope[r] = 1.0 when max_bias == 0 (no ALiBi)
- // Pointer into hmx_fa_context.vtcm_slopes (sized to g_br)
- __fp16 * slopes;
-
- // Mask info (preloaded before softmax)
- const struct htp_tensor * mask;
- const __fp16 * mask_vtcm; // VTCM mask buffer base (NULL = DDR fallback)
- size_t mask_vtcm_row_stride; // elements (__fp16) per row in VTCM mask buffer
-} fa_softmax_args_t;
-
-static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) {
- fa_softmax_args_t * args = (fa_softmax_args_t *) data;
- struct hmx_fa_context * factx = args->factx;
-
- const size_t n_rows_g = args->n_rows_g;
- const size_t kv_rows = args->kv_rows;
- const size_t Bc = args->Bc;
- const size_t G = args->G;
- const size_t n_tiles_per_bc = args->n_tiles_per_bc;
- const size_t n_row_vec_cnt = hmx_ceil_div(n_rows_g, 64);
-
- // Partition r_vec_idx across threads
- const size_t vecs_per_t = hmx_ceil_div(n_row_vec_cnt, n);
- const size_t vec_start = i * vecs_per_t;
- const size_t vec_end = hex_smin(vec_start + vecs_per_t, n_row_vec_cnt);
-
- if (vec_start >= n_row_vec_cnt) {
- return;
- }
-
- struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, vec_start);
-
- // Per-thread row scratch: thread i uses bufs at offset i * 2 * stride
- const size_t row_buf_stride = factx->row_buf_stride;
- HVX_Vector * my_row_buf0 = factx->vtcm_row_bufs + i * 2 * row_buf_stride;
- HVX_Vector * my_row_buf1 = my_row_buf0 + row_buf_stride;
-
- const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
-
- // Per-row accumulators: each fp16 lane in a 64-lane vector holds one row's scalar.
- // CONTRACT: lane bits must be IEEE fp16 (hf), never qf16 — qf16 uses a different
- // bit layout, so a later hf-domain read would silently produce wrong values.
- // Convert first via Q6_Vhf_equals_Vqf16(). For reference: vtcm_m_vec/vtcm_s_rowmax
- // are hf; vtcm_l_vec is qf16 — don't mix them up.
-
- for (size_t r_vec_idx = vec_start; r_vec_idx < vec_end; ++r_vec_idx) {
- HVX_Vector rowmax_acc_v = v_neg_inf;
- HVX_Vector rowsum_acc_v = Q6_V_vzero();
- HVX_Vector m_prev_v = factx->vtcm_m_vec[r_vec_idx];
-
- for (int r_vec_off = 0; r_vec_off < 64; r_vec_off += 2) {
- int r = r_vec_idx * 64 + r_vec_off;
- if (r >= (int) hex_align_up(n_rows_g, 2)) {
- break;
- }
-
- int r0 = r / HMX_FP16_TILE_N_ROWS;
- int 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;
-
- // Decode 2 rows from S tiles into per-thread row buffers
- HVX_Vector * pv_row_buf0 = my_row_buf0;
- HVX_Vector * pv_row_buf1 = my_row_buf1;
- for (size_t c = 0; c < kv_rows; c += 64) {
- const __fp16 * in_dual_tile = s_ld_base + (c / 64) * HMX_FP16_TILE_N_ELMS * 2;
- const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
- const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
-
- HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
- *pv_row_buf0++ = Q6_V_lo_W(vp_s_dual_row);
- *pv_row_buf1++ = Q6_V_hi_W(vp_s_dual_row);
- }
-
- // Apply softcap if enabled (in F32 precision)
- if (factx->logit_softcap != 0.0f) {
- // When EXP2_HF is on, fold log2(e) into v_cap so the output lands in
- // log2(e)-scaled space for the downstream exp2. log2(e) is kept OUT
- // of qk_scale in this configuration (see scale setup) so tanh sees
- // the physical QK/(√d·c) argument.
- float cap = factx->logit_softcap;
-#ifdef HMX_FA_USE_EXP2_HF
- cap *= 1.44269504f; // log2(e)
-#endif
- const HVX_Vector v_cap = hvx_vec_splat_f32(cap);
- for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 64;
-
- HVX_VectorPair r0_f32 = hvx_vec_f16_to_f32(my_row_buf0[ci]);
- HVX_Vector t0_lo = hvx_vec_tanh_f32(Q6_V_lo_W(r0_f32));
- HVX_Vector t0_hi = hvx_vec_tanh_f32(Q6_V_hi_W(r0_f32));
- t0_lo = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(t0_lo, v_cap));
- t0_hi = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(t0_hi, v_cap));
- my_row_buf0[ci] = hvx_vec_f32_to_f16(t0_lo, t0_hi);
-
- HVX_VectorPair r1_f32 = hvx_vec_f16_to_f32(my_row_buf1[ci]);
- HVX_Vector t1_lo = hvx_vec_tanh_f32(Q6_V_lo_W(r1_f32));
- HVX_Vector t1_hi = hvx_vec_tanh_f32(Q6_V_hi_W(r1_f32));
- t1_lo = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(t1_lo, v_cap));
- t1_hi = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(t1_hi, v_cap));
- my_row_buf1[ci] = hvx_vec_f32_to_f16(t1_lo, t1_hi);
- }
- }
-
- // Apply mask & compute rowmax(S)
- //
- // Optimizations over baseline:
- // A. No-ALiBi fast path: when max_bias==0 (slope≡1.0), skip the
- // slope multiplication — still add mask (additive bias) but
- // avoid the mul_f16_f16. Saves 2 ops/dual-row vs ALiBi path.
- // B. GQA mask row dedup: G consecutive Q rows share one mask row
- // (qi = r / G). Reuse mask vector when qi is unchanged between
- // row0 and row1 (saves ~75% of VTCM loads for G=4).
-
- // ALiBi slopes — only needed when has_alibi (scheme A)
- HVX_Vector v_slope0, v_slope1;
- if (args->has_alibi) {
- HVX_Vector v_s = hvx_vmemu(args->slopes + r);
- v_slope0 = hvx_vec_repl_f16(v_s);
- v_slope1 = (r + 1 < (int) n_rows_g) ? hvx_vec_repl_f16(Q6_V_vror_VR(v_s, 2)) : Q6_V_vzero();
- }
-
- const HVX_Vector v_threshold = Q6_Vh_vsplat_R(0xcc00); // fp16 -16.0 (hoisted outside for-c)
-
- 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 (args->mask) {
- HVX_Vector v_mask0, v_mask1;
-
- if (args->mask_vtcm) {
- // Read mask from VTCM buffer (DMA'd per KV block).
- // GQA dedup (scheme B): skip load when qi unchanged.
- const size_t qi0 = fastdiv(r + 0, &factx->div_G);
- v_mask0 = *(const HVX_UVector *) (args->mask_vtcm + qi0 * args->mask_vtcm_row_stride + c);
- v_mask1 = v_neg_inf;
- if (r + 1 < (int) n_rows_g) {
- const size_t qi1 = fastdiv(r + 1, &factx->div_G);
- if (qi1 == qi0) {
- v_mask1 = v_mask0; // scheme B: reuse — same mask row
- } else {
- v_mask1 = *(const HVX_UVector *) (args->mask_vtcm + qi1 * args->mask_vtcm_row_stride + c);
- }
- }
- } else {
- // Fallback: read mask directly from DDR (when mask->ne[2] > 1).
- const struct htp_tensor * mask = args->mask;
- const size_t q_idx0 = args->q_start + fastdiv(r + 0, &factx->div_G);
- const size_t h_idx0 = args->kv_head * G + fastmodulo(r + 0, G, &factx->div_G);
- const uint32_t im2_0 = h_idx0 % mask->ne[2];
- const uint32_t im3_0 = args->ib3 % mask->ne[3];
-
- const __fp16 * m0_ptr = (const __fp16 *) ((const uint8_t *) mask->data + q_idx0 * mask->nb[1] +
- im2_0 * mask->nb[2] + im3_0 * mask->nb[3]) + args->kv_start + c;
- v_mask0 = *(const HVX_UVector *) m0_ptr;
- v_mask1 = v_neg_inf;
-
- if (r + 1 < (int) n_rows_g) {
- const size_t q_idx1 = args->q_start + fastdiv(r + 1, &factx->div_G);
- if (q_idx1 == q_idx0) {
- // scheme B: same mask row in DDR path
- v_mask1 = v_mask0;
- } else {
- const size_t h_idx1 = args->kv_head * G + fastmodulo(r + 1, G, &factx->div_G);
- const uint32_t im2_1 = h_idx1 % mask->ne[2];
- const uint32_t im3_1 = args->ib3 % mask->ne[3];
- const __fp16 * m1_ptr = (const __fp16 *) ((const uint8_t *) mask->data + q_idx1 * mask->nb[1] +
- im2_1 * mask->nb[2] + im3_1 * mask->nb[3]) + args->kv_start + c;
- v_mask1 = *(const HVX_UVector *) m1_ptr;
- }
- }
- }
-
- // Threshold: mask values below -16.0 are treated as -inf (causal mask).
- HVX_VectorPred q_keep0 = Q6_Q_and_QQ(Q6_Q_vcmp_gt_VhfVhf(v_mask0, v_threshold), q_tail_keep);
- HVX_VectorPred q_keep1 = Q6_Q_and_QQ(Q6_Q_vcmp_gt_VhfVhf(v_mask1, v_threshold), q_tail_keep);
-
- if (args->has_alibi) {
- // ALiBi path: S += slope * mask (full mul + add)
- HVX_Vector v_sm0 = hvx_vec_mul_f16_f16(v_mask0, v_slope0);
- HVX_Vector v_sm1 = hvx_vec_mul_f16_f16(v_mask1, v_slope1);
- my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_sm0), v_neg_inf);
- my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_sm1), v_neg_inf);
- } else {
- // No-ALiBi fast path (scheme A): slope≡1.0, skip the mul
- // but still add mask (additive positional bias). vmux
- // clamps mask < -16 to -inf as a numerical safeguard.
- my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_mask0), v_neg_inf);
- my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_mask1), v_neg_inf);
- }
- } else {
- 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 = hvx_vec_reduce_max_f16(v_s_rowmax0);
- v_s_rowmax1 = hvx_vec_reduce_max_f16(v_s_rowmax1);
-
- // Splat m_prev[r], m_prev[r+1] from the per-row accumulator.
- // vror brings the target lane to lane 0, then vdelta replicates it
- // across all lanes — stays in the vector domain (no store/reload).
- HVX_Vector v_m_prev0 = hvx_vec_repl_f16(Q6_V_vror_VR(m_prev_v, r_vec_off * 2));
- HVX_Vector v_m_prev1 = hvx_vec_repl_f16(Q6_V_vror_VR(m_prev_v, (r_vec_off + 1) * 2));
-
- // HVX max — both operands are splats, so result is splat of m_new.
- HVX_Vector v_dup_m0 = Q6_Vhf_vmax_VhfVhf(v_m_prev0, v_s_rowmax0);
- HVX_Vector v_dup_m1 = Q6_Vhf_vmax_VhfVhf(v_m_prev1, v_s_rowmax1);
-
- // Insert row r, r+1 rowmax into rowmax_acc_v via 2-byte-wide vmux.
- // Byte ranges: lane0 = [r_vec_off*2 .. r_vec_off*2+1], lane1 shifted by 2.
- // vsetq2 handles the n=128 corner case when r_vec_off reaches 62.
- {
- HVX_VectorPred p_start = Q6_Q_vsetq_R(r_vec_off * 2);
- HVX_VectorPred p_mid = Q6_Q_vsetq_R((r_vec_off + 1) * 2);
- HVX_VectorPred p_end = Q6_Q_vsetq2_R((r_vec_off + 2) * 2);
- HVX_VectorPred p_lane0 = Q6_Q_and_QQn(p_mid, p_start);
- HVX_VectorPred p_lane1 = Q6_Q_and_QQn(p_end, p_mid);
- rowmax_acc_v = Q6_V_vmux_QVV(p_lane0, v_dup_m0, rowmax_acc_v);
- rowmax_acc_v = Q6_V_vmux_QVV(p_lane1, v_dup_m1, rowmax_acc_v);
- }
-
- // Compute P = exp(S - m_new), using HVX exp
- const HVX_Vector v_zero = Q6_V_vzero();
- HVX_Vector v_p_rowsum0 = v_zero;
- HVX_Vector v_p_rowsum1 = v_zero;
-
-#ifdef HMX_FA_USE_EXP2_HF
- // FP16 exp2 polynomial path (matches htp-ops-lib flash_attn.c):
- // P = exp2(S - m_new)
- for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 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);
-
- HVX_Vector v_p_row0_hf = hvx_exp2_hf(Q6_Vhf_equals_Vqf16(v_s_minus_m0));
- HVX_Vector v_p_row1_hf = hvx_exp2_hf(Q6_Vhf_equals_Vqf16(v_s_minus_m1));
-#else
- // F32 exp path: qf16 → f32 → exp → f32 → f16. Higher precision,
- for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 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);
-
- HVX_VectorPair vp0 = hvx_vec_f16_to_f32_shuff(Q6_Vhf_equals_Vqf16(v_s_minus_m0));
- HVX_Vector p0_lo = hvx_vec_exp_f32(Q6_V_lo_W(vp0));
- HVX_Vector p0_hi = hvx_vec_exp_f32(Q6_V_hi_W(vp0));
- HVX_Vector v_p_row0_hf = hvx_vec_f32_to_f16_shuff(p0_lo, p0_hi);
-
- HVX_VectorPair vp1 = hvx_vec_f16_to_f32_shuff(Q6_Vhf_equals_Vqf16(v_s_minus_m1));
- HVX_Vector p1_lo = hvx_vec_exp_f32(Q6_V_lo_W(vp1));
- HVX_Vector p1_hi = hvx_vec_exp_f32(Q6_V_hi_W(vp1));
- HVX_Vector v_p_row1_hf = hvx_vec_f32_to_f16_shuff(p1_lo, p1_hi);
-#endif
- // Write P to tile format. Dual-tile pattern assumes Bc is a
- // multiple of 64 (enforced by bc_unit=64 in hmx_fa_find_chunk_size),
- // so both tile halves are always in the current r0 block.
- __fp16 * out_dual_tile = p_st_base + (c / 64) * HMX_FP16_TILE_N_ELMS * 2;
- HVX_Vector * pv_p_out0 = ((HVX_Vector *) out_dual_tile) + r1 / 2;
- HVX_Vector * pv_p_out1 = pv_p_out0 + 16;
-
- HVX_VectorPair vp_p_dual = Q6_W_vshuff_VVR(v_p_row1_hf, v_p_row0_hf, -2);
- *pv_p_out0 = Q6_V_lo_W(vp_p_dual);
- *pv_p_out1 = Q6_V_hi_W(vp_p_dual);
-
- HVX_VectorPair vp_p0 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf);
- HVX_VectorPair vp_p1 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf);
-
- v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0), Q6_V_hi_W(vp_p0)));
- v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1), Q6_V_hi_W(vp_p1)));
- }
-
- HVX_Vector rowsum0_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(v_p_rowsum0));
- HVX_Vector rowsum1_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(v_p_rowsum1));
- {
- // Both inputs are f32 splats, so the f32->f16 output is an fp16 splat.
- HVX_Vector rv0_v = hvx_vec_f32_to_f16(rowsum0_sf, rowsum0_sf);
- HVX_Vector rv1_v = hvx_vec_f32_to_f16(rowsum1_sf, rowsum1_sf);
-
- HVX_VectorPred p_start = Q6_Q_vsetq_R(r_vec_off * 2);
- HVX_VectorPred p_mid = Q6_Q_vsetq_R((r_vec_off + 1) * 2);
- HVX_VectorPred p_end = Q6_Q_vsetq2_R((r_vec_off + 2) * 2);
- HVX_VectorPred p_lane0 = Q6_Q_and_QQn(p_mid, p_start);
- HVX_VectorPred p_lane1 = Q6_Q_and_QQn(p_end, p_mid);
- rowsum_acc_v = Q6_V_vmux_QVV(p_lane0, rv0_v, rowsum_acc_v);
- rowsum_acc_v = Q6_V_vmux_QVV(p_lane1, rv1_v, rowsum_acc_v);
- }
- }
-
- factx->vtcm_s_rowmax[r_vec_idx] = rowmax_acc_v;
- factx->vtcm_p_rowsum[r_vec_idx] = rowsum_acc_v;
- }
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, vec_start);
-}
-
-// Serial m/l update + build_D. Must run after softmax barrier (s_rowmax written by all threads).
-//
-// noinline: function boundary acts as a hard compiler barrier so the (size_t)addr scatter
-// intrinsics inside cannot be hoisted past the call site. Mirrors the structural protection
-// matmul gets for free via worker_pool function-pointer dispatch. Without this, the compiler
-// can reorder the scatter past the subsequent hmx_queue_push and the HMX-queue worker thread
-// reads stale VTCM (PPL → ~vocab-size).
-static __attribute__((noinline)) void fa_ml_update_and_build_d(struct hmx_fa_context * factx,
- size_t n_rows_g,
- size_t n_row_tiles,
- size_t n_row_tiles_g_br) {
- // Reuse s_rowmax buffer for exp(m_diff) — safe because softmax is fully complete
- HVX_Vector * const mvec_exp_m_diff = factx->vtcm_s_rowmax;
-
- const size_t n_row_vec_cnt = hmx_ceil_div(n_rows_g, 64);
- for (size_t i = 0; i < n_row_vec_cnt; ++i) {
- HVX_Vector v_m_prev = factx->vtcm_m_vec[i];
- HVX_Vector v_m_curr = Q6_Vhf_vmax_VhfVhf(v_m_prev, factx->vtcm_s_rowmax[i]);
- HVX_Vector v_m_diff = Q6_Vqf16_vsub_VhfVhf(v_m_prev, v_m_curr);
-
-#ifdef HMX_FA_USE_EXP2_HF
- // Base-2 path: must match P = exp2(S - m_new) in fa_softmax_thread.
- HVX_Vector v_exp_m_diff = hvx_exp2_hf(Q6_Vhf_equals_Vqf16(v_m_diff));
-#else
- HVX_VectorPair vp_diff = hvx_vec_f16_to_f32_shuff(Q6_Vhf_equals_Vqf16(v_m_diff));
- HVX_Vector exp_lo = hvx_vec_exp_f32(Q6_V_lo_W(vp_diff));
- HVX_Vector exp_hi = hvx_vec_exp_f32(Q6_V_hi_W(vp_diff));
- HVX_Vector v_exp_m_diff = hvx_vec_f32_to_f16_shuff(exp_lo, exp_hi);
-#endif
-
- HVX_Vector v_l_curr = Q6_Vqf16_vmpy_Vqf16Vhf(factx->vtcm_l_vec[i], v_exp_m_diff);
- v_l_curr = Q6_Vqf16_vadd_Vqf16Vhf(v_l_curr, factx->vtcm_p_rowsum[i]);
-
- factx->vtcm_m_vec[i] = v_m_curr;
- factx->vtcm_l_vec[i] = v_l_curr;
- mvec_exp_m_diff[i] = v_exp_m_diff;
- }
-
- // Build diagonal tile D = diag(exp(m_diff))
- const HVX_Vector v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets;
- const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16));
- for (size_t i = 0; i < n_row_tiles; ++i) {
- const HVX_Vector v_content = Q6_V_vror_VR(mvec_exp_m_diff[i / 2], (i % 2) * 64);
- __fp16 * out_base = factx->vtcm_d_tiles + 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);
- // Compiler barrier — Q6_vscatter takes (size_t)addr; without this the
- // compiler may not recognize the volatile read below as aliasing and
- // could reorder it before the scatter, defeating the HW drain.
- __asm__ __volatile__("" ::: "memory");
- // Per-tile drain: scatter regions are disjoint (stride > tile size),
- // so a single drain at tile 0 does NOT retire later tiles' entries.
- (void) *(volatile HVX_Vector *) out_base;
- }
-}
-
-// Build D = diag(1/l) tile for the final O = D @ O normalization.
-//
-// noinline: same rationale as fa_ml_update_and_build_d — keeps Q6_vscatter from
-// being hoisted past the subsequent hmx_queue_push at the o_norm call site.
-static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_context * factx,
- size_t n_row_tiles,
- size_t n_row_tiles_g_br) {
- const HVX_Vector v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets;
- const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16));
- const HVX_Vector one = hvx_vec_splat_f32(1.0f);
-
- HVX_Vector v_content = Q6_V_vzero();
- for (size_t i = 0; i < n_row_tiles; ++i) {
- if ((i % 2) == 0) {
- HVX_Vector v_l_hf = Q6_Vhf_equals_Vqf16(factx->vtcm_l_vec[i / 2]);
- HVX_VectorPair vp_l = hvx_vec_f16_to_f32_shuff(v_l_hf);
- HVX_Vector inv_lo = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(one, hvx_vec_inverse_f32(Q6_V_lo_W(vp_l))));
- HVX_Vector inv_hi = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(one, hvx_vec_inverse_f32(Q6_V_hi_W(vp_l))));
- v_content = hvx_vec_f32_to_f16_shuff(inv_lo, inv_hi);
- } else {
- 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;
- Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
- // Compiler barrier — see fa_ml_update_and_build_d for rationale.
- __asm__ __volatile__("" ::: "memory");
- (void) *(volatile HVX_Vector *) out_base;
- }
-}
-
-// Combined: multi-thread softmax -> barrier -> serial m/l update + build_D
-static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx,
- 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;
- const size_t n_row_vec_cnt = hmx_ceil_div(sargs->n_rows_g, 64);
-
- 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);
- worker_pool_run_func(wp, fa_softmax_thread, sargs, n_use);
- } else {
- fa_softmax_thread(1, 0, sargs);
- }
- // barrier implicit in worker_pool_run_func return
-
- fa_ml_update_and_build_d(factx, sargs->n_rows_g, n_row_tiles, n_row_tiles_g_br);
-}
-
-// ============================================================================
-// HMX job structs and worker functions
-// ============================================================================
-
-typedef struct {
- const __fp16 * q_tiles;
- const __fp16 * k_tiles;
- __fp16 * s_tiles;
- size_t n_row_tiles;
- size_t n_col_tiles;
- size_t n_dot_tiles; // DK / 32
- size_t n_tiles_per_bc;
- uint8_t * hmx_scales;
-} hmx_fa_qk_job_t;
-
-static void hmx_fa_qk_dot_worker(void * data) {
- hmx_fa_qk_job_t * job = (hmx_fa_qk_job_t *) data;
- const size_t n_row_tiles = job->n_row_tiles;
- const size_t n_col_tiles = job->n_col_tiles;
- const size_t n_dot_tiles = job->n_dot_tiles;
- const size_t n_tiles_per_bc = job->n_tiles_per_bc;
- const __fp16 * restrict q_tiles = job->q_tiles;
- const __fp16 * restrict k_tiles = job->k_tiles;
- __fp16 * restrict s_tiles = job->s_tiles;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(n_col_tiles > 0);
- __builtin_assume(n_dot_tiles > 0);
-
- Q6_bias_mxmem2_A((void *) job->hmx_scales);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < n_col_tiles; ++c) {
- const __fp16 * row_tiles = q_tiles + r * HMX_FP16_TILE_N_ROWS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
- const __fp16 * col_tiles = k_tiles + c * HMX_FP16_TILE_N_COLS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
- __fp16 * out_tile = s_tiles + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
-
- for (size_t k = 0; k < n_dot_tiles; ++k) {
- Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047);
- row_tiles += HMX_FP16_TILE_N_ELMS;
- col_tiles += HMX_FP16_TILE_N_ELMS;
- }
- Q6_mxmem_AR_after_hf(out_tile, 0);
- }
- }
-}
-
-typedef struct {
- __fp16 * o_curr;
- const __fp16 * o_prev;
- const __fp16 * p_tiles;
- const __fp16 * v_tiles;
- const __fp16 * d_tiles;
- uint8_t * hmx_scales;
- size_t n_row_tiles;
- size_t n_col_tiles;
- size_t n_row_tiles_g_br;
- size_t n_tiles_per_bc;
- size_t DV;
-} hmx_fa_o_update_job_t;
-
-static void hmx_fa_o_update_worker(void * data) {
- hmx_fa_o_update_job_t * job = (hmx_fa_o_update_job_t *) data;
- const size_t n_row_tiles = job->n_row_tiles;
- const size_t n_col_tiles = job->n_col_tiles;
- const size_t n_row_tiles_g_br = job->n_row_tiles_g_br;
- const size_t n_tiles_per_bc = job->n_tiles_per_bc;
- const size_t DV_tiles = job->DV / 32;
- const __fp16 * restrict d_tiles = job->d_tiles;
- const __fp16 * restrict p_tiles = job->p_tiles;
- const __fp16 * restrict v_tiles = job->v_tiles;
- const __fp16 * restrict o_prev = job->o_prev;
- __fp16 * restrict o_curr = job->o_curr;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(n_col_tiles > 0);
- __builtin_assume(DV_tiles > 0);
-
- Q6_bias_mxmem2_A((void *) job->hmx_scales);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- // D[r,r] @ O_prev[r,c] — only the diagonal tile
- const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
-
- // P @ V (accumulate on same accumulator)
- const __fp16 * p_tile_in = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * v_tile_in = v_tiles + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- for (size_t k = 0; k < n_col_tiles; ++k) {
- Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047);
- p_tile_in += HMX_FP16_TILE_N_ELMS;
- v_tile_in += HMX_FP16_TILE_N_ELMS;
- }
-
- __fp16 * o_tile_out = o_curr + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- Q6_mxmem_AR_after_hf(o_tile_out, 0);
- }
- }
-}
-
-typedef struct {
- __fp16 * o_curr; // output (row-major tile layout)
- const __fp16 * o_prev; // input (column-major tile layout)
- const __fp16 * d_tiles; // diag(1/l) tiles
- uint8_t * hmx_scales;
- size_t n_row_tiles;
- size_t n_row_tiles_g_br;
- size_t DV;
-} hmx_fa_o_norm_job_t;
-
-static void hmx_fa_o_norm_worker(void * data) {
- hmx_fa_o_norm_job_t * job = (hmx_fa_o_norm_job_t *) data;
- const size_t n_row_tiles = job->n_row_tiles;
- const size_t n_row_tiles_g_br = job->n_row_tiles_g_br;
- const size_t DV_tiles = job->DV / 32;
- const __fp16 * restrict d_tiles = job->d_tiles;
- const __fp16 * restrict o_prev = job->o_prev;
- __fp16 * restrict o_curr = job->o_curr;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(DV_tiles > 0);
-
- Q6_bias_mxmem2_A((void *) job->hmx_scales);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- __fp16 * o_out = o_curr + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
-
- Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
- Q6_mxmem_AR_after_hf(o_out, 0);
- }
- }
-}
-
-// Populate per-GQA-row ALiBi slopes for a given KV head.
-// Row r in the GQA-merged block maps to Q head h = kv_head * G + r % G.
-// slope(h) = m0^(h+1) when h < n_head_log2, else m1^(2*(h-n_head_log2)+1).
-// When max_bias == 0, all slopes are 1.0 (no ALiBi).
-static __attribute__((noinline)) void fa_compute_slopes(
- const struct hmx_fa_context * factx,
- uint32_t kv_head,
- size_t n_rows_g) {
- __fp16 * slopes = factx->vtcm_slopes;
- if (factx->max_bias == 0.0f) {
- hvx_splat_f16_a(slopes, 1.0f, n_rows_g);
- return;
- }
-
- const uint32_t G = factx->G;
- const uint32_t n_head_log2 = factx->n_head_log2;
- const float m0 = factx->m0;
- const float m1 = factx->m1;
-
- __fp16 temp_slopes[512] __attribute__((aligned(128)));
- if (G <= 32) {
- // Fast path: Compute G unique slope values in vector registers
- HVX_Vector v_val = hvx_alibi_slopes(kv_head, G, n_head_log2, m0, m1);
-
- __fp16 temp_slopes_aligned[64] __attribute__((aligned(128)));
- hvx_vmem(temp_slopes_aligned) = hvx_vec_f32_to_f16(v_val, Q6_V_vzero());
-
- for (uint32_t i = 0; i < G; ++i) {
- temp_slopes[i] = temp_slopes_aligned[i];
- }
- } else {
- // Fallback path: G > 32 (rare configurations)
- for (uint32_t i = 0; i < G; ++i) {
- temp_slopes[i] = (__fp16)alibi_slope(kv_head * G + i, n_head_log2, m0, m1);
- }
- }
-
- // Allocate stack buffer to avoid scalar writes to VTCM (which generates L2 misses)
- __fp16 local_slopes[n_rows_g] __attribute__((aligned(128)));
- for (size_t r = 0; r < n_rows_g; ++r) {
- local_slopes[r] = temp_slopes[fastmodulo(r, G, &factx->div_G)];
- }
-
- // Copy to VTCM slopes using HVX block copy (both are aligned to 128 bytes)
- hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)local_slopes, n_rows_g);
-}
-
-// ============================================================================
-// Core HMX flash attention algorithm (GQA-merged)
-// ============================================================================
-
-int hmx_flash_attn_ext(struct htp_ops_context * octx) {
- struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[HTP_MAX_NTHREADS] : NULL;
- const struct htp_tensor * q = octx->src[0];
- const struct htp_tensor * k = octx->src[1];
- const struct htp_tensor * v = octx->src[2];
- const struct htp_tensor * mask = (octx->src[3] && octx->src[3]->data) ? octx->src[3] : NULL;
- const struct htp_tensor * dst = octx->dst;
-
- struct htp_context * const ctx = octx->ctx;
-
- if (!ctx->hmx_enabled) {
- return HTP_STATUS_NO_SUPPORT;
- }
-
- // Dimensions
- const uint32_t neq0 = q->ne[0]; // head_dim (DK)
- const uint32_t neq1 = q->ne[1]; // n_tokens
- const uint32_t neq2 = q->ne[2]; // n_heads
- const uint32_t neq3 = q->ne[3]; // n_seqs
-
- const uint32_t nek0 = k->ne[0]; // head_dim
- const uint32_t nek1 = k->ne[1]; // kv_len
-
- const uint32_t nev0 = v->ne[0]; // head_dim (DV)
-
- const uint32_t DK = neq0;
- const uint32_t DV = nev0;
-
- // HMX requires head_dim to be multiple of 32
- if (DK % 32 != 0 || DV % 32 != 0) {
- return HTP_STATUS_NO_SUPPORT;
- }
-
- // GQA factor
- const uint32_t n_kv_heads = k->ne[2];
- const uint32_t G = neq2 / n_kv_heads;
-
- // Thread count for multi-thread HVX phases
- const uint32_t n_threads_init = octx->n_threads;
-
- // Compute dynamic block sizes (GQA-aware, accounting for per-thread row bufs)
- size_t Br, Bc;
- const size_t vtcm_budget = ctx->vtcm_size;
- if (hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, vtcm_budget, n_threads_init) != 0) {
- return HTP_STATUS_VTCM_TOO_SMALL;
- }
-
- const size_t g_br = hex_align_up(G * Br, HMX_FP16_TILE_N_ROWS);
-
- const uint32_t n_kv_blocks = (nek1 + Bc - 1) / Bc;
- const bool pipeline = (n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads_init >= 2);
-
- // Bypass thread pool dispatch for small prompts/non-pipelined prefill by setting n_threads = 1
- const uint32_t n_threads = pipeline ? n_threads_init : 1;
-
- FARF(HIGH, "hmx-fa: neq1=%u nek1=%u DK=%u DV=%u G=%u Br=%zu Bc=%zu g_br=%zu n_kv_blocks=%u pipeline=%d vtcm=%zu",
- neq1, nek1, DK, DV, G, Br, Bc, g_br, n_kv_blocks, pipeline, vtcm_budget);
-
- // ======== Build context ========
- struct hmx_fa_context factx;
- memset(&factx, 0, sizeof(factx));
- factx.octx = octx;
- factx.n_threads = n_threads;
- factx.DK = DK;
- factx.DV = DV;
- factx.n_kv = nek1;
- factx.n_kv_heads = n_kv_heads;
- factx.n_heads = neq2;
- factx.G = G;
- factx.div_G = init_fastdiv_values(G);
- factx.neq1 = neq1;
- factx.Br = (uint32_t) Br;
- factx.Bc = (uint32_t) Bc;
- factx.g_br = (uint32_t) g_br;
- factx.n_kv_blocks = n_kv_blocks;
- factx.is_q_fp32 = (q->type == HTP_TYPE_F32);
- factx.is_dst_fp32 = (dst->type == HTP_TYPE_F32);
- factx.pipeline = pipeline;
- factx.mask_broadcast = (mask != NULL && mask->ne[2] == 1);
-
- // Extract op parameters (mutable during softcap adjustment, then stored as const in factx)
- float scale = 1.0f, max_bias = 0.0f, logit_softcap = 0.0f;
- memcpy(&scale, (float *) octx->op_params + 0, sizeof(float));
- memcpy(&max_bias, (float *) octx->op_params + 1, sizeof(float));
- memcpy(&logit_softcap, (float *) octx->op_params + 2, sizeof(float));
-
- if (logit_softcap != 0.0f) {
- scale /= logit_softcap;
- }
-
-#ifdef HMX_FA_USE_EXP2_HF
- // Pre-bake log2(e) into qk_scale so HMX-produced S tiles are in log2(e)-scaled
- // space. Then exp2(S - m) in the softmax equals base-e exp((S - m) / log2(e)),
- // preserving ggml's base-e softmax semantics. Matches htp-ops-lib flash_attn.c.
- //
- // When softcap is active we cannot pre-bake log2(e) here — it would land inside
- // the tanh argument and shift the softcap knee from x≈c to x≈c/log2(e), giving
- // numerically wrong softcapped values. Instead fold log2(e) into the post-tanh
- // multiplier (see softcap block: v_cap absorbs log2(e)).
- if (logit_softcap == 0.0f) {
- scale *= 1.44269504f; // log2(e)
- }
-#endif
-
- factx.scale = scale;
- factx.max_bias = max_bias;
- factx.logit_softcap = logit_softcap;
-
- factx.n_head_log2 = 1u << (uint32_t) floor(log2(neq2));
- factx.m0 = powf(2.0f, -(max_bias) / factx.n_head_log2);
- factx.m1 = powf(2.0f, -(max_bias / 2.0f) / factx.n_head_log2);
-
- // ======== VTCM allocation (GQA-aware) ========
- const size_t size_k_row = DK * sizeof(__fp16);
- const size_t size_v_row = DV * sizeof(__fp16);
- const size_t size_k_row_padded = hex_round_up(size_k_row, 128);
- const size_t size_v_row_padded = hex_round_up(size_v_row, 128);
-
- const size_t q_tile_bytes = hex_align_up(g_br * DK * sizeof(__fp16), 4096);
- const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
- const size_t k_dma_bytes = hex_align_up(Bc * size_k_row_padded, 4096);
- const size_t v_dma_bytes = hex_align_up(Bc * size_v_row_padded, 4096);
- const size_t k_tile_bytes = hex_align_up(Bc * DK * sizeof(__fp16), 4096);
- const size_t v_tile_bytes = hex_align_up(Bc * DV * sizeof(__fp16), 4096);
- const size_t s_tile_bytes = hex_align_up(g_br * Bc * sizeof(__fp16), 4096);
- const size_t d_tile_bytes = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
- const size_t col_vec_bytes = hex_align_up(g_br * sizeof(__fp16), 256);
- const size_t row_vec_bytes = hex_align_up(Bc * sizeof(__fp16), 256);
- const size_t m_line_bytes = hex_align_up(Bc * sizeof(__fp16), 128);
- const size_t m_buf_bytes = hex_align_up(Br * m_line_bytes, 4096);
- const size_t slopes_bytes = hex_align_up(g_br * sizeof(__fp16), 128);
-
- uint8_t * vtcm_cur = ctx->vtcm_base;
-
- factx.vtcm_q_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, q_tile_bytes);
- factx.vtcm_o_tiles[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
- factx.vtcm_o_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
- factx.vtcm_k_fp16[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
- factx.vtcm_k_fp16[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
- factx.vtcm_v_fp16[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
- factx.vtcm_v_fp16[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
- factx.vtcm_k_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_tile_bytes);
- factx.vtcm_v_tiles[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
- if (pipeline) {
- factx.vtcm_v_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
- } else {
- factx.vtcm_v_tiles[1] = NULL;
- }
- factx.vtcm_s_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
- factx.vtcm_p_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
- factx.vtcm_d_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, d_tile_bytes);
- factx.vtcm_m_vec = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_l_vec = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_s_rowmax = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_p_rowsum = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_row_bufs = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, row_vec_bytes * 2 * n_threads);
- factx.row_buf_stride = row_vec_bytes / sizeof(HVX_Vector);
- factx.vtcm_hmx_scales_id = vtcm_seq_alloc(&vtcm_cur, 256);
- factx.vtcm_hmx_scales_qk = vtcm_seq_alloc(&vtcm_cur, 256);
- factx.vtcm_mask_buf = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, m_buf_bytes);
- factx.mask_buf_row_stride = m_line_bytes / sizeof(__fp16);
- factx.vtcm_slopes = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, slopes_bytes);
-
- if ((size_t) (vtcm_cur - ctx->vtcm_base) > ctx->vtcm_size) {
- return HTP_STATUS_VTCM_TOO_SMALL;
- }
-
- // ======== Initialize HMX output scales ========
- // Identity scale (1.0) for O updates and normalization
- hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0
-
- // QK scale embedded in HMX output
- hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale));
-
- // ======== Skip compute if profiling ========
- if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
- return HTP_STATUS_OK;
- }
-
-
- // ======== DMA setup ========
- dma_queue * const dma = ctx->dma[0];
-
- // Padded row sizes for DMA (defined in outer scope)
-
- const size_t n_row_tiles_g_br = g_br / HMX_FP16_TILE_N_ROWS;
- const size_t n_tiles_per_bc = Bc / HMX_FP16_TILE_N_COLS;
-
- // Q/O element size for Q load and O store
- const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
-
- // ======== HMX lock strategy ========
- // Pipeline: queue thread auto-acquires HMX lock on first push; released by suspend.
- // Fallback: main thread holds the lock (original behavior).
- if (!factx.pipeline) {
- HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
- }
-
- // ======== Reusable job descriptors for pipeline ========
- hmx_fa_qk_job_t qk_job;
- hmx_fa_o_update_job_t ou_job;
- hmx_fa_o_norm_job_t on_job;
-
- // ======== Main loop: per batch, per KV head, per Q block ========
- for (uint32_t ib3 = 0; ib3 < neq3; ++ib3) {
- for (uint32_t kv_head = 0; kv_head < n_kv_heads; ++kv_head) {
- const uint32_t ik2 = kv_head;
- const uint32_t ik3 = ib3 / (neq3 / k->ne[3]);
- const uint32_t iv2 = kv_head;
- const uint32_t iv3 = ib3 / (neq3 / v->ne[3]);
-
- for (uint32_t q_start = 0; q_start < neq1; q_start += Br) {
- const uint32_t n_q_rows = hex_smin(Br, neq1 - q_start);
- const size_t n_rows_g = n_q_rows * G;
- const size_t g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
- const size_t n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
-
- // ---- Load Q block [g_br, D] -> tiles, interleaving G heads ----
- if (n_rows_g < g_br) {
- hvx_splat_u8_a(factx.vtcm_q_tiles, 0, q_tile_bytes);
- }
- fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
-
- // ---- Initialize per-block state ----
- hvx_splat_u8_a(factx.vtcm_l_vec, 0, col_vec_bytes);
- hvx_splat_u8_a(factx.vtcm_d_tiles, 0, d_tile_bytes);
- hvx_splat_u16_a(factx.vtcm_m_vec, 0xfbff, col_vec_bytes/2);
-
- __fp16 * o_tile_prev = factx.vtcm_o_tiles[0];
- __fp16 * o_tile_curr = factx.vtcm_o_tiles[1];
- hvx_splat_u8_a(o_tile_prev, 0, o_tile_bytes);
-
- // ---- KV block loop with DMA double-buffering ----
- size_t buf_idx = 0;
-
- fa_compute_slopes(&factx, kv_head, n_rows_g);
-
- // Prefetch first KV block
- if (factx.n_kv_blocks > 0) {
- const uint32_t kv_rows0 = hex_smin(Bc, nek1);
-
- 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);
- }
-
- // Mask DMA: single 2D transfer of n_q_rows unique mask rows into VTCM buffer.
- // Only when mask is head-broadcast (ne[2]==1); otherwise softmax reads DDR directly.
- #define MASK_DMA_PUSH(kv_start_val, kv_rows_val, has_mask_dma_var) \
- do { \
- has_mask_dma_var = false; \
- if (mask && factx.mask_broadcast) { \
- const uint32_t _im3 = ib3 % mask->ne[3]; \
- const uint8_t * _ms = (const uint8_t *) mask->data + q_start * mask->nb[1] + _im3 * mask->nb[3] + \
- (kv_start_val) * sizeof(__fp16); \
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_mask_buf, _ms), m_line_bytes, mask->nb[1], \
- (kv_rows_val) * sizeof(__fp16), n_q_rows); \
- has_mask_dma_var = true; \
- } \
- } while (0)
-
- #define MASK_DMA_POP(has_mask_dma_var) \
- do { \
- if (has_mask_dma_var) { \
- dma_queue_pop(dma); \
- } \
- } while (0)
-
- #define DMA_PREFETCH_KV(blk_val) \
- do { \
- if ((blk_val) < factx.n_kv_blocks) { \
- const uint32_t _ns = (blk_val) * Bc; \
- const uint32_t _nr = hex_smin(Bc, nek1 - _ns); \
- size_t _nb = 1 - buf_idx; \
- const uint8_t * _ks = (const uint8_t *) k->data + _ns * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3]; \
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[_nb], _ks), size_k_row_padded, k->nb[1], size_k_row, _nr); \
- const uint8_t * _vs = (const uint8_t *) v->data + _ns * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3]; \
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[_nb], _vs), size_v_row_padded, v->nb[1], size_v_row, _nr); \
- } \
- } while (0)
-
- const size_t k_src_stride = size_k_row_padded / sizeof(__fp16);
- const size_t v_src_stride = size_v_row_padded / sizeof(__fp16);
-
- if (factx.pipeline) {
- // ==================================================================
- // Pipeline path: HVX phases ‖ HMX queue worker
- // ==================================================================
- struct hmx_queue * hmx_q = ctx->hmx_queue;
-
- 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);
-
- // Wait for current KV DMA
- dma_queue_pop(dma); // K
- dma_queue_pop(dma); // V
-
- // Push mask DMA for this block (single 2D DMA when broadcast)
- bool has_mask_dma = false;
- MASK_DMA_PUSH(kv_start, kv_rows, has_mask_dma);
-
- // ---- Phase 1: K_int(blk) ‖ O_update(blk-1) ----
- if (kv_blk > 0) {
- // Submit O_update for previous block (HMX worker)
- 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));
- }
- fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx);
-
- // ---- Phase 2: qk_dot(blk) on HMX ‖ V_int(blk) + DMA prefetch on HVX ----
- 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));
-
- // DMA push next block (non-blocking, before worker_pool)
- DMA_PREFETCH_KV(kv_blk + 1);
- fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc);
-
- // Pop and swap previous block's output update (deferred HMX pop)
- if (kv_blk > 0) {
- hmx_queue_pop(hmx_q);
- hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
- }
-
- // Pop current block's dot product job
- hmx_queue_pop(hmx_q);
-
- // ---- Phase 3: softmax(blk) + build_D(blk) | HMX idle ----
- // Pop mask DMA before softmax (ensures VTCM buffer is ready)
- MASK_DMA_POP(has_mask_dma);
-
- fa_softmax_args_t sargs;
- memset(&sargs, 0, sizeof(sargs));
- sargs.factx = &factx;
- sargs.kv_rows = kv_rows;
- sargs.n_rows_g = n_rows_g;
- sargs.n_col_tiles = n_col_tiles;
- sargs.n_tiles_per_bc = n_tiles_per_bc;
- sargs.n_row_tiles = n_row_tiles;
- sargs.n_row_tiles_g_br = n_row_tiles_g_br;
- sargs.Bc = Bc;
- sargs.G = G;
- sargs.kv_head = kv_head;
- sargs.kv_start = kv_start;
- sargs.q_start = q_start;
- sargs.ib3 = ib3;
- sargs.has_alibi = (factx.max_bias != 0.0f);
- sargs.mask = mask;
- sargs.mask_vtcm = has_mask_dma ? (const __fp16 *) factx.vtcm_mask_buf : NULL;
- sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
- sargs.slopes = factx.vtcm_slopes;
- fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
-
- buf_idx = 1 - buf_idx;
- } // end KV block loop (pipeline)
-
- // Epilogue: O_update for last block
- 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));
- hmx_queue_pop(hmx_q);
-
- hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
- }
-
- } else {
- // ==================================================================
- // Fallback path: sequential with multi-thread HVX phases
- // Main thread holds HMX lock, runs HMX inline.
- // ==================================================================
-
- 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);
- dma_queue_pop(dma); // K
- dma_queue_pop(dma); // V
-
- bool has_mask_dma = false;
- MASK_DMA_PUSH(kv_start, kv_rows, has_mask_dma);
- DMA_PREFETCH_KV(kv_blk + 1);
- fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx);
-
- {
- const size_t n_dot_tiles = (size_t) (DK / 32);
- const __fp16 * restrict q_base = factx.vtcm_q_tiles;
- const __fp16 * restrict k_base = factx.vtcm_k_tiles;
- __fp16 * restrict s_base = factx.vtcm_s_tiles;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(n_col_tiles > 0);
- __builtin_assume(n_dot_tiles > 0);
-
- htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, HTP_MAX_NTHREADS);
- Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_qk);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < n_col_tiles; ++c) {
- const __fp16 * row_tiles = q_base + r * HMX_FP16_TILE_N_ROWS * DK;
- const __fp16 * col_tiles = k_base + c * HMX_FP16_TILE_N_COLS * DK;
- __fp16 * out_tile = s_base + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
- for (size_t k = 0; k < n_dot_tiles; ++k) {
- Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047);
- row_tiles += HMX_FP16_TILE_N_ELMS;
- col_tiles += HMX_FP16_TILE_N_ELMS;
- }
- Q6_mxmem_AR_after_hf(out_tile, 0);
- }
- }
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, HTP_MAX_NTHREADS);
- }
-
- // Pop mask DMA
- MASK_DMA_POP(has_mask_dma);
-
- // Softmax + build_D (multi-thread HVX + serial m/l update)
- fa_softmax_args_t sargs;
- memset(&sargs, 0, sizeof(sargs));
- sargs.factx = &factx;
- sargs.kv_rows = kv_rows;
- sargs.n_rows_g = n_rows_g;
- sargs.n_col_tiles = n_col_tiles;
- sargs.n_tiles_per_bc = n_tiles_per_bc;
- sargs.n_row_tiles = n_row_tiles;
- sargs.n_row_tiles_g_br = n_row_tiles_g_br;
- sargs.Bc = Bc;
- sargs.G = G;
- sargs.kv_head = kv_head;
- sargs.kv_start = kv_start;
- sargs.q_start = q_start;
- sargs.ib3 = ib3;
- sargs.has_alibi = (factx.max_bias != 0.0f);
- sargs.mask = mask;
- sargs.mask_vtcm = has_mask_dma ? (const __fp16 *) factx.vtcm_mask_buf : NULL;
- sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
- sargs.slopes = factx.vtcm_slopes;
- fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
- fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc);
-
- {
- const size_t DV_tiles = (size_t) (DV / 32);
- const __fp16 * restrict d_base = factx.vtcm_d_tiles;
- const __fp16 * restrict p_base = factx.vtcm_p_tiles;
- const __fp16 * restrict v_base = factx.vtcm_v_tiles[0];
- const __fp16 * restrict op_base = o_tile_prev;
- __fp16 * restrict oc_base = o_tile_curr;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(n_col_tiles > 0);
- __builtin_assume(DV_tiles > 0);
-
- htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, HTP_MAX_NTHREADS);
- Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
-
- const __fp16 * p_tile_in = p_base + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * v_tile_in = v_base + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- for (size_t k = 0; k < n_col_tiles; ++k) {
- Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047);
- p_tile_in += HMX_FP16_TILE_N_ELMS;
- v_tile_in += HMX_FP16_TILE_N_ELMS;
- }
-
- __fp16 * o_tile_out = oc_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- Q6_mxmem_AR_after_hf(o_tile_out, 0);
- }
- }
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, HTP_MAX_NTHREADS);
- hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
- }
-
- buf_idx = 1 - buf_idx;
- } // end KV block loop (fallback)
- }
-
- // ---- Final normalization: O = diag(1/l) @ O ----
- {
- fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
-
- // HMX: O_final = diag(1/l) @ O_prev
- if (factx.pipeline) {
- on_job.o_curr = o_tile_curr;
- on_job.o_prev = o_tile_prev;
- on_job.d_tiles = factx.vtcm_d_tiles;
- 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;
- on_job.DV = DV;
- hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
- hmx_queue_pop(ctx->hmx_queue);
- } else {
- const size_t DV_tiles = (size_t) (DV / 32);
- const __fp16 * restrict d_base = factx.vtcm_d_tiles;
- const __fp16 * restrict op_base = o_tile_prev;
- __fp16 * restrict oc_base = o_tile_curr;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(DV_tiles > 0);
-
- htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, HTP_MAX_NTHREADS);
- Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- __fp16 * o_out = oc_base + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
-
- Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
- Q6_mxmem_AR_after_hf(o_out, 0);
- }
- }
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, HTP_MAX_NTHREADS);
- }
- }
-
- // ---- Store O block ----
- fa_phase_o_store(&factx, dst, o_tile_curr, q_start, kv_head, ib3, n_rows_g);
-
-#undef MASK_DMA_PUSH
-#undef MASK_DMA_POP
-#undef DMA_PREFETCH_KV
-
- } // end Q block loop
- } // end KV head loop
- } // end batch loop
-
- if (factx.pipeline) {
- hmx_queue_suspend(ctx->hmx_queue);
- } else {
- HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
- }
-
-
-
- return HTP_STATUS_OK;
-}
// output : fp16 -> f32p
-static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 *restrict vtcm_src, uint32_t start_row, uint32_t n_rows, uint32_t n_cols, uint32_t dst_stride, uint32_t dst_cols) {
+static 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
+) {
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;
const size_t r1 = (r_idx0 % 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;
float *output_row_base = dst + r * dst_stride; // global memory row base for row r (and r+1)
+ const float *src2_row_base = src2 ? (src2 + r * src2_stride) : NULL;
#pragma unroll(4)
for (size_t c = 0; c < limit_c_aligned; c += HTP_MM_HMX_TILE_N_COLS) {
HVX_Vector *pv_out0 = (HVX_Vector *) (output_row_base + c + 0);
HVX_Vector *pv_out1 = (HVX_Vector *) (output_row_base + c + dst_stride);
- *pv_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
+ HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
+ if (src2_row_base) {
+ HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0);
+ v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0);
+ }
+ *pv_out0 = v_out0;
+
if (r + 1 < n_rows) {
- *pv_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
+ HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
+ if (src2_row_base) {
+ HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride);
+ v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1);
+ }
+ *pv_out1 = v_out1;
}
}
HVX_Vector v = ((const HVX_Vector *) tile)[r1];
HVX_VectorPair vp = Q6_Wqf32_vmpy_VhfVhf(v, one);
- hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp)));
+ HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
+ if (src2_row_base) {
+ HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0);
+ v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0);
+ }
+ hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), v_out0);
+
if (r + 1 < n_rows) {
- hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp)));
+ HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
+ if (src2_row_base) {
+ HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride);
+ v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1);
+ }
+ hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), v_out1);
}
}
}
typedef struct {
const __fp16 *vtcm_src;
float *dst;
+ const float *src2;
uint32_t n_tasks;
uint32_t n_tot_chunks;
uint32_t n_chunks_per_task;
uint32_t n_cols;
uint32_t dst_stride; // DDR row stride
+ uint32_t src2_stride; // DDR row stride for residual
uint32_t dst_cols; // Actual output columns
struct htp_thread_trace * traces;
} output_transfer_task_state_t;
// Full range: start_row=0, end_row=n_cols.
static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
const __fp16 * restrict vtcm_src,
- int n_cols,
- int k,
- int src_stride,
- int start_row,
- int end_row) {
+ uint32_t n_cols,
+ uint32_t k,
+ size_t src_stride,
+ uint32_t start_row,
+ uint32_t end_row) {
assert(k % HMX_FP16_TILE_N_COLS == 0);
- const int n_k_tiles = k / HMX_FP16_TILE_N_COLS;
+ const uint32_t n_k_tiles = k / HMX_FP16_TILE_N_COLS;
const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets);
const HVX_Vector v_scat_step = Q6_V_vsplat_R(4);
const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64);
if (pair_scatter) {
// Step c by 64 fp16 (two K-tiles per scatter), advance dst by 2 tiles per iter.
- const int c_step = 2 * HMX_FP16_TILE_N_COLS;
- const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
- const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS;
- const int n_c_iters = k / c_step;
-
- for (int r = start_row; r < end_row; r += 2) {
- const int ct = r / HMX_FP16_TILE_N_ROWS;
- const int local_r = r % HMX_FP16_TILE_N_ROWS;
+ const uint32_t c_step = 2 * HMX_FP16_TILE_N_COLS;
+ const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
+ const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS;
+ const uint32_t n_c_iters = k / c_step;
+
+ for (uint32_t r = start_row; r < end_row; r += 2) {
+ const uint32_t ct = r / HMX_FP16_TILE_N_ROWS;
+ const uint32_t local_r = r % HMX_FP16_TILE_N_ROWS;
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
assert(c_byte_step % 128 == 0);
if (p1) {
- for (int i = 0; i < n_c_iters; ++i) {
+ for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
HVX_Vector v1 = hvx_vmem(p1); p1 += c_byte_step;
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
}
} else {
const HVX_Vector vzero = Q6_V_vzero();
- for (int i = 0; i < n_c_iters; ++i) {
+ for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero);
}
} else {
// Fallback: scatter one K-tile per call (region 2047, masked).
- const int c_step = HMX_FP16_TILE_N_COLS;
- const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
- const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS;
- const int n_c_iters = k / c_step;
-
- for (int r = start_row; r < end_row; r += 2) {
- const int ct = r / HMX_FP16_TILE_N_ROWS;
- const int local_r = r % HMX_FP16_TILE_N_ROWS;
+ const uint32_t c_step = HMX_FP16_TILE_N_COLS;
+ const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
+ const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS;
+ const uint32_t n_c_iters = k / c_step;
+
+ for (uint32_t r = start_row; r < end_row; r += 2) {
+ const uint32_t ct = r / HMX_FP16_TILE_N_ROWS;
+ const uint32_t local_r = r % HMX_FP16_TILE_N_ROWS;
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
if (p1) {
- for (int i = 0; i < n_c_iters; ++i) {
+ for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
HVX_Vector v1 = hvx_vmemu(p1); p1 += c_byte_step;
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
}
} else {
const HVX_Vector vzero = Q6_V_vzero();
- for (int i = 0; i < n_c_iters; ++i) {
+ for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero);
// Full range: start_row=0, end_row=n_rows.
static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
const __fp16 * restrict src,
- int n_rows,
- int head_dim,
- int src_stride,
- int n_row_tiles,
- int start_row,
- int end_row) {
+ uint32_t n_rows,
+ uint32_t head_dim,
+ size_t src_stride,
+ uint32_t n_row_tiles,
+ uint32_t start_row,
+ uint32_t end_row) {
__builtin_assume(head_dim > 0);
const size_t tile_stride_elms = (size_t) n_row_tiles * HMX_FP16_TILE_N_ELMS;
- for (int r = start_row; r < end_row; r += 2) {
+ for (uint32_t r = start_row; r < end_row; r += 2) {
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_rows;
const HVX_Vector * pv_in0 = (const HVX_Vector *) (src + r * src_stride);
const HVX_Vector * pv_in1 = next_row_valid ? (const HVX_Vector *) (src + (r + 1) * src_stride) : NULL;
// Row-pair invariants hoisted out of the c loop.
- const int r0 = r / HMX_FP16_TILE_N_ROWS;
- const int r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2;
+ const uint32_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ const uint32_t r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2;
// tb0 starts at tile (c0=0, r0); tb1 at the adjacent dim-tile (c0=1, r0).
// Each c step (+= 64) advances both by 2 dim-tiles worth of fp16.
const size_t tb_step = 2 * tile_stride_elms;
if (pv_in1) {
- for (int c = 0; c < head_dim; c += 64) {
+ for (uint32_t c = 0; c < head_dim; c += 64) {
HVX_Vector v0 = *pv_in0++;
HVX_Vector v1 = *pv_in1++;
HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
}
} else {
const HVX_Vector vzero = Q6_V_vzero();
- for (int c = 0; c < head_dim; c += 64) {
+ for (uint32_t c = 0; c < head_dim; c += 64) {
HVX_Vector v0 = *pv_in0++;
HVX_VectorPair vp = Q6_W_vshuff_VVR(vzero, v0, -2);
((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp);
HTP_OP_MUL_MAT_ID,
HTP_OP_MUL_MAT_QKV,
HTP_OP_MUL_MAT_FFN,
+ HTP_OP_MUL_MAT_ADD,
HTP_OP_RMS_NORM,
HTP_OP_RMS_NORM_MUL,
HTP_OP_UNARY_SILU,
HTP_TRACE_EVT_HVX_W_DEQUANT = 23,
HTP_TRACE_EVT_HVX_W_PREP = 24,
HTP_TRACE_EVT_HVX_O_PROC = 25,
+ HTP_TRACE_EVT_HVX_FA_QK = 26,
+ HTP_TRACE_EVT_HVX_FA_SFM = 27,
+ HTP_TRACE_EVT_HVX_FA_Q_PREP = 28,
+ HTP_TRACE_EVT_HVX_FA_K_PREP = 29,
+ HTP_TRACE_EVT_HVX_FA_V_PREP = 30,
HTP_TRACE_EVT_HMX_COMP = 40,
};
}
static inline HVX_Vector hvx_vec_f32_to_f16(HVX_Vector v0, HVX_Vector v1) {
- HVX_Vector v = Q6_Vh_vdeal_Vh(hvx_vec_f32_to_f16_shuff(v0, v1));
-
-#if __HVX_ARCH__ < 79
- // replace NaNs with -INF, older arches produce NaNs for (-INF + 0.0)
- const HVX_Vector neg_inf = hvx_vec_splat_f16(-INFINITY);
- HVX_VectorPred nan = hvx_vec_is_nan_f16(v);
- v = Q6_V_vmux_QVV(nan, neg_inf, v);
-#endif
-
- return v;
+ return Q6_Vh_vdeal_Vh(hvx_vec_f32_to_f16_shuff(v0, v1));
}
#if __HVX_ARCH__ >= 79
}
#endif
-
-
static inline HVX_Vector hvx_vec_i16_from_hf_rnd_sat(HVX_Vector vin) {
// This looks complicated.
// Ideally should just be Q6_Vh_equals_Vhf(vin)
#define EXP_COEFF_0 (0x3F000000) // 0.5 = 1/(2!)
#define EXP_LOGN2 (0x3F317218) // ln(2) = 0.6931471805
#define EXP_LOG2E (0x3FB8AA3B) // log2(e) = 1/ln(2) = 1.4426950408
+#define EXP_LOG2E_F 1.44269504f
#define EXP_ONE (0x3f800000) // 1.0
#define EXP_RANGE_R (0x42B17218) // ln(FLT_MAX) approx = 88.7228
#define EXP_RANGE_L (0xC2B00000) // -88.0 (approx log(FLT_MIN))
}
}
+static inline HVX_Vector hvx_vec_exp2_f16(HVX_Vector x_v) {
+ const HVX_Vector zero_v = Q6_V_vzero();
+ const HVX_Vector half_hf_v = Q6_Vh_vsplat_R(0x3800); // fp16 0.5
+
+ // Clamp input to prevent integer underflow in FP16-to-INT16 conversion
+ const HVX_Vector v_clamp_min = hvx_vec_splat_f16(-24.0f);
+ x_v = Q6_Vhf_vmax_VhfVhf(v_clamp_min, x_v);
+
+ // k = round_toward_neg_inf(x); f = (float)k; frac = x - f
+ HVX_Vector x_minus_half = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(x_v, half_hf_v));
+ HVX_Vector k_v = Q6_Vh_equals_Vhf(x_minus_half); // truncate to int16
+ HVX_Vector f_v = Q6_Vhf_equals_Vh(k_v); // back to fp16
+
+ HVX_Vector x_qf16 = Q6_Vqf16_vsub_VhfVhf(x_v, f_v); // fractional part in qf16
+
+ // Horner: y = ((((E5*x + E4)*x + E3)*x + E2)*x + E1)*x + E0
+ HVX_Vector y = Q6_Vqf16_vmpy_Vqf16Vqf16(Q6_Vh_vsplat_R(0x5082), x_qf16); // E5*x
+ y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x157d)); // + E4
+ y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
+ y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x20ed)); // + E3
+ y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
+ y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x2b1b)); // + E2
+ y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
+ y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x33b0)); // + E1
+ y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
+ y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x398c)); // + E0
+ y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16); // y = y * x
+ y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x3c00)); // + 1.0
+
+ // Combine polynomial (mantissa) with integer part (exponent): result = y * 2^k
+ y = Q6_Vhf_equals_Vqf16(y);
+ HVX_Vector y_exp = Q6_Vuh_vlsr_VuhR(Q6_Vh_vasl_VhR(y, 1), 11);
+ y_exp = Q6_Vh_vadd_VhVh(k_v, y_exp);
+ HVX_VectorPred q_underflow = Q6_Q_vcmp_gt_VhVh(zero_v, y_exp);
+ y = Q6_Vh_vaslacc_VhVhR(y, k_v, 10);
+ return Q6_V_vmux_QVV(q_underflow, zero_v, y);
+}
+
#endif /* HVX_EXP_H */
--- /dev/null
+#ifndef HVX_FA_KERNELS_H
+#define HVX_FA_KERNELS_H
+
+#include <assert.h>
+#include <math.h>
+#include "hvx-utils.h"
+
+// Little inner kernels for HVX
+
+#if __HVX_ARCH__ < 79
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
+#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b))
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
+#else
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
+#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b)
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
+#endif
+
+// This is a bit of a hack because the compiler is struggling to properly inline
+// the default hvx_vec_f32_to_f16 with output into the local array.
+static __attribute__((unused)) __attribute__((noinline)) void hvx_vec_f32_to_f16_a(void *ptr, HVX_Vector v0, HVX_Vector v1)
+{
+ *(HVX_Vector *) ptr = hvx_vec_f32_to_f16(v0, v1);
+}
+
+// Dot product of two F16 vectors, accumulating to float
+static inline void hvx_dot_f16_f16_aa(float * restrict r, const void * restrict x, const void * restrict y, unsigned int n, float s) {
+ const HVX_Vector * restrict vx = (const HVX_Vector * restrict) x; // fp16
+ const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
+
+ uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+ HVX_VectorPair rsum_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
+
+ uint32_t i = 0;
+
+ #pragma unroll(4)
+ for (i = 0; i < nvec; i++) {
+ rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, vx[i], vy[i]);
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+ HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
+ HVX_Vector x_hf = Q6_V_vand_QV(bmask, vx[i]);
+
+ rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
+ }
+
+ HVX_Vector rsum = HVX_OP_ADD_F32(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p));
+ rsum = HVX_OP_MUL_F32(hvx_vec_splat_f32(s), hvx_vec_reduce_sum_f32(rsum));
+ hvx_vec_store_u(r, 4, rsum);
+}
+
+static inline HVX_Vector hvx_dot_f16_f16_aa_rx4(const void * restrict y,
+ const uint8_t * restrict x,
+ const size_t stride_x,
+ const size_t nvec,
+ const size_t nloe) {
+ const HVX_Vector * restrict vx0 = (const HVX_Vector * restrict) x; // fp16
+ const HVX_Vector * restrict vx1 = (const HVX_Vector * restrict) (x + stride_x); // fp16
+ const HVX_Vector * restrict vx2 = (const HVX_Vector * restrict) (x + stride_x * 2); // fp16
+ const HVX_Vector * restrict vx3 = (const HVX_Vector * restrict) (x + stride_x * 3); // fp16
+ const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
+
+ HVX_VectorPair rsum0_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
+ HVX_VectorPair rsum1_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
+ HVX_VectorPair rsum2_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
+ HVX_VectorPair rsum3_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
+
+ uint32_t i = 0;
+
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector y_hf = vy[i];
+ HVX_Vector x0_hf = vx0[i];
+ HVX_Vector x1_hf = vx1[i];
+ HVX_Vector x2_hf = vx2[i];
+ HVX_Vector x3_hf = vx3[i];
+
+ rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
+ rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
+ rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
+ rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
+ }
+
+ if (nloe) {
+ // Load x (fp16) and zero-out unused elements
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+ HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
+ HVX_Vector x0_hf = Q6_V_vand_QV(bmask, vx0[i]);
+ HVX_Vector x1_hf = Q6_V_vand_QV(bmask, vx1[i]);
+ HVX_Vector x2_hf = Q6_V_vand_QV(bmask, vx2[i]);
+ HVX_Vector x3_hf = Q6_V_vand_QV(bmask, vx3[i]);
+
+ rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
+ rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
+ rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
+ rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
+ }
+
+ HVX_Vector rsum0 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p));
+ HVX_Vector rsum1 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p));
+ HVX_Vector rsum2 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum2_p), Q6_V_hi_W(rsum2_p));
+ HVX_Vector rsum3 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum3_p), Q6_V_hi_W(rsum3_p));
+
+ HVX_Vector_x4 rsum0123 = { .v = { rsum0, rsum1, rsum2, rsum3 } };
+ return hvx_vec_reduce_sum_f32x4(rsum0123);
+}
+
+static inline HVX_Vector hvx_dot_f16_f16_aa_rx32(const void * restrict y,
+ const uint8_t * restrict x,
+ const size_t stride_x,
+ const size_t n,
+ float s) {
+
+ const size_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ const size_t nloe = n % VLEN_FP16; // leftover elements
+
+ HVX_Vector sums = Q6_V_vzero();
+ const size_t stride_x_4 = stride_x * 4;
+ for (uint32_t j = 0; j < VLEN_FP32; j += 4) {
+ HVX_Vector sums_x4 = hvx_dot_f16_f16_aa_rx4(y, x, stride_x, nvec, nloe);
+ HVX_VectorPred pred = Q6_Q_vsetq_R(j * SIZEOF_FP32);
+ sums = Q6_V_vmux_QVV(pred, sums, sums_x4);
+ x += stride_x_4;
+ }
+
+ return HVX_OP_MUL_F32(hvx_vec_splat_f32(s), sums);
+}
+
+// MAD: y (F32) += x (F16) * s (F16)
+static inline void hvx_mad_f32_f16_aa(float * restrict y, const void * restrict x, const __fp16 * restrict s, 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
+
+ HVX_Vector S0 = hvx_vec_splat_f16(*s);
+
+ 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);
+ }
+ }
+}
+
+// MAD: y (F32) += x0 (F16) * s0 (F16) + x1 (F16) * s1 (F16)
+static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restrict x0, const void * restrict x1,
+ const __fp16 * restrict s0, const __fp16 * restrict 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
+
+ HVX_Vector S0 = hvx_vec_splat_f16(*s0);
+ HVX_Vector S1 = hvx_vec_splat_f16(*s1);
+
+ 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);
+ assert((size_t) src % 128 == 0);
+
+ const HVX_Vector * restrict vsrc = (const HVX_Vector * restrict) src;
+ HVX_Vector * restrict vdst = (HVX_Vector * restrict) dst;
+
+ const uint32_t nvec = n / VLEN_FP32;
+ const uint32_t nloe = n % VLEN_FP32;
+
+ uint32_t i = 0;
+ #pragma unroll(4)
+ for (; i < nvec; ++i) {
+ vdst[i] = HVX_OP_MUL_F32(vsrc[i], vs);
+ }
+ if (nloe) {
+ hvx_vec_store_a(&vdst[i], nloe * sizeof(float), HVX_OP_MUL_F32(vsrc[i], vs));
+ }
+}
+
+#endif /* HVX_FA_KERNELS_H */
// Dot kernels that consume flat (non-tiled) activations
-static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
+}
+
+#if __HVX_ARCH__ < 79
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
+#else
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
+#endif
+
+static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+ const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+ const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+ uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
+ uint32_t nloe = n % VLEN_FP32; // leftover elements
+
+ HVX_Vector rsum = Q6_V_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(4)
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]);
+ rsum = HVX_OP_ADD_F32(rsum, prod);
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]);
+ HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
+ HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf);
+ rsum = HVX_OP_ADD_F32(rsum, prod);
+ }
+
+ *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
+ const void * restrict vx0, const void * restrict vx1,
+ const void * restrict vy0) {
+ const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+ const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+ const HVX_Vector * restrict y = (const HVX_Vector *) vy0;
+
+ uint32_t nvec = n / VLEN_FP32;
+ uint32_t nloe = n % VLEN_FP32;
+
+ HVX_Vector rsum0 = Q6_V_vzero();
+ HVX_Vector rsum1 = Q6_V_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector y_sf = y[i];
+ HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf);
+ HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf);
+ rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
+ rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
+ HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]);
+ HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]);
+ HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf);
+ HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf);
+ rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
+ rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
+ }
+
+ HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+ hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+ const void * restrict vx0, const void * restrict vx1,
+ const void * restrict vy0, const void * restrict vy1) {
+ const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+ const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+ const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+ const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+ uint32_t nvec = n / VLEN_FP32;
+ uint32_t nloe = n % VLEN_FP32;
+
+ HVX_Vector r0_c0_sum = Q6_V_vzero();
+ HVX_Vector r0_c1_sum = Q6_V_vzero();
+ HVX_Vector r1_c0_sum = Q6_V_vzero();
+ HVX_Vector r1_c1_sum = Q6_V_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector r0_sf = x0[i];
+ HVX_Vector r1_sf = x1[i];
+ HVX_Vector c0_sf = y0[i];
+ HVX_Vector c1_sf = y1[i];
+
+ r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
+ r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
+ r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
+ r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+
+ HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]);
+ HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]);
+ HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]);
+ HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]);
+
+ r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
+ r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
+ r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
+ r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
+ }
+
+ // Reduce and store results
+ HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+ HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+ hvx_vec_store_u(s0, 8, r0_r1_c0_sum);
+ hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
}
+
+static inline void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
+ const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
+ const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
+
+ uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
+ uint32_t nloe = n % VLEN_FP32; // leftover elements
+
+ HVX_Vector rsum = Q6_V_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector x_sf = vx[i];
+ HVX_Vector y_sf = vy[i];
+
+ rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
+ }
+
+ if (nloe) {
+ HVX_Vector x_sf = vx[i];
+ HVX_Vector y_sf = vy[i];
+
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ x_sf = Q6_V_vand_QV(bmask, x_sf);
+ y_sf = Q6_V_vand_QV(bmask, y_sf);
+
+ rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
+ }
+
+ rsum = hvx_vec_reduce_sum_f32(rsum);
+ hvx_vec_store_u(&s[0], 4, rsum);
+}
+
+#undef HVX_OP_ADD_F32
+#undef HVX_OP_MUL_F32
+
+static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+ const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+ const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+ uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+ HVX_VectorPair rsum_p = Q6_W_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(4)
+ for (i = 0; i < nvec; i++) {
+ rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]);
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+ HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
+ HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
+ rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
+ }
+
+ HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p)));
+ hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
+ const void * restrict vx0, const void * restrict vx1,
+ const void * restrict vy0) {
+ const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+ const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+ const HVX_Vector * restrict y = (const HVX_Vector *) vy0;
+
+ uint32_t nvec = n / VLEN_FP16;
+ uint32_t nloe = n % VLEN_FP16;
+
+ HVX_VectorPair rsum0_p = Q6_W_vzero();
+ HVX_VectorPair rsum1_p = Q6_W_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector y_hf = y[i];
+ rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf);
+ rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf);
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+ HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
+ HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]);
+ HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]);
+ rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
+ rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
+ }
+
+ HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p)));
+ HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p)));
+ HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+ hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+ const void * restrict vx0, const void * restrict vx1,
+ const void * restrict vy0, const void * restrict vy1) {
+ const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+ const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+ const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+ const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+ uint32_t nvec = n / VLEN_FP16;
+ uint32_t nloe = n % VLEN_FP16;
+
+ // Row sums (sf) - 4 accumulators for 2x2 tile
+ HVX_VectorPair r0_c0_sum_p = Q6_W_vzero();
+ HVX_VectorPair r0_c1_sum_p = Q6_W_vzero();
+ HVX_VectorPair r1_c0_sum_p = Q6_W_vzero();
+ HVX_VectorPair r1_c1_sum_p = Q6_W_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; i++) {
+ HVX_Vector r0_hf = x0[i];
+ HVX_Vector r1_hf = x1[i];
+ HVX_Vector c0_hf = y0[i];
+ HVX_Vector c1_hf = y1[i];
+
+ // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1
+ r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
+ r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
+ r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
+ r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+
+ HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]);
+ HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]);
+ HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]);
+ HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]);
+
+ r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
+ r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
+ r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
+ r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
+ }
+
+ HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p)));
+ HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p)));
+ HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p)));
+ HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p)));
+
+ // Reduce and store results
+ HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+ HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+ hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum); // row0,col0 row1,col0
+ hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum); // row0,col1 row1,col1
+}
+
+static inline void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+ const HVX_UVector * restrict x = (const HVX_UVector *) vx;
+ const HVX_UVector * restrict y = (const HVX_UVector *) vy;
+
+ uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+ HVX_Vector rsum = Q6_V_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(4)
+ for (i = 0; i < nvec; i++) {
+ HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]);
+ rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
+ }
+
+ if (nloe) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+ HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
+ HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
+
+ HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
+ rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
+ }
+
+ rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
+ hvx_vec_store_u(&s[0], 4, rsum);
+}
+
+static inline void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
+ const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
+ const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
+
+ uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+ uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+ const HVX_Vector zero = Q6_V_vzero();
+
+ HVX_Vector rsum = Q6_V_vzero();
+
+ uint32_t i = 0;
+
+ #pragma unroll(2)
+ for (i = 0; i < nvec; i++) {
+ // Load y (fp32) and convert into fp16
+ HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements
+ HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements
+ HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
+
+ // Load x (fp16)
+ HVX_Vector x_hf = vx[i];
+
+ HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
+
+ rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
+ }
+
+ if (nloe) {
+ // Load y (fp32) and convert into fp16
+ HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements
+ HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements
+ HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
+
+ // Load x (fp16)
+ HVX_Vector x_hf = vx[i];
+
+ // Zero-out unused elements
+ // Note that we need to clear both x and y because they may contain NANs
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+ x_hf = Q6_V_vand_QV(bmask, x_hf);
+ y_hf = Q6_V_vand_QV(bmask, y_hf);
+
+ HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
+
+ rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
+ }
+
+ // Convert into fp32 and reduce
+ rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
+ hvx_vec_store_u(&s[0], 4, rsum);
+}
+
+static inline void hvx_tensor_add_f32_grid(
+ const struct htp_tensor * restrict dst,
+ const struct htp_tensor * restrict src2,
+ uint32_t start_row,
+ uint32_t end_row,
+ uint32_t start_col,
+ uint32_t end_col,
+ const struct fastdiv_values * div_ne11_12,
+ const struct fastdiv_values * div_ne11
+) {
+ if (start_row >= end_row || start_col >= end_col) return;
+ const uint32_t nb1 = dst->nb[1]; // row stride in bytes
+
+ const uint32_t ne11 = dst->ne[1];
+ const uint32_t ne12 = dst->ne[2];
+ const uint32_t ne11_12 = ne11 * ne12;
+
+ const bool is_broadcast1 = (src2->ne[1] == 1);
+ const bool is_broadcast2 = (src2->ne[2] == 1);
+ const bool is_broadcast3 = (src2->ne[3] == 1);
+
+ for (uint32_t r = start_row; r < end_row; r++) {
+ float * dst_row = (float *) ((uint8_t *) dst->data + r * nb1);
+
+ uint32_t i13 = fastdiv(r, div_ne11_12);
+ uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11);
+ uint32_t i11 = r - i13 * ne11_12 - i12 * ne11;
+
+ uint32_t i23 = is_broadcast3 ? 0 : i13;
+ uint32_t i22 = is_broadcast2 ? 0 : i12;
+ uint32_t i21 = is_broadcast1 ? 0 : i11;
+
+ const float * src2_row = (const float *) ((const uint8_t *) src2->data +
+ i21 * src2->nb[1] + i22 * src2->nb[2] + i23 * src2->nb[3]);
+
+ float * dst_ptr = &dst_row[start_col];
+ const float * src2_ptr = &src2_row[start_col];
+ int remaining = end_col - start_col;
+ while (remaining >= 32) {
+ HVX_Vector v_out = hvx_vmemu(dst_ptr);
+ HVX_Vector v_z = hvx_vmemu(src2_ptr);
+ hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z);
+ dst_ptr += 32;
+ src2_ptr += 32;
+ remaining -= 32;
+ }
+ if (remaining > 0) {
+ HVX_Vector v_out = hvx_vmemu(dst_ptr);
+ HVX_Vector v_z = hvx_vmemu(src2_ptr);
+ hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
+ }
+ }
+}
+
return Q6_W_vcombine_VV(v_sum1, v_sum0);
}
-static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
}
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
}
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
-static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y_q = vy;
v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
- hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ if (sz) {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+ } else {
+ hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+ }
}
-static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
const uint8_t * restrict tile_ptr = vx;
const uint8_t * restrict y0_q = vy0;
const uint8_t * restrict y1_q = vy1;
v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
- hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
- hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ if (sz0) {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+ } else {
+ hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+ }
+ if (sz1) {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+ } else {
+ hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+ }
}
static inline void quantize_f32_q8_0_tiled_kernel(
#include "hvx-base.h"
#include "hvx-inverse.h"
+#include "hvx-exp.h"
#define FAST_SIGMOID_LOG2F (0x3fb8aa3b) // 1.442695022
#define FAST_SIGMOID_C1 (0x3d009076) // 0.03138777
hvx_tanh_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
}
+static inline HVX_Vector hvx_vec_fast_sigmoid_f16(HVX_Vector x_v) {
+ const HVX_Vector v_one = hvx_vec_splat_f16(1.0f);
+ const HVX_Vector v_neg_log2e = hvx_vec_splat_f16(-EXP_LOG2E_F);
+ const HVX_Vector em_mask = Q6_Vh_vsplat_R(0x7FFF);
+
+ // Compute absolute value of x_v
+ HVX_Vector abs_x = Q6_V_vand_VV(x_v, em_mask);
+
+ // Compute u = -abs_x * log2(e) <= 0.
+ HVX_Vector u = hvx_vec_mul_f16_f16(abs_x, v_neg_log2e);
+
+ // Clamp input to prevent underflow in exp2
+ const HVX_Vector v_clamp_min = hvx_vec_splat_f16(-24.0f);
+ u = Q6_Vhf_vmax_VhfVhf(v_clamp_min, u);
+
+ HVX_Vector exp_val = hvx_vec_exp2_f16(u);
+ HVX_Vector denom = hvx_vec_add_f16_f16(v_one, exp_val);
+ HVX_Vector sig_abs = hvx_vec_inverse_f16(denom);
+
+ // check if x_v < 0 (using integer comparison on absolute value)
+ HVX_VectorPred is_neg = Q6_Q_vcmp_gt_VhVh(abs_x, x_v);
+
+ // If x_v < 0, return 1.0f - sig_abs
+ HVX_Vector sig_neg = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(v_one, sig_abs));
+ return Q6_V_vmux_QVV(is_neg, sig_neg, sig_abs);
+}
+
+static inline HVX_Vector hvx_vec_tanh_f16(HVX_Vector x) {
+ // tanh(x) = 2 * sigmoid(2x) - 1
+ const HVX_Vector v_two = hvx_vec_splat_f16(2.0f);
+
+ HVX_Vector x2 = hvx_vec_mul_f16_f16(x, v_two);
+ HVX_Vector sig2x = hvx_vec_fast_sigmoid_f16(x2);
+
+ const HVX_Vector v_neg_one = hvx_vec_splat_f16(-1.0f);
+ return hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(sig2x, v_two), v_neg_one);
+}
+
#endif /* HVX_SIGMOID_H */
static int execute_op(struct htp_ops_context * octx) {
switch (octx->op) {
case HTP_OP_MUL_MAT:
+ case HTP_OP_MUL_MAT_ADD:
return op_matmul(octx);
case HTP_OP_MUL_MAT_ID:
#include <math.h>
#include <string.h>
+#include <stdatomic.h>
#include "hex-dma.h"
#include "hvx-utils.h"
#include "matmul-ops.h"
#include "vtcm-utils.h"
+static void hvx_tensor_add_f32_grid(
+ const struct htp_tensor * restrict dst,
+ const struct htp_tensor * restrict src2,
+ uint32_t start_row,
+ uint32_t end_row,
+ uint32_t start_col,
+ uint32_t end_col,
+ const struct fastdiv_values * div_ne11_12,
+ const struct fastdiv_values * div_ne11
+);
+
typedef struct {
float *dst;
+ const float *src2;
const float *activation;
const __fp16 *weight;
int m;
int act_stride;
int weight_stride;
int dst_stride;
+ uint32_t src2_stride;
int ne02;
int ne03;
int ne12;
size_t src1_nb3;
size_t dst_nb2;
size_t dst_nb3;
+ size_t src2_nb2;
+ size_t src2_nb3;
} hmx_mm_f16_f32_batched_params_t;
struct htp_mm_context {
void (*vec_dot_32x1)(const uint32_t n, float * restrict s,
const void * restrict vx,
- const void * restrict vy, uint32_t valid_rows);
+ const void * restrict vy, uint32_t valid_rows,
+ const float * restrict sz);
// Precomputed values
uint32_t src0_nrows_per_thread;
- uint32_t src1_nrows_per_thread;
struct fastdiv_values mm_div_ne12_ne1;
struct fastdiv_values mm_div_ne1;
struct fastdiv_values mm_div_r3;
struct fastdiv_values mm_div_ne11;
+ // Per thread quant tasks
// Precomputed block-parallel quantization values
- 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];
+ 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 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;
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
};
-#if __HVX_ARCH__ < 79
-#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
-#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
-#else
-#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
-#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
-#endif
-
-static void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
- const HVX_Vector * restrict x = (const HVX_Vector *) vx;
- const HVX_Vector * restrict y = (const HVX_Vector *) vy;
-
- uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
- uint32_t nloe = n % VLEN_FP32; // leftover elements
-
- HVX_Vector rsum = Q6_V_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(4)
- for (i = 0; i < nvec; i++) {
- HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]);
- rsum = HVX_OP_ADD_F32(rsum, prod);
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]);
- HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
- HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf);
- rsum = HVX_OP_ADD_F32(rsum, prod);
- }
-
- *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
-}
-
-static void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
- const void * restrict vx0, const void * restrict vx1,
- const void * restrict vy0) {
- const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
- const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
- const HVX_Vector * restrict y = (const HVX_Vector *) vy0;
-
- uint32_t nvec = n / VLEN_FP32;
- uint32_t nloe = n % VLEN_FP32;
-
- HVX_Vector rsum0 = Q6_V_vzero();
- HVX_Vector rsum1 = Q6_V_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(2)
- for (i = 0; i < nvec; i++) {
- HVX_Vector y_sf = y[i];
- HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf);
- HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf);
- rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
- rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
- HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]);
- HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]);
- HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf);
- HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf);
- rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
- rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
- }
-
- HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
- HVX_VectorAlias va;
- va.v = rsum;
- s0[0] = va.fp32[0];
- s0[1] = va.fp32[1];
-}
-
-static void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
- const void * restrict vx0, const void * restrict vx1,
- const void * restrict vy0, const void * restrict vy1) {
- const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
- const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
- const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
- const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
-
- uint32_t nvec = n / VLEN_FP32;
- uint32_t nloe = n % VLEN_FP32;
-
- HVX_Vector r0_c0_sum = Q6_V_vzero();
- HVX_Vector r0_c1_sum = Q6_V_vzero();
- HVX_Vector r1_c0_sum = Q6_V_vzero();
- HVX_Vector r1_c1_sum = Q6_V_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(2)
- for (i = 0; i < nvec; i++) {
- HVX_Vector r0_sf = x0[i];
- HVX_Vector r1_sf = x1[i];
- HVX_Vector c0_sf = y0[i];
- HVX_Vector c1_sf = y1[i];
-
- r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
- r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
- r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
- r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
-
- HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]);
- HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]);
- HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]);
- HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]);
-
- r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
- r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
- r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
- r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
- }
-
- // Reduce and store results
- HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
- HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
-
- HVX_VectorAlias va0, va1;
- va0.v = r0_r1_c0_sum;
- va1.v = r0_r1_c1_sum;
- s0[0] = va0.fp32[0];
- s0[1] = va0.fp32[1];
- s1[0] = va1.fp32[0];
- s1[1] = va1.fp32[1];
-}
-
-static void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
- const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
- const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
-
- uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
- uint32_t nloe = n % VLEN_FP32; // leftover elements
-
- HVX_Vector rsum = Q6_V_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(2)
- for (i = 0; i < nvec; i++) {
- HVX_Vector x_sf = vx[i];
- HVX_Vector y_sf = vy[i];
-
- rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
- }
-
- if (nloe) {
- HVX_Vector x_sf = vx[i];
- HVX_Vector y_sf = vy[i];
-
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
- x_sf = Q6_V_vand_QV(bmask, x_sf);
- y_sf = Q6_V_vand_QV(bmask, y_sf);
-
- rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
- }
-
- rsum = hvx_vec_reduce_sum_f32(rsum);
- hvx_vec_store_u(&s[0], 4, rsum);
-}
-
-static void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
- const HVX_Vector * restrict x = (const HVX_Vector *) vx;
- const HVX_Vector * restrict y = (const HVX_Vector *) vy;
-
- uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
- uint32_t nloe = n % VLEN_FP16; // leftover elements
-
- HVX_VectorPair rsum_p = Q6_W_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(4)
- for (i = 0; i < nvec; i++) {
- rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]);
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
- HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
- HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
- rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
- }
-
- HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p)));
- hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
-}
-
-static void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
- const void * restrict vx0, const void * restrict vx1,
- const void * restrict vy0) {
- const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
- const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
- const HVX_Vector * restrict y = (const HVX_Vector *) vy0;
-
- uint32_t nvec = n / VLEN_FP16;
- uint32_t nloe = n % VLEN_FP16;
-
- HVX_VectorPair rsum0_p = Q6_W_vzero();
- HVX_VectorPair rsum1_p = Q6_W_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(2)
- for (i = 0; i < nvec; i++) {
- HVX_Vector y_hf = y[i];
- rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf);
- rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf);
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
- HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
- HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]);
- HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]);
- rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
- rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
- }
-
- HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p)));
- HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p)));
- HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
- hvx_vec_store_u(s0, 8, rsum);
-}
-
-static void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
- const void * restrict vx0, const void * restrict vx1,
- const void * restrict vy0, const void * restrict vy1) {
- const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
- const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
- const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
- const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
-
- uint32_t nvec = n / VLEN_FP16;
- uint32_t nloe = n % VLEN_FP16;
-
- // Row sums (sf) - 4 accumulators for 2×2 tile
- HVX_VectorPair r0_c0_sum_p = Q6_W_vzero();
- HVX_VectorPair r0_c1_sum_p = Q6_W_vzero();
- HVX_VectorPair r1_c0_sum_p = Q6_W_vzero();
- HVX_VectorPair r1_c1_sum_p = Q6_W_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(2)
- for (i = 0; i < nvec; i++) {
- HVX_Vector r0_hf = x0[i];
- HVX_Vector r1_hf = x1[i];
- HVX_Vector c0_hf = y0[i];
- HVX_Vector c1_hf = y1[i];
-
- // Compute 4 dot products: r0×c0, r0×c1, r1×c0, r1×c1
- r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
- r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
- r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
- r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
-
- HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]);
- HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]);
- HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]);
- HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]);
-
- r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
- r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
- r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
- r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
- }
-
- HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p)));
- HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p)));
- HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p)));
- HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p)));
- // Reduce and store results
- HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
- HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
-
- hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum); // row0,col0 row1,col0
- hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum); // row0,col1 row1,col1
-}
-
-static void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
- const HVX_UVector * restrict x = (const HVX_UVector *) vx;
- const HVX_UVector * restrict y = (const HVX_UVector *) vy;
-
- uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
- uint32_t nloe = n % VLEN_FP16; // leftover elements
-
- HVX_Vector rsum = Q6_V_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(4)
- for (i = 0; i < nvec; i++) {
- HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]);
- rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
- }
-
- if (nloe) {
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
- HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
- HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
-
- HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
- rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
- }
-
- rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
- hvx_vec_store_u(&s[0], 4, rsum);
-}
-
-static void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
- const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
- const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
-
- uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
- uint32_t nloe = n % VLEN_FP16; // leftover elements
-
- const HVX_Vector zero = Q6_V_vzero();
-
- HVX_Vector rsum = Q6_V_vzero();
-
- uint32_t i = 0;
-
- #pragma unroll(2)
- for (i = 0; i < nvec; i++) {
- // Load y (fp32) and convert into fp16
- HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements
- HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements
- HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
-
- // Load x (fp16)
- HVX_Vector x_hf = vx[i];
-
- HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
-
- rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
- }
-
- if (nloe) {
- // Load y (fp32) and convert into fp16
- HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements
- HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements
- HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
-
- // Load x (fp16)
- HVX_Vector x_hf = vx[i];
-
- // Zero-out unused elements
- // Note that we need to clear both x and y because they may contain NANs
- HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
- x_hf = Q6_V_vand_QV(bmask, x_hf);
- y_hf = Q6_V_vand_QV(bmask, y_hf);
-
- HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
-
- rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
- }
-
- // Convert into fp32 and reduce
- rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
- hvx_vec_store_u(&s[0], 4, rsum);
-}
#define htp_matmul_tensors_preamble \
const struct htp_tensor * restrict src0 = octx->src[0]; \
const uint32_t ne12 = src1->ne[2]; \
const uint32_t ne13 = src1->ne[3]; \
\
- const uint32_t ne20 = src2->ne[0]; \
- const uint32_t ne21 = src2->ne[1]; \
- const uint32_t ne22 = src2->ne[2]; \
- const uint32_t ne23 = src2->ne[3]; \
+ const uint32_t ne20 = src2 ? src2->ne[0] : 0; \
+ const uint32_t ne21 = src2 ? src2->ne[1] : 0; \
+ const uint32_t ne22 = src2 ? src2->ne[2] : 0; \
+ const uint32_t ne23 = src2 ? src2->ne[3] : 0; \
\
const uint32_t ne0 = dst->ne[0]; \
const uint32_t ne1 = dst->ne[1]; \
uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; \
htp_matmul_tensors_preamble;
+static inline void hvx_mm_run_quant_task(struct htp_mm_context * mmctx, unsigned int ith) {
+ if (mmctx->quant_task_func) {
+ if (ith < mmctx->n_quant_tasks) {
+ mmctx->quant_task_func(mmctx->n_quant_tasks, ith, mmctx);
+ atomic_fetch_sub(&mmctx->quant_barrier, 1);
+ }
+ while (atomic_load(&mmctx->quant_barrier) > 0) {
+ // spin
+ }
+ }
+}
+
// *** matmul with support for 4d tensors and full broadcasting
static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) {
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start);
+ if (src2) {
+ hvx_tensor_add_f32_grid(dst, src2, ir1_start, ir1_end, ir0_start, ir0_end, &mmctx->mm_div_ne12_ne1, &mmctx->mm_div_ne1);
+ }
}
#include "hmx-mm-kernels-tiled.h"
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; \
- } \
- \
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
\
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \
const size_t dst_row_size = nb1; \
const size_t src1_row_size = nb11; \
const size_t src1_stride = mmctx->vtcm_src1_stride; \
+ const size_t src2_stride = src2 ? ((src2->ne[1] == 1) ? 0 : src2->nb[1]) : 0; \
\
uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \
uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \
uint32_t ct_end = (src0_end_row + 31) / 32; \
\
uint32_t push_ct = ct_start; \
- 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); \
+ if (src0_start_row < src0_end_row) { \
+ 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); \
+ } \
+ } \
+ \
+ hvx_mm_run_quant_task(mmctx, ith); \
+ \
+ if (src0_start_row >= src0_end_row) { \
+ return; \
} \
\
for (uint32_t ct = ct_start; ct < ct_end; ct++) { \
float * dst_ptr0 = &dst_row0[ct * 32]; \
float * dst_ptr1 = &dst_row1[ct * 32]; \
\
- DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows); \
+ const float * src2_ptr0 = NULL; \
+ const float * src2_ptr1 = NULL; \
+ if (src2) { \
+ const float * restrict src2_row0 = (const float *) ((const uint8_t *) src2->data + ((ir1+0) * src2_stride)); \
+ const float * restrict src2_row1 = (const float *) ((const uint8_t *) src2->data + ((ir1+1) * src2_stride)); \
+ src2_ptr0 = &src2_row0[ct * 32]; \
+ src2_ptr1 = &src2_row1[ct * 32]; \
+ } \
+ DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, src2_ptr0, src2_ptr1); \
} \
\
for (; ir1 < src1_nrows; ++ir1) { \
float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \
float * dst_ptr = &dst_row[ct * 32]; \
\
- DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows); \
+ const float * src2_ptr = NULL; \
+ if (src2) { \
+ const float * restrict src2_row = (const float *) ((const uint8_t *) src2->data + (ir1 * src2_stride)); \
+ src2_ptr = &src2_row[ct * 32]; \
+ } \
+ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, src2_ptr); \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\
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; \
- } \
- \
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL; \
\
const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \
float * tmp = (float *) vtcm_dst_ptr; \
\
const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \
+ \
const uint8_t * restrict src1_col = (const uint8_t *) src1_data; \
float * restrict dst_col = (float *) dst->data; \
\
uint32_t ct_end = (src0_end_row + 31) / 32; \
\
uint32_t push_ct = ct_start; \
- 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); \
+ if (src0_start_row < src0_end_row) { \
+ 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); \
+ } \
+ } \
+ \
+ hvx_mm_run_quant_task(mmctx, ith); \
+ \
+ if (src0_start_row >= src0_end_row) { \
+ return; \
} \
\
for (uint32_t ct = ct_start; ct < ct_end; ct++) { \
valid_rows = MIN(32, MAX(0, valid_rows)); \
\
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
- DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows); \
+ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\
if (push_ct < ct_end) { \
\
int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \
if (copy_cnt > 0) { \
- hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \
+ 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; \
+ } \
+ } else { \
+ hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \
+ } \
} \
}
const uint32_t start_row_kv = src0_nrows_per_thread_kv * ith; \
const uint32_t end_row_kv = MIN(start_row_kv + src0_nrows_per_thread_kv, src0_nrows_kv); \
\
- if (start_row_kv < end_row_kv) { \
- uint32_t ct_start_kv = start_row_kv / 32; \
- uint32_t ct_end_kv = (end_row_kv + 31) / 32; \
+ uint32_t ct_start_kv = start_row_kv / 32; \
+ uint32_t ct_end_kv = (end_row_kv + 31) / 32; \
\
- uint32_t push_ct = ct_start_kv; \
+ uint32_t push_ct = ct_start_kv; \
+ if (start_row_kv < end_row_kv) { \
for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end_kv; 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); \
dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \
src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \
} \
+ } \
+ \
+ hvx_mm_run_quant_task(mmctx, ith); \
+ \
+ if (start_row_kv < end_row_kv) { \
\
for (uint32_t ct = ct_start_kv; ct < ct_end_kv; ct++) { \
const uint8_t * w_tile_k = dma_queue_pop(dma_queue).dst; \
float * dst_ptr0_v = &dst_row0_v[ct * 32]; \
float * dst_ptr1_v = &dst_row1_v[ct * 32]; \
\
- DOT_2X2(ne10, dst_ptr0_k, dst_ptr1_k, w_tile_k, src1_col0, src1_col1, valid_rows); \
- DOT_2X2(ne10, dst_ptr0_v, dst_ptr1_v, w_tile_v, src1_col0, src1_col1, valid_rows); \
+ DOT_2X2(ne10, dst_ptr0_k, dst_ptr1_k, w_tile_k, src1_col0, src1_col1, valid_rows, NULL, NULL); \
+ DOT_2X2(ne10, dst_ptr0_v, dst_ptr1_v, w_tile_v, src1_col0, src1_col1, valid_rows, NULL, NULL); \
} \
\
for (; ir1 < src1_nrows; ++ir1) { \
float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); \
float * dst_ptr_v = &dst_row_v[ct * 32]; \
\
- DOT_2X1(ne10, dst_ptr_k, w_tile_k, src1_col, valid_rows); \
- DOT_2X1(ne10, dst_ptr_v, w_tile_v, src1_col, valid_rows); \
+ DOT_2X1(ne10, dst_ptr_k, w_tile_k, src1_col, valid_rows, NULL); \
+ DOT_2X1(ne10, dst_ptr_v, w_tile_v, src1_col, valid_rows, NULL); \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ith); \
\
float * dst_ptr0_q = &dst_row0_q[ct * 32]; \
float * dst_ptr1_q = &dst_row1_q[ct * 32]; \
\
- DOT_2X2(ne10, dst_ptr0_q, dst_ptr1_q, w_tile_q, src1_col0, src1_col1, valid_rows); \
+ DOT_2X2(ne10, dst_ptr0_q, dst_ptr1_q, w_tile_q, src1_col0, src1_col1, valid_rows, NULL, NULL); \
} \
\
for (; ir1 < src1_nrows; ++ir1) { \
float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); \
float * dst_ptr_q = &dst_row_q[ct * 32]; \
\
- DOT_2X1(ne10, dst_ptr_q, w_tile_q, src1_col, valid_rows); \
+ DOT_2X1(ne10, dst_ptr_q, w_tile_q, src1_col, valid_rows, NULL); \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\
uint32_t ct_end = (src0_end_row + 31) / 32; \
\
uint32_t push_ct = ct_start; \
- 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); \
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \
- src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \
+ if (src0_start_row < src0_end_row) { \
+ 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); \
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \
+ src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \
+ } \
+ } \
+ \
+ hvx_mm_run_quant_task(mmctx, ith); \
+ \
+ if (src0_start_row >= src0_end_row) { \
+ return; \
} \
\
for (uint32_t ct = ct_start; ct < ct_end; ct++) { \
float * dst_ptr0_up = &dst_row0_up[ct * 32]; \
float * dst_ptr1_up = &dst_row1_up[ct * 32]; \
\
- DOT_2X2(ne10, dst_ptr0_gate, dst_ptr1_gate, w_tile_gate, src1_col0, src1_col1, valid_rows); \
- DOT_2X2(ne10, dst_ptr0_up, dst_ptr1_up, w_tile_up, src1_col0, src1_col1, valid_rows); \
+ DOT_2X2(ne10, dst_ptr0_gate, dst_ptr1_gate, w_tile_gate, src1_col0, src1_col1, valid_rows, NULL, NULL); \
+ DOT_2X2(ne10, dst_ptr0_up, dst_ptr1_up, w_tile_up, src1_col0, src1_col1, valid_rows, NULL, NULL); \
} \
\
for (; ir1 < src1_nrows; ++ir1) { \
float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); \
float * dst_ptr_up = &dst_row_up[ct * 32]; \
\
- DOT_2X1(ne10, dst_ptr_gate, w_tile_gate, src1_col, valid_rows); \
- DOT_2X1(ne10, dst_ptr_up, w_tile_up, src1_col, valid_rows); \
+ DOT_2X1(ne10, dst_ptr_gate, w_tile_gate, src1_col, valid_rows, NULL); \
+ DOT_2X1(ne10, dst_ptr_up, w_tile_up, src1_col, valid_rows, NULL); \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \
\
const uint32_t ne2 = src->ne[2]; \
const uint32_t ne3 = src->ne[3]; \
const uint32_t nrows = ne1 * ne2 * ne3; \
- const uint32_t nrows_per_thread = mmctx->src1_nrows_per_thread; \
+ const uint32_t nrows_per_thread = mmctx->n_quant_rows_per_thread; \
\
const uint32_t ir_first = nrows_per_thread * ith; \
if (ir_first >= nrows) { \
const size_t dst_row_size = (dst_row_size_expr); \
const uint8_t * restrict src_data = (const uint8_t *) src->data + (src_row_size * ir_first); \
uint8_t * restrict dst_data = (uint8_t *) dst + (dst_row_size * ir_first); \
- uint8_t * restrict tmp_data = (uint8_t *) mmctx->vtcm_src0 + (mmctx->vtcm_src0_size_per_thread * ith); \
+ uint8_t * restrict tmp_data = (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith); \
kernel_fn(src_data, dst_data, tmp_data, ne0, ir_last - ir_first, src_row_size, dst_row_size); \
\
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \
quantize_f32_q8_0_tiled_block_kernel(
(const float *) src->data,
mmctx->vtcm_src1,
- (uint8_t *) mmctx->vtcm_src0 + (mmctx->vtcm_src0_size_per_thread * ith),
+ (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith),
src->ne[0],
mmctx->quant_ib_first[ith],
mmctx->quant_ib_last[ith],
quantize_f32_q8_1_tiled_block_kernel(
(const float *) src->data,
mmctx->vtcm_src1,
- (uint8_t *) mmctx->vtcm_src0 + (mmctx->vtcm_src0_size_per_thread * ith),
+ (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith),
src->ne[0],
mmctx->quant_ib_first[ith],
mmctx->quant_ib_last[ith],
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);
- // no work for this thread
- if (src0_start_row >= src0_end_row) {
- return;
- }
-
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
const size_t dst_row_size = nb1;
const uint8_t * restrict src0_row = (const uint8_t *) src0->data;
// Prefill vtcm with src0 rows
- #pragma unroll(4)
- for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
- const int is0 = (ir0 - src0_start_row);
- if (is0 >= (int)n_prefetch) {
- break;
+ if (src0_start_row < src0_end_row) {
+ for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
+ const int is0 = (ir0 - src0_start_row);
+ if (is0 >= (int)n_prefetch) {
+ break;
+ }
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
+ src0_stride, src0_row_size, src0_row_size, 2);
}
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
- src0_stride, src0_row_size, src0_row_size, 2);
+ }
+
+ hvx_mm_run_quant_task(mmctx, ith);
+
+ if (src0_start_row >= src0_end_row) {
+ return;
}
// Process src0 rows
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0);
}
+ if (src2) {
+ hvx_tensor_add_f32_grid(dst, src2, 0, src1_nrows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1);
+ }
}
static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
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;
- }
-
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
const size_t dst_row_size = nb1;
const uint32_t prefetch_mask = n_prefetch - 1;
// Prefill vtcm with 2x src0 rows
- #pragma unroll(2)
- 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) {
- break;
+ if (src0_start_row < src0_end_row) {
+ 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) {
+ break;
+ }
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
+ src0_stride, src0_row_size, src0_row_size, 2);
}
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
- src0_stride, src0_row_size, src0_row_size, 2);
+ }
+
+ hvx_mm_run_quant_task(mmctx, ith);
+
+ if (src0_start_row >= src0_end_row) {
+ return;
}
// Process src0 rows
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0);
}
- hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, src0_end_row - src0_start_row);
+ 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;
+ }
+ } else {
+ hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
+ }
}
#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ids->ne[0] * ids->ne[1] + (i1)]
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
+ hvx_mm_run_quant_task(mmctx, ith);
+
if (src0_start_row >= src0_end_row) {
return;
}
const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (ir1 + rm2 * ne11 + 0) * src1_stride);
float * restrict dst_row = (float *) (dst->data + (rm1 * nb1 + rm2 * nb2 + 0));
- mmctx->vec_dot_32x1(ne10, &dst_row[ct * 32], w_tile, src1_col, valid_rows);
+ mmctx->vec_dot_32x1(ne10, &dst_row[ct * 32], w_tile, src1_col, valid_rows, NULL);
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct);
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
+ hvx_mm_run_quant_task(mmctx, ith);
+
if (src0_start_row >= src0_end_row) {
return;
}
valid_rows = MIN(32, MAX(0, valid_rows));
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct);
- mmctx->vec_dot_32x1(ne10, &dst_row[ct * 32], w_tile, src1_col, valid_rows);
+ mmctx->vec_dot_32x1(ne10, &dst_row[ct * 32], w_tile, src1_col, valid_rows, NULL);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct);
if (push_ct < ct_end) {
const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
size_t src1_row_size_padded;
- worker_callback_t quant_job_func;
+ worker_callback_t quant_task_func;
worker_callback_t matmul_job_func;
- uint32_t n_quant_jobs = 1;
+ uint32_t n_quant_tasks = 1;
if (src1_nrows > 1) {
if (is_repacked) {
switch (src0->type) {
switch (kparams->kernel_type) {
case HTP_MM_KERNEL_HVX_F16_F16_VTCM:
- quant_job_func = (src1->type == HTP_TYPE_F32) ? quantize_f32_f16_flat : quantize_f16_f16_flat;
+ quant_task_func = (src1->type == HTP_TYPE_F32) ? quantize_f32_f16_flat : quantize_f16_f16_flat;
mmctx->type = "f16-f16";
mmctx->vec_dot_1x1 = vec_dot_f16_f16_aa_1x1;
mmctx->vec_dot_2x1 = vec_dot_f16_f16_aa_2x1;
mmctx->mm_div_r2 = kparams->div_r2;
mmctx->mm_div_r3 = kparams->div_r3;
need_quant = false;
- quant_job_func = NULL;
+ quant_task_func = NULL;
src1_row_size = nb11;
break;
mmctx->mm_div_r3 = kparams->div_r3;
src1_row_size = nb11;
need_quant = false;
- quant_job_func = NULL;
+ quant_task_func = NULL;
break;
case HTP_MM_KERNEL_HVX_F32_F32_VTCM:
- quant_job_func = quantize_f32_f32_flat;
+ quant_task_func = quantize_f32_f32_flat;
mmctx->type = "f32-f32";
mmctx->vec_dot_1x1 = vec_dot_f32_f32_aa_1x1;
mmctx->vec_dot_2x1 = vec_dot_f32_f32_aa_2x1;
break;
case HTP_MM_KERNEL_HVX_F32_F32_DDR:
- quant_job_func = NULL;
+ quant_task_func = NULL;
mmctx->type = "f32-f32";
mmctx->vec_dot_1x1 = vec_dot_f32_f32_uu_1x1;
mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1;
break;
case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat;
src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
if (src1_nrows > 1) {
const uint32_t total_nb = src1_nrows * nb;
if (src1_nrows < octx->n_threads) {
- n_quant_jobs = MIN(total_nb, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
- for (uint32_t ith = 0; ith < n_quant_jobs; ++ith) {
- uint32_t ib_first = (total_nb * ith) / n_quant_jobs;
- uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_jobs;
+ n_quant_tasks = MIN(total_nb, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
+ for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) {
+ uint32_t ib_first = (total_nb * ith) / n_quant_tasks;
+ uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks;
mmctx->quant_ib_first[ith] = ib_first;
mmctx->quant_ib_last[ith] = ib_last;
mmctx->quant_r[ith] = ib_first / nb;
mmctx->quant_c[ith] = ib_first % nb;
}
} else {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
}
src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
break;
return HTP_STATUS_OK;
if (need_quant) {
- mmctx->src1_nrows_per_thread = (src1_nrows + n_quant_jobs - 1) / n_quant_jobs;
- worker_pool_run_func(octx->ctx->worker_pool, quant_job_func, mmctx, n_quant_jobs);
+ 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;
+ atomic_init(&mmctx->quant_barrier, n_quant_tasks);
+ } else {
+ mmctx->quant_task_func = NULL;
+ mmctx->n_quant_tasks = 0;
}
const uint32_t n_matmul_jobs = octx->n_threads;
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);
- if (src0_start_row >= src0_end_row) {
- return;
- }
-
const size_t dst_k_row_size = dst_k->nb[1]; // K and V share output width
const size_t dst_q_row_size = dst_q->nb[1]; // Q may be wider (GQA)
const size_t src0_row_size = src0->nb[1];
const uint8_t * restrict src3_row = (const uint8_t *) src3->data;
// Prefill spad with src0, src2, src3 rows
- for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
- const int is0 = (ir0 - src0_start_row);
- if (is0 >= (int)n_prefetch) {
- break;
+ if (src0_start_row < src0_end_row) {
+ for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
+ const int is0 = (ir0 - src0_start_row);
+ if (is0 >= (int)n_prefetch) {
+ break;
+ }
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
+ src0_stride, src0_row_size, src0_row_size, 2);
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size),
+ src2_stride, src2_row_size, src2_row_size, 2);
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + ir0 * src3_row_size),
+ src3_stride, src3_row_size, src3_row_size, 2);
}
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
- src0_stride, src0_row_size, src0_row_size, 2);
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size),
- src2_stride, src2_row_size, src2_row_size, 2);
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + ir0 * src3_row_size),
- src3_stride, src3_row_size, src3_row_size, 2);
+ }
+
+ hvx_mm_run_quant_task(mmctx, ith);
+
+ if (src0_start_row >= src0_end_row) {
+ return;
}
// Process rows
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);
- if (src0_start_row >= src0_end_row) {
- return;
- }
-
const size_t dst_row_size = dst_gate->nb[1];
const size_t src0_row_size = src0->nb[1];
const size_t src2_row_size = src2->nb[1];
const uint8_t * restrict src2_row = (const uint8_t *) src2->data;
// Prefill spad with src0, src2 rows
- for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
- const int is0 = (ir0 - src0_start_row);
- if (is0 >= (int)n_prefetch) {
- break;
+ if (src0_start_row < src0_end_row) {
+ for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
+ const int is0 = (ir0 - src0_start_row);
+ if (is0 >= (int)n_prefetch) {
+ break;
+ }
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
+ src0_stride, src0_row_size, src0_row_size, 2);
+ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size),
+ src2_stride, src2_row_size, src2_row_size, 2);
}
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size),
- src0_stride, src0_row_size, src0_row_size, 2);
- dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size),
- src2_stride, src2_row_size, src2_row_size, 2);
+ }
+
+ hvx_mm_run_quant_task(mmctx, ith);
+
+ if (src0_start_row >= src0_end_row) {
+ return;
}
// Process rows
size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task);
float *dst = st->dst + chunk_idx * st->dst_stride;
- transfer_output_chunk_fp16_to_fp32(dst, st->vtcm_src, chunk_idx, chunk_size, st->n_cols, st->dst_stride, st->dst_cols);
+ const float *src2 = st->src2 ? (st->src2 + chunk_idx * st->src2_stride) : NULL;
+ transfer_output_chunk_fp16_to_fp32(dst, src2, st->vtcm_src, chunk_idx, chunk_size, st->n_cols, st->dst_stride, st->src2_stride, st->dst_cols);
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx);
}
}
-static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const __fp16 *vtcm_src,
- int n_rows, int n_cols, int dst_stride, int dst_cols, int n_threads) {
+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;
state.n_tot_chunks = n_tot_chunks;
state.n_chunks_per_task = n_chunks_per_task;
state.dst = dst;
+ state.src2 = src2;
state.vtcm_src = vtcm_src;
state.n_cols = n_cols;
state.dst_stride = dst_stride;
+ state.src2_stride = src2_stride;
state.dst_cols = dst_cols;
state.traces = ctx->trace;
static int hmx_mm_2d_f32(struct htp_context *ctx,
float *restrict dst,
+ const float *restrict src2,
const float *activation,
const uint8_t *weight,
int m, int k, int n,
int weight_type,
int k_valid,
int dst_stride,
+ uint32_t src2_stride,
int dst_cols,
int m_chunk,
int n_chunk,
const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles);
- const bool is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32);
+ const bool is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32);
const size_t vec_dot_size = k * sizeof(__fp16);
const size_t vtcm_budget = ctx->vtcm_size;
+ const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride;
+ const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride;
+ const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride;
+
size_t m_chunk_n_rows = m_chunk;
size_t n_chunk_n_cols = n_chunk;
size_t vtcm_used = vtcm_size;
vtcm_weight_raw[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
}
}
- __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size);
- float *vtcm_f32_act = (float *) vtcm_seq_alloc(&vtcm_ptr, act_f32_size);
+
+ __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size);
+ float *vtcm_f32_act = (float *) vtcm_seq_alloc(&vtcm_ptr, act_f32_size);
__fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch0_size);
void *vtcm_scratch1 = scratch1_size ? vtcm_seq_alloc(&vtcm_ptr, scratch1_size) : NULL;
// 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);
- if (is_quant) {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), aligned_tile_size, tile_size, tile_size, (n_cols_A0 / 32) * n_k_tiles);
- } else {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), row_stride, weight_stride, row_stride, n_cols_A0);
- }
+ const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0);
if (1 < n_chunk_cnt) {
const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols);
- if (is_quant) {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), aligned_tile_size, tile_size, tile_size, (n_cols_A1 / 32) * n_k_tiles);
- } else {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), row_stride, weight_stride, row_stride, n_cols_A1);
- }
+ const uint32_t height_A1 = is_quant ? (n_cols_A1 / 32) * n_k_tiles : n_cols_A1;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1);
}
// pop A0 -> dequantize A0 -> submit C0
// 2. push A_{i+2} (if i+2 < n_chunk_cnt)
if (i + 2 < n_chunk_cnt) {
- if (is_quant) {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[(i + 2) % 2], weight + nc_p2 * weight_stride), aligned_tile_size, tile_size, tile_size, (n_cols_p2 / 32) * n_k_tiles);
- } else {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[(i + 2) % 2], weight + nc_p2 * weight_stride), row_stride, weight_stride, row_stride, n_cols_p2);
- }
+ const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[(i + 2) % 2], weight + nc_p2 * weight_stride),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2);
}
// 3. submit C_{i+1} (if i+1 < n_chunk_cnt)
// 4. wait C_i and store D_i (multi-thread HVX, parallel with C_{i+1})
hmx_queue_pop(ctx->hmx_queue);
float *output_chunk = dst + (mr * dst_stride + nc);
+ const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc) : NULL;
int chunk_dst_cols = dst_cols - (int)nc;
if (chunk_dst_cols > 0) {
- transfer_output_chunk_threaded(ctx, output_chunk, vtcm_output_bufs[i % 2], n_rows, n_cols, dst_stride, chunk_dst_cols, n_threads);
+ transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[i % 2], n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads);
}
}
}
const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS);
// A: Weight DMA (Synchronous)
- if (is_quant) {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc * weight_stride), aligned_tile_size, tile_size, tile_size, (n_cols / 32) * n_k_tiles);
- } else {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc * weight_stride), row_stride, weight_stride, row_stride, n_cols);
- }
+ const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc * weight_stride),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height);
dma_queue_pop(ctx->dma[0]);
// B: Weight Dequantize (Threaded)
// D: Output Store
float *output_chunk = dst + (mr * dst_stride + nc);
+ const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc) : NULL;
int chunk_dst_cols = dst_cols - (int)nc;
if (chunk_dst_cols > 0) {
- transfer_output_chunk_threaded(ctx, output_chunk, vtcm_output, n_rows, n_cols, dst_stride, chunk_dst_cols, n_threads);
+ transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output, n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads);
}
}
}
(size_t) dst_b3 * params->dst_nb3);
}
+static inline const float *hmx_mm_src2_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params,
+ int src2_b2, int src2_b3) {
+ return params->src2 ? (const float *) ((const uint8_t *) params->src2 +
+ (size_t) src2_b2 * params->src2_nb2 +
+ (size_t) src2_b3 * params->src2_nb3) : NULL;
+}
+
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) {
for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) {
for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) {
ret = hmx_mm_2d_f32(ctx, hmx_mm_dst_batch_ptr(params, b2, b3),
+ hmx_mm_src2_batch_ptr(params, b2, b3),
hmx_mm_activation_batch_ptr(params, b2, b3),
(const uint8_t *)hmx_mm_weight_batch_ptr(params, b2, b3),
params->m, params->k, params->n,
params->act_stride, params->weight_stride * (int)sizeof(__fp16),
- HTP_TYPE_F16, params->k, params->n, params->n,
+ 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);
}
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 * act_head_stride;
- if (use_dma_activation) {
- 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);
- } else {
- transfer_activation_chunk_threaded(ctx, vtcm_act_g,
- activation_chunk, (int) n_rows,
- params->k, params->act_stride, act_threads, params->k, NULL);
- }
+ 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);
}
void *buf_curr = vtcm_scratch0;
{
float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc;
+ const float *src2_chunk = params->src2 ? (hmx_mm_src2_batch_ptr(params, b2_base + g, b3) + mr * params->src2_stride + nc) : NULL;
int chunk_dst_cols = params->n - (int)nc;
if (chunk_dst_cols > 0) {
- transfer_output_chunk_threaded(ctx, output, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, chunk_dst_cols, ctx->n_threads);
+ transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
}
}
}
int tile_size = htp_mm_get_weight_tile_size(weight_type);
int aligned_tile_size = htp_mm_get_weight_aligned_tile_size(weight_type);
+ const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride;
+ const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride;
+ const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride;
+
const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0;
const size_t weight_row_stride = is_quant ? qweight_row_stride : row_stride;
size_t scratch0_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- uint8_t *vtcm_ptr = (uint8_t *) ctx->vtcm_base;
- __fp16 *vtcm_weight = weight_area_size ? (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size) : NULL;
- __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size);
- __fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
- void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch0_size);
- __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256);
+ uint8_t *vtcm_ptr = (uint8_t *) ctx->vtcm_base;
+ __fp16 *vtcm_weight = weight_area_size ? (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size) : NULL;
+ __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size);
+ __fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
+ void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch0_size);
+ __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256);
vtcm_used = vtcm_ptr - (uint8_t *) ctx->vtcm_base;
if (vtcm_used > vtcm_budget) {
const size_t n_cols = hex_smin((size_t) n - nc, n_chunk_n_cols);
const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS);
- if (is_quant) {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc * weight_stride), aligned_tile_size, tile_size, tile_size, (n_cols / 32) * n_k_tiles);
- } else {
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc * weight_stride), row_stride, weight_stride, row_stride, n_cols);
- }
+ const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc * weight_stride),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height);
dma_queue_pop(ctx->dma[0]);
dequantize_tiled_weight_chunk_to_fp16_tiles(
return HTP_STATUS_OK;
}
+ const float * src2_ptr = NULL;
+ uint32_t src2_stride = 0;
+ size_t src2_nb2 = 0;
+ size_t src2_nb3 = 0;
+ if (src2) {
+ src2_ptr = (const float *) src2->data;
+ src2_stride = (src2->ne[1] == 1) ? 0 : (uint32_t) (src2->nb[1] / sizeof(float));
+ src2_nb2 = (src2->ne[2] == 1) ? 0 : src2->nb[2];
+ src2_nb3 = (src2->ne[3] == 1) ? 0 : src2->nb[3];
+ }
+
int ret = -1;
const int n_threads = MIN(kparams->n_threads, (int) octx->n_threads);
if (kparams->kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) {
hmx_mm_f16_f32_batched_params_t batch_params = {
.dst = (float *) dst->data,
+ .src2 = src2_ptr,
.activation = (float *) src1->data,
.weight = (const __fp16 *) src0->data,
.m = m_total,
.act_stride = act_stride,
.weight_stride = wgt_stride,
.dst_stride = (int) (dst->nb[1] / sizeof(float)),
+ .src2_stride = src2_stride,
.ne02 = ne02,
.ne03 = ne03,
.ne12 = ne12,
.src1_nb3 = src1->nb[3],
.dst_nb2 = dst->nb[2],
.dst_nb3 = dst->nb[3],
+ .src2_nb2 = src2_nb2,
+ .src2_nb3 = src2_nb3,
};
ret = hmx_mm_f16_f32_batched(octx->ctx, &batch_params,
kparams->m_chunk, kparams->n_chunk,
kparams->vtcm_size);
} else {
ret = hmx_mm_2d_f32(
- octx->ctx, (float*) dst->data, (float*) src1->data, (const uint8_t *) src0->data,
+ octx->ctx, (float*) dst->data, src2_ptr, (float*) src1->data, (const uint8_t *) src0->data,
m_total, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0],
- (int)(dst->nb[1] / sizeof(float)), (int)dst->ne[0],
+ (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->tile_size, kparams->aligned_tile_size, kparams->vtcm_size
return HTP_STATUS_OK;
}
-static int hvx_mm_op_matmul_id(
+static int hvx_mm_matmul_id(
struct htp_ops_context * octx,
struct htp_mm_context * mmctx,
size_t src0_row_size_padded,
const uint32_t nb = (ne10 + qk - 1) / qk;
const uint32_t total_nb = src1_nrows * nb;
- worker_callback_t quant_job_func;
- uint32_t n_quant_jobs = 1;
+ worker_callback_t quant_task_func;
+ uint32_t n_quant_tasks = 1;
if (src1_nrows < octx->n_threads) {
- n_quant_jobs = MIN(total_nb, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
- for (uint32_t ith = 0; ith < n_quant_jobs; ++ith) {
- uint32_t ib_first = (total_nb * ith) / n_quant_jobs;
- uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_jobs;
+ n_quant_tasks = MIN(total_nb, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
+ for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) {
+ uint32_t ib_first = (total_nb * ith) / n_quant_tasks;
+ uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks;
mmctx->quant_ib_first[ith] = ib_first;
mmctx->quant_ib_last[ith] = ib_last;
mmctx->quant_r[ith] = ib_first / nb;
mmctx->quant_c[ith] = ib_first % nb;
}
} else {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
}
size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
size_t src0_sz = kparams->vtcm_src0_size;
size_t src1_sz = kparams->vtcm_src1_size;
size_t src2_sz = 0; // mapping lives in DDR
- size_t dst_sz = 0; // ID kernels scatter straight to DDR
+ size_t dst_sz = kparams->vtcm_dst_size;
size_t vtcm_size = kparams->vtcm_size;
size_t src0_sz_per_thread = src0_sz / octx->n_threads;
size_t src1_sz_per_thread = src1_sz;
size_t src2_sz_per_thread = 0;
- size_t dst_sz_per_thread = 0;
+ size_t dst_sz_per_thread = dst_sz / octx->n_threads;
FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size %zu dst-spad-size %zu (%zu)\n", mmctx->type,
src0_sz, src1_sz, src2_sz, dst_sz, vtcm_size);
mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread;
- mmctx->src1_nrows_per_thread = (src1_nrows + n_quant_jobs - 1) / n_quant_jobs;
- worker_pool_run_func(octx->ctx->worker_pool, quant_job_func, mmctx, n_quant_jobs);
+ 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;
+ 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);
const uint32_t src0_nrows = ne01; // per expert
const uint32_t src1_nrows = ne11 * ne12 * ne13;
- worker_callback_t quant_job_func;
+ 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
return hmx_mm_op_matmul_id(octx, mmctx, matrix_row_counts, matrix_rows, mapping_buf, must_free_mapping);
}
- return hvx_mm_op_matmul_id(octx, mmctx, src0_row_size_padded, src1_nrows, matmul_id_job_func, 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);
}
int op_matmul_qkv(struct htp_ops_context * octx) {
const uint32_t nb = (src1->ne[0] + qk - 1) / qk;
const uint32_t total_nb = src1_nrows * nb;
- worker_callback_t quant_job_func;
- uint32_t n_quant_jobs = 1;
+ worker_callback_t quant_task_func;
+ uint32_t n_quant_tasks = 1;
if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat;
} else if (src1_nrows < octx->n_threads) {
- n_quant_jobs = MIN(total_nb, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
- for (uint32_t ith = 0; ith < n_quant_jobs; ++ith) {
- uint32_t ib_first = (total_nb * ith) / n_quant_jobs;
- uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_jobs;
+ n_quant_tasks = MIN(total_nb, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
+ for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) {
+ uint32_t ib_first = (total_nb * ith) / n_quant_tasks;
+ uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks;
mmctx->quant_ib_first[ith] = ib_first;
mmctx->quant_ib_last[ith] = ib_last;
mmctx->quant_r[ith] = ib_first / nb;
mmctx->quant_c[ith] = ib_first % nb;
}
} else {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
}
size_t src1_row_size;
size_t src1_sz = kparams->vtcm_src1_size;
size_t src2_sz = kparams->vtcm_src2_size;
size_t src3_sz = kparams->vtcm_src3_size;
+ size_t dst_sz = kparams->vtcm_dst_size;
size_t vtcm_size = kparams->vtcm_size;
size_t src0_sz_per_thread = src0_sz / octx->n_threads;
size_t src1_sz_per_thread = src1_sz;
size_t src2_sz_per_thread = src2_sz / octx->n_threads;
size_t src3_sz_per_thread = src3_sz / octx->n_threads;
+ size_t dst_sz_per_thread = dst_sz / octx->n_threads;
if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-qkv: current VTCM reservation %zu is too small, needed %zu\n",
mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
mmctx->vtcm_src3 = vtcm_seq_alloc(&vtcm_ptr, src3_sz);
+ mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
octx->src1_spad.src = NULL;
octx->src0_spad.src = NULL;
octx->src2_spad.src = NULL;
octx->src3_spad.src = NULL;
+ octx->dst_spad.src = NULL;
mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
mmctx->vtcm_src3_size_per_thread = src3_sz_per_thread;
+ mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread;
if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
return HTP_STATUS_OK;
- // Run quantization once
- mmctx->src1_nrows_per_thread = (src1_nrows + n_quant_jobs - 1) / n_quant_jobs;
- worker_pool_run_func(octx->ctx->worker_pool, quant_job_func, mmctx, n_quant_jobs);
+ 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;
+ atomic_init(&mmctx->quant_barrier, n_quant_tasks);
// Run fused matmul
const uint32_t n_matmul_jobs = octx->n_threads;
const uint32_t nb = (src1->ne[0] + qk - 1) / qk;
const uint32_t total_nb = src1_nrows * nb;
- worker_callback_t quant_job_func;
- uint32_t n_quant_jobs = 1;
+ worker_callback_t quant_task_func;
+ uint32_t n_quant_tasks = 1;
if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat;
} else if (src1_nrows < octx->n_threads) {
- n_quant_jobs = MIN(total_nb, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
- for (uint32_t ith = 0; ith < n_quant_jobs; ++ith) {
- uint32_t ib_first = (total_nb * (ith + 0)) / n_quant_jobs;
- uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_jobs;
+ n_quant_tasks = MIN(total_nb, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block;
+ for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) {
+ uint32_t ib_first = (total_nb * (ith + 0)) / n_quant_tasks;
+ uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks;
mmctx->quant_ib_first[ith] = ib_first;
mmctx->quant_ib_last[ith] = ib_last;
mmctx->quant_r[ith] = ib_first / nb;
mmctx->quant_c[ith] = ib_first % nb;
}
} else {
- n_quant_jobs = MIN(src1_nrows, octx->n_threads);
- quant_job_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
+ n_quant_tasks = MIN(src1_nrows, octx->n_threads);
+ quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled;
}
size_t src1_row_size;
size_t src0_sz = kparams->vtcm_src0_size;
size_t src1_sz = kparams->vtcm_src1_size;
size_t src2_sz = kparams->vtcm_src2_size;
+ size_t dst_sz = kparams->vtcm_dst_size;
size_t vtcm_size = kparams->vtcm_size;
size_t src0_sz_per_thread = src0_sz / octx->n_threads;
size_t src1_sz_per_thread = src1_sz;
size_t src2_sz_per_thread = src2_sz / octx->n_threads;
+ size_t dst_sz_per_thread = dst_sz / octx->n_threads;
if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-ffn: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size);
mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
+ mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
octx->src1_spad.src = NULL;
octx->src0_spad.src = NULL;
octx->src2_spad.src = NULL;
+ octx->dst_spad.src = NULL;
mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
+ mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread;
if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
return HTP_STATUS_OK;
- // Run quantization once
- mmctx->src1_nrows_per_thread = (src1_nrows + n_quant_jobs - 1) / n_quant_jobs;
- worker_pool_run_func(octx->ctx->worker_pool, quant_job_func, mmctx, n_quant_jobs);
+ 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;
+ atomic_init(&mmctx->quant_barrier, n_quant_tasks);
// Run fused matmul
const uint32_t n_matmul_jobs = octx->n_threads;
case HTP_MM_KERNEL_HVX_QUANT_ROW: {
size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
- vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
- // src0 spad is also used in dynamic quantizer to store padded src1 rows
- size_t src1_row_size_padded = htp_mm_round_up(q_src1_row_size, QK_Q8_0_TILED * sizeof(float));
- if (vtcm_src0_size < src1_row_size_padded) {
- vtcm_src0_size = src1_row_size_padded;
- }
-
vtcm_src0_size = vtcm_src0_size * n_threads;
- vtcm_dst_size = vtcm_dst_size * n_threads;
if (is_repack) {
uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
uint32_t n_k_tiles = ne10 / 32;
uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
vtcm_src0_size = repacked_vtcm_size * n_threads;
}
+
+ size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+ size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+ if (dst_size_per_thread < quant_scratch_size_per_thread) {
+ dst_size_per_thread = quant_scratch_size_per_thread;
+ }
+ vtcm_dst_size = dst_size_per_thread * n_threads;
break;
}
case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
- vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
- size_t src1_row_size_padded = htp_mm_round_up(q_src1_row_size, 256);
- if (vtcm_src0_size < src1_row_size_padded) {
- vtcm_src0_size = src1_row_size_padded;
- }
-
vtcm_src0_size = vtcm_src0_size * n_threads;
- vtcm_dst_size = vtcm_dst_size * n_threads;
if (is_repack) {
uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
uint32_t n_k_tiles = ne10 / 32;
uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
vtcm_src0_size = repacked_vtcm_size * n_threads;
}
+
+ size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+ size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+ if (dst_size_per_thread < quant_scratch_size_per_thread) {
+ dst_size_per_thread = quant_scratch_size_per_thread;
+ }
+ vtcm_dst_size = dst_size_per_thread * n_threads;
break;
}
default:
size_t src0_row_size, // nb01
uint32_t n_prefetch,
size_t * vtcm_src0_size_out,
- size_t * vtcm_src1_size_out
+ size_t * vtcm_src1_size_out,
+ size_t * vtcm_dst_size_out
) {
const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
size_t src1_sz = htp_mm_round_up(src1_row_size * src1_nrows, 256);
- // src0 spad also holds temporary transposed src1 columns during dynamic quantization.
- const size_t src1_row_size_padded = htp_mm_round_up(src1_row_size, QK_Q8_0_TILED * sizeof(float));
- if (src0_sz_per_thread < src1_row_size_padded) {
- src0_sz_per_thread = src1_row_size_padded;
- }
-
if (is_repack) {
const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
const uint32_t n_k_tiles = ne10 / 32;
const uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
- if (repacked_vtcm_size < src1_row_size_padded) {
- repacked_vtcm_size = src1_row_size_padded;
- }
src0_sz_per_thread = repacked_vtcm_size;
}
const size_t vtcm_src0_size = src0_sz_per_thread * n_threads;
+ const size_t vtcm_dst_size = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
*vtcm_src0_size_out = vtcm_src0_size;
*vtcm_src1_size_out = src1_sz;
+ *vtcm_dst_size_out = vtcm_dst_size;
- return vtcm_src0_size + src1_sz;
+ return vtcm_src0_size + src1_sz + vtcm_dst_size;
}
#ifdef __cplusplus
mmsel=
[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM"
+fasel=
+[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA"
+
set -x
adb $adbserial $adbhost shell " \
cd $basedir; ulimit -c unlimited; \
LD_LIBRARY_PATH=$basedir/$branch/lib \
ADSP_LIBRARY_PATH=$basedir/$branch/lib \
- $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opflt $opfuse $vmem $mbuf $mmsel \
+ $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opflt $opfuse $vmem $mbuf $mmsel $fasel \
./$branch/bin/llama-completion --no-mmap -m $basedir/../gguf/$model \
--poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \
--ctx-size 8192 --ubatch-size 1024 -fa on \
mmsel=
[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM"
+fasel=
+[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA"
+
set -x
tool=$1; shift
cd $basedir; ulimit -c unlimited; \
LD_LIBRARY_PATH=$basedir/$branch/lib \
ADSP_LIBRARY_PATH=$basedir/$branch/lib \
- $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opfuse $mmsel ./$branch/bin/$tool $@ \
+ $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opfuse $mmsel $fasel ./$branch/bin/$tool $@ \
"
char = 'Q'
elif norm_evt == 'A-PREP':
char = 'A'
+ elif norm_evt == 'Q-PREP':
+ char = 'q'
+ elif norm_evt == 'K-PREP':
+ char = 'k'
+ elif norm_evt == 'V-PREP':
+ char = 'v'
elif norm_evt == 'W-DEQUANT':
char = 'D'
elif norm_evt == 'O-PROC':
overview / pkg / ggml / sources / llama.cpp / commitdiff