From 13e673863b36993d9e80c0a9ee3c1e02f8b9ace4 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Wed, 1 Jul 2026 06:59:19 -0700 Subject: [PATCH] hexagon: flash attention rework (optimizations, accuracy improvements, etc) (#25085) * hex-mm: fold mm quant tasks into the main matmul threads * hex-mm: minor formatting fixes * hex-mm: cleanup is_quant checks in dma dispatch * hex-mm: fix dst-spad alignment * hex-mm: move fp kernels in the hvx-mm-kernels header * hex-mm: fuse with ADD * hex-fa: factor out ukernels into separate headers and unify the rest * hex-fa: move kernel-params compute into the host * hex-fa: refactor vtcm alloc for consistency * hex-fa: add support for FA_SELECT * hex-fa: update tracing insrumentation to cover all functions * hex-fa: update hvx fallback thresholds to recover t/g regressions * hex-fa: update tracing instrumentation * hex-fa: improved tracing with additional events * hex-fa: optimize mask processing (fastdiv, etc) * hex-fa: improve mask dma caching * hmx-fa: change loop order to maximize mask cache hits * hex-fa: remove over instrumentation * hex-fa: breakdown QKV prep trace events * hmx-fa: further mask proc optimizations * hex-fa: mask broadcast is the common case, optimize for that * hex-fa: use aligned loads where possible * hex-fa: update loops to use uint32_t indices * hmx-fa: fold vtcm init into q prep task * hex-fa: update rest of the hmx funcs to use uint32_t * hmx-fa: fold build_d into the main softmax loop * hmx-fa: start kv dmas earlier * hmx-fa: start mask dma a bit earlier * hex-fa: precompute rows per task to avoid divs * hmx-fa: specialize fa_o_store for f16 and f32 * hmx-fa: prelim support for Sinks * hmx-fa: keep softmax accumulators in fp32 * hex-fa: add tanh_f16 and exp2_f16 and use that in FA * hex-fa: use fp16 math in the hvx kernel * hex-fa: avoid expensive float -> __fp16 cast for slopes and softcap * hex-fa: replace most vec_exp_f32 with vec_exp2_f16 * hmx-fa: vectorize sinks update * hex-fa: minor formatting * hmx-fa: fold softcap loop into the tile load * hmx-fa: use vectoralias to populate sinks * hex-fa: remove redudant check * hex-fa: fix vtcm size compute to use fp32 for accumulators * hex-mm: fix trailing spaces * hmx-fa: dont use -inf to init mask to avoid conversion overflows * hex-fa: no need to explicitly guard -inf in the f16->f32 converter now * hmx-fa: cleanup fa sinks handling * hex-mm: fixed src2 stride handling when mm is fused with add * hex-fa: make lto happy --- ggml/src/ggml-hexagon/CMakeLists.txt | 1 - ggml/src/ggml-hexagon/ggml-hexagon.cpp | 247 +- ggml/src/ggml-hexagon/htp-opnode.h | 14 +- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 9 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 2171 ++++++++++++++--- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 253 ++ ggml/src/ggml-hexagon/htp/hex-dma.h | 30 +- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 96 + .../src/ggml-hexagon/htp/hmx-flash-attn-ops.c | 1838 -------------- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 45 +- ggml/src/ggml-hexagon/htp/hmx-utils.h | 70 +- ggml/src/ggml-hexagon/htp/htp-ops.h | 6 + ggml/src/ggml-hexagon/htp/hvx-base.h | 13 +- ggml/src/ggml-hexagon/htp/hvx-exp.h | 39 + ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h | 232 ++ .../ggml-hexagon/htp/hvx-mm-kernels-flat.h | 537 +++- .../ggml-hexagon/htp/hvx-mm-kernels-tiled.h | 110 +- ggml/src/ggml-hexagon/htp/hvx-sigmoid.h | 39 + ggml/src/ggml-hexagon/htp/main.c | 1 + ggml/src/ggml-hexagon/htp/matmul-ops.c | 982 +++----- ggml/src/ggml-hexagon/htp/matmul-ops.h | 51 +- scripts/snapdragon/adb/run-completion.sh | 5 +- scripts/snapdragon/adb/run-tool.sh | 5 +- scripts/snapdragon/ggml-hexagon-profile.py | 6 + 24 files changed, 3777 insertions(+), 3023 deletions(-) create mode 100644 ggml/src/ggml-hexagon/htp/flash-attn-ops.h create mode 100644 ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h delete mode 100644 ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c create mode 100644 ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h diff --git a/ggml/src/ggml-hexagon/CMakeLists.txt b/ggml/src/ggml-hexagon/CMakeLists.txt index c6e49a71d..845b922bf 100644 --- a/ggml/src/ggml-hexagon/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/CMakeLists.txt @@ -23,7 +23,6 @@ include(${HEXAGON_SDK_ROOT}/build/cmake/hexagon_fun.cmake) 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 diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 3d41c47b6..4f9dafed0 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -43,6 +43,7 @@ #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" @@ -62,6 +63,7 @@ static int opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu) 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 @@ -125,6 +127,11 @@ static const char * htp_event_name(uint16_t id) { 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"; } @@ -1879,6 +1886,162 @@ ggml_hexagon_session::~ggml_hexagon_session() noexcept(true) { // ** 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]; @@ -1912,6 +2075,17 @@ static bool ggml_hexagon_supported_flash_attn_ext(const struct ggml_hexagon_sess 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; } @@ -2211,14 +2385,14 @@ static void ggml_hexagon_precompute_hvx_mm_params( 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; @@ -2228,14 +2402,14 @@ static void ggml_hexagon_precompute_hvx_mm_params( 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) { @@ -2441,11 +2615,12 @@ static void ggml_hexagon_precompute_fused_qkv_params( 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; @@ -2453,13 +2628,10 @@ static void ggml_hexagon_precompute_fused_qkv_params( 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; @@ -2471,9 +2643,6 @@ static void ggml_hexagon_precompute_fused_qkv_params( } 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); @@ -2492,7 +2661,7 @@ static void ggml_hexagon_precompute_fused_qkv_params( 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; @@ -2500,6 +2669,7 @@ static void ggml_hexagon_precompute_fused_qkv_params( 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 { @@ -2510,7 +2680,8 @@ static void ggml_hexagon_precompute_fused_qkv_params( 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; } } @@ -2536,11 +2707,12 @@ static void ggml_hexagon_precompute_fused_ffn_params( 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; @@ -2548,12 +2720,9 @@ static void ggml_hexagon_precompute_fused_ffn_params( 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; @@ -2564,9 +2733,6 @@ static void ggml_hexagon_precompute_fused_ffn_params( } 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); } @@ -2582,13 +2748,14 @@ static void ggml_hexagon_precompute_fused_ffn_params( 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 { @@ -2598,7 +2765,8 @@ static void ggml_hexagon_precompute_fused_ffn_params( 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; } } @@ -3243,7 +3411,7 @@ static inline bool op_is_compute(ggml_tensor *node) 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]; @@ -3262,7 +3430,7 @@ static bool is_hmx_eligible(const ggml_tensor * t) { 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) { @@ -3357,6 +3525,26 @@ static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph * } } + 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; } @@ -3393,6 +3581,11 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg 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)); } @@ -4079,6 +4272,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { 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"); @@ -4120,6 +4314,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { 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; diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index 6fe23b0d6..19a2504c7 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -11,6 +11,7 @@ #include #include "htp-ops.h" #include "htp/matmul-ops.h" +#include "htp/flash-attn-ops.h" struct htp_opnode { ggml_tensor * node = nullptr; @@ -335,7 +336,8 @@ struct htp_opformat { } 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; @@ -350,6 +352,16 @@ struct htp_opformat { 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, "----"); } diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index c48a5b86e..9f9471f69 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -20,9 +20,6 @@ add_library(${HTP_LIB} SHARED 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 @@ -42,16 +39,14 @@ add_library(${HTP_LIB} SHARED 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 $,HTP_DEBUG=1,NDEBUG=1> $,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}) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 65f7844ae..a1d15fc7a 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -3,229 +3,40 @@ #pragma clang diagnostic ignored "-Wunused-but-set-variable" #include +#include #include #include #include +#include +#include +#include #include #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; @@ -242,12 +53,12 @@ struct htp_fa_context { 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; @@ -264,28 +75,76 @@ struct htp_fa_context { bool is_q_fp32; + size_t size_q_block; + size_t size_vkq_acc; + + uint8_t * spad_q; + uint8_t * spad_k; + uint8_t * spad_v; + uint8_t * spad_m; + uint8_t * spad_a; + 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); +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; - const HVX_Vector * restrict vsrc = (const HVX_Vector * restrict) src; - HVX_Vector * restrict vdst = (HVX_Vector * restrict) dst; + // Op parameters + __fp16 scale; + float max_bias; + __fp16 logit_softcap; + uint32_t n_head_log2; + float m0, m1; - const uint32_t nvec = n / VLEN_FP32; - const uint32_t nloe = n % VLEN_FP32; + // 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 - 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)); - } -} + // 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; @@ -341,7 +200,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * 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]; @@ -353,16 +211,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * 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); @@ -379,9 +235,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * 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); @@ -410,18 +263,13 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * // 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); @@ -433,6 +281,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * } 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); @@ -442,113 +291,101 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * 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); - } - - // 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 - } - - // 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)); - } - - sb_scores[iv] = scores; - v_max = hvx_vec_reduce_max2_f32(scores, v_max); // All lanes have block max + // 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); } + // 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); + } + + HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(current_block_size * sizeof(__fp16)); + + // 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; - } - - 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; - } - - // 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; - } - - 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); - } - } - - 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); + + // 5. Accumulate V (F16 * F16 -> F32 accumulator) + __fp16 __attribute__((aligned(128))) p_arr[VLEN_FP16]; + hvx_vec_store_a(p_arr, 128, P); + + 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; + } + + if (p_arr[j] == 0.0f && p_arr[j + 1] == 0.0f) { + continue; + } + + 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) { @@ -569,14 +406,10 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * 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); @@ -605,9 +438,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * // 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] @@ -618,8 +451,1563 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * } 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); } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); +} + +// ============================================================================ +// 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); + } +} + +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) { @@ -634,13 +2022,14 @@ 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; @@ -648,94 +2037,68 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { 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; + factx.n_blocks = kparams->n_kv_blocks; - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; + factx.scale = kparams->scale; + factx.max_bias = kparams->max_bias; + factx.logit_softcap = (__fp16) kparams->logit_softcap; - 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.scale = scale; - factx.max_bias = max_bias; - factx.logit_softcap = 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); } diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h new file mode 100644 index 000000000..fbae247e0 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -0,0 +1,253 @@ +#ifndef HTP_FLASH_ATTN_OPS_H +#define HTP_FLASH_ATTN_OPS_H + +#include +#include +#include + +#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 */ diff --git a/ggml/src/ggml-hexagon/htp/hex-dma.h b/ggml/src/ggml-hexagon/htp/hex-dma.h index 8031a5679..98fcc9fda 100644 --- a/ggml/src/ggml-hexagon/htp/hex-dma.h +++ b/ggml/src/ggml-hexagon/htp/hex-dma.h @@ -138,27 +138,28 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t } 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; } @@ -320,7 +321,7 @@ static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_ 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; @@ -352,20 +353,19 @@ static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * sr 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 diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h new file mode 100644 index 000000000..1ab286e1c --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -0,0 +1,96 @@ +#ifndef HMX_FA_KERNELS_H +#define HMX_FA_KERNELS_H + +#include +#include +#include +#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 */ diff --git a/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c deleted file mode 100644 index 996fd5975..000000000 --- a/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c +++ /dev/null @@ -1,1838 +0,0 @@ -// 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 -#include -#include -#include -#include -#include -#include -#include - -#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; -} diff --git a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h index b7fba22a8..4a0ca7885 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -712,7 +712,17 @@ static inline void hmx_matmul_job_init(hmx_matmul_job_t * job, // 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; @@ -727,6 +737,7 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 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) { @@ -738,9 +749,20 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 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; } } @@ -752,9 +774,20 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 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); } } } @@ -763,11 +796,13 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 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; diff --git a/ggml/src/ggml-hexagon/htp/hmx-utils.h b/ggml/src/ggml-hexagon/htp/hmx-utils.h index f448ee337..561d26944 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-utils.h +++ b/ggml/src/ggml-hexagon/htp/hmx-utils.h @@ -42,14 +42,14 @@ static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VL // 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); @@ -65,14 +65,14 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, 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; + 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 (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; + 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); @@ -86,7 +86,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, 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); @@ -95,7 +95,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, } } 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); @@ -105,14 +105,14 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, } } 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; + 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 (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; + 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); @@ -122,7 +122,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, 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); @@ -131,7 +131,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, } } 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); @@ -148,24 +148,24 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, // 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. @@ -174,7 +174,7 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out, 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); @@ -185,7 +185,7 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out, } } 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); diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index d04090135..c9d0b3539 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -60,6 +60,7 @@ enum htp_op_code { 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, @@ -175,6 +176,11 @@ enum htp_trace_event_id { 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, }; diff --git a/ggml/src/ggml-hexagon/htp/hvx-base.h b/ggml/src/ggml-hexagon/htp/hvx-base.h index 493b26c6e..d7575b9ff 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-base.h +++ b/ggml/src/ggml-hexagon/htp/hvx-base.h @@ -134,16 +134,7 @@ static inline HVX_Vector hvx_vec_f32_to_f16_shuff(HVX_Vector v0, HVX_Vector v1) } 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 @@ -170,8 +161,6 @@ static inline HVX_VectorPair hvx_vec_f16_to_f32(HVX_Vector v) { } #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) diff --git a/ggml/src/ggml-hexagon/htp/hvx-exp.h b/ggml/src/ggml-hexagon/htp/hvx-exp.h index e71ec4909..bcd3d2d32 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-exp.h +++ b/ggml/src/ggml-hexagon/htp/hvx-exp.h @@ -16,6 +16,7 @@ #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)) @@ -213,4 +214,42 @@ static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict } } +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 */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h new file mode 100644 index 000000000..c05bd0b85 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/hvx-fa-kernels.h @@ -0,0 +1,232 @@ +#ifndef HVX_FA_KERNELS_H +#define HVX_FA_KERNELS_H + +#include +#include +#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 */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h index 52351b103..328a83118 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h @@ -256,7 +256,7 @@ static inline void quantize_f16_f16_flat_kernel( // 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; @@ -312,10 +312,14 @@ static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const v 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; @@ -397,11 +401,19 @@ static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float 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; @@ -464,10 +476,14 @@ static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const v 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; @@ -561,11 +577,19 @@ static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float 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; @@ -620,10 +644,14 @@ static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const v 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; @@ -704,11 +732,19 @@ static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float 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; @@ -765,10 +801,14 @@ static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const 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; @@ -851,11 +891,19 @@ static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float 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; @@ -921,10 +969,14 @@ static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const 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; @@ -1019,6 +1071,441 @@ static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float 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)); + } + } +} + diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h index bcb0b8f9e..40b65aa3b 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h @@ -378,7 +378,7 @@ static inline HVX_VectorPair accum_q8_0_32x2( 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; @@ -401,10 +401,14 @@ static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const 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; @@ -484,11 +488,19 @@ static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float 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; @@ -519,10 +531,14 @@ static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const 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; @@ -637,11 +653,19 @@ static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float 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; @@ -663,10 +687,14 @@ static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const 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; @@ -745,11 +773,19 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float 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; @@ -773,10 +809,14 @@ static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const 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; @@ -857,11 +897,19 @@ static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, floa 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; @@ -896,10 +944,14 @@ static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const 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; @@ -1013,8 +1065,16 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa 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( diff --git a/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h b/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h index 37f3e7b6f..dd66dd84c 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h +++ b/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h @@ -3,6 +3,7 @@ #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 @@ -139,4 +140,42 @@ static inline void hvx_tanh_f32_aa(uint8_t * restrict dst, const uint8_t * restr 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 */ diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index d76512ea4..684625e42 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -575,6 +575,7 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) { 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: diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 81a0ffbeb..79c292dc2 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -9,6 +9,7 @@ #include #include +#include #include "hex-dma.h" #include "hvx-utils.h" @@ -21,8 +22,20 @@ #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; @@ -31,6 +44,7 @@ typedef struct { int act_stride; int weight_stride; int dst_stride; + uint32_t src2_stride; int ne02; int ne03; int ne12; @@ -41,6 +55,8 @@ typedef struct { 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 { @@ -61,11 +77,11 @@ 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; @@ -73,11 +89,16 @@ struct htp_mm_context { 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; @@ -133,382 +154,7 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = { 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]; \ @@ -526,10 +172,10 @@ static void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const v 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]; \ @@ -558,6 +204,18 @@ static void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const v 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) { @@ -631,6 +289,9 @@ 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" @@ -648,10 +309,6 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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; \ @@ -661,6 +318,7 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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; \ @@ -680,9 +338,17 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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++) { \ @@ -702,7 +368,15 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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) { \ @@ -710,7 +384,12 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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); \ \ @@ -731,10 +410,6 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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; \ @@ -752,6 +427,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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; \ \ @@ -767,9 +443,17 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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++) { \ @@ -780,7 +464,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void 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) { \ @@ -792,7 +476,24 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void \ 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); \ + } \ } \ } @@ -850,17 +551,22 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, 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; \ @@ -885,8 +591,8 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, 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) { \ @@ -898,8 +604,8 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, 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); \ \ @@ -948,7 +654,7 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int 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) { \ @@ -957,7 +663,7 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, 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); \ \ @@ -1019,11 +725,19 @@ static void hvx_mm_ffn_2d_repacked_##SUFFIX(unsigned int nth, unsigned int 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); \ - 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++) { \ @@ -1049,8 +763,8 @@ static void hvx_mm_ffn_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, 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) { \ @@ -1062,8 +776,8 @@ static void hvx_mm_ffn_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, 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); \ \ @@ -1099,7 +813,7 @@ static void name(unsigned int nth, unsigned int ith, void * data) { 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) { \ @@ -1115,7 +829,7 @@ static void name(unsigned int nth, unsigned int ith, void * data) { 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); \ @@ -1140,7 +854,7 @@ static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, vo 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], @@ -1164,7 +878,7 @@ static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, vo 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], @@ -1230,11 +944,6 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { 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; @@ -1252,14 +961,21 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { 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 @@ -1311,6 +1027,9 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { } 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) { @@ -1321,11 +1040,6 @@ 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; @@ -1354,14 +1068,21 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { 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 @@ -1392,7 +1113,25 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { 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)] @@ -1410,7 +1149,8 @@ static void hvx_mm_id(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 + hvx_mm_run_quant_task(mmctx, ith); + if (src0_start_row >= src0_end_row) { return; } @@ -1476,7 +1216,7 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { 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); @@ -1498,7 +1238,8 @@ static void hvx_mv_id(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 + hvx_mm_run_quant_task(mmctx, ith); + if (src0_start_row >= src0_end_row) { return; } @@ -1555,7 +1296,7 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { 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) { @@ -1625,9 +1366,9 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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) { @@ -1660,7 +1401,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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; @@ -1677,7 +1418,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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; @@ -1691,11 +1432,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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; @@ -1704,7 +1445,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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; @@ -1717,8 +1458,8 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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) { @@ -1755,19 +1496,19 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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; @@ -1831,8 +1572,13 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { 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; @@ -1870,10 +1616,6 @@ static void hvx_mm_qkv_2d(unsigned int nth, unsigned int ith, void * data) { 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]; @@ -1902,17 +1644,25 @@ static void hvx_mm_qkv_2d(unsigned int nth, unsigned int ith, void * data) { 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 @@ -2024,10 +1774,6 @@ static void hvx_mm_ffn_2d(unsigned int nth, unsigned int ith, void * data) { 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]; @@ -2051,15 +1797,23 @@ static void hvx_mm_ffn_2d(unsigned int nth, unsigned int ith, void * data) { 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 @@ -2187,7 +1941,8 @@ static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void 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); @@ -2316,8 +2071,8 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles( } } -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; @@ -2333,9 +2088,11 @@ static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, 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; @@ -2384,6 +2141,7 @@ static void transfer_activation_chunk_threaded( 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, @@ -2392,6 +2150,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, int weight_type, int k_valid, int dst_stride, + uint32_t src2_stride, int dst_cols, int m_chunk, int n_chunk, @@ -2425,10 +2184,14 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, 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; @@ -2458,8 +2221,9 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, 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; @@ -2494,19 +2258,15 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // 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 @@ -2543,11 +2303,9 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // 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) @@ -2562,9 +2320,10 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // 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); } } } @@ -2583,11 +2342,9 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, 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) @@ -2601,9 +2358,10 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // 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); } } } @@ -2644,6 +2402,13 @@ static inline float *hmx_mm_dst_batch_ptr(const hmx_mm_f16_f32_batched_params_t (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) { @@ -2651,11 +2416,12 @@ static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, 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); } @@ -2740,15 +2506,9 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ 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; @@ -2793,9 +2553,10 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ { 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); } } } @@ -2956,6 +2717,10 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, 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; @@ -2978,12 +2743,12 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, 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) { @@ -3007,11 +2772,9 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, 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( @@ -3051,11 +2814,23 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k 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, @@ -3064,6 +2839,7 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k .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, @@ -3074,6 +2850,8 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k .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, @@ -3082,9 +2860,9 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k 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 @@ -3145,7 +2923,7 @@ static int hmx_mm_op_matmul_id( 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, @@ -3163,22 +2941,22 @@ static int hvx_mm_op_matmul_id( 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); @@ -3188,13 +2966,13 @@ static int hvx_mm_op_matmul_id( 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); @@ -3230,8 +3008,10 @@ static int hvx_mm_op_matmul_id( 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); @@ -3259,7 +3039,7 @@ int op_matmul_id(struct htp_ops_context * octx) { 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 @@ -3329,7 +3109,7 @@ int op_matmul_id(struct htp_ops_context * octx) { 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) { @@ -3373,25 +3153,25 @@ 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; @@ -3406,12 +3186,14 @@ int op_matmul_qkv(struct htp_ops_context * octx) { 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", @@ -3424,11 +3206,13 @@ int op_matmul_qkv(struct htp_ops_context * octx) { 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; @@ -3439,13 +3223,15 @@ int op_matmul_qkv(struct htp_ops_context * octx) { 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; @@ -3517,25 +3303,25 @@ int op_matmul_ffn(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 + 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; @@ -3549,11 +3335,13 @@ int op_matmul_ffn(struct htp_ops_context * octx) { 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); @@ -3564,10 +3352,12 @@ int op_matmul_ffn(struct htp_ops_context * octx) { 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; @@ -3576,13 +3366,15 @@ int op_matmul_ffn(struct htp_ops_context * octx) { 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; diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index a94d5430d..96369825a 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -392,56 +392,49 @@ static inline size_t htp_mm_hvx_get_vtcm_sizes( 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: @@ -463,7 +456,8 @@ static inline size_t htp_mm_hvx_id_get_vtcm_sizes( 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 || @@ -476,29 +470,22 @@ static inline size_t htp_mm_hvx_id_get_vtcm_sizes( 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 diff --git a/scripts/snapdragon/adb/run-completion.sh b/scripts/snapdragon/adb/run-completion.sh index f7622eb52..2130b9a74 100755 --- a/scripts/snapdragon/adb/run-completion.sh +++ b/scripts/snapdragon/adb/run-completion.sh @@ -69,13 +69,16 @@ mbuf= 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 \ diff --git a/scripts/snapdragon/adb/run-tool.sh b/scripts/snapdragon/adb/run-tool.sh index f6332391b..8bf7ba3a5 100755 --- a/scripts/snapdragon/adb/run-tool.sh +++ b/scripts/snapdragon/adb/run-tool.sh @@ -57,6 +57,9 @@ opfuse= mmsel= [ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM" +fasel= +[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA" + set -x tool=$1; shift @@ -65,5 +68,5 @@ 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 $opfuse $mmsel ./$branch/bin/$tool $@ \ + $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opfuse $mmsel $fasel ./$branch/bin/$tool $@ \ " diff --git a/scripts/snapdragon/ggml-hexagon-profile.py b/scripts/snapdragon/ggml-hexagon-profile.py index c53ad7779..0f9240ddc 100755 --- a/scripts/snapdragon/ggml-hexagon-profile.py +++ b/scripts/snapdragon/ggml-hexagon-profile.py @@ -230,6 +230,12 @@ def print_ascii_timeline(op_name, dims, types, usec, cycles, events, evt_val=Non 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':