From b540ec78377ca4666010fd322436a02ab01d4109 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Tue, 30 Jun 2026 15:00:33 -0700 Subject: [PATCH 01/32] hex-fa: refactor kernel param compute to use common layout builder --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 89 +++++++-------- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 125 +++++++++++++++++---- 2 files changed, 144 insertions(+), 70 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index a1d15fc7a55c..b873e606d345 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1582,62 +1582,55 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const uint32_t G = factx.G; // ======== VTCM allocation (GQA-aware) ======== + // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used + // throughout the KV loop below. 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) { + // Build the VTCM layout once (shared with the host estimator) and place every + // scratch buffer at its computed offset. This is the single source of truth for + // sizes/offsets/strides -- host and device can no longer disagree. + _Static_assert(sizeof(HVX_Vector) == HMX_FA_HVX_VECTOR_BYTES, "HVX vector size mismatch"); + struct hmx_fa_layout L; + hmx_fa_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline); + + if (L.total_bytes > ctx->vtcm_size) { return HTP_STATUS_VTCM_TOO_SMALL; } + uint8_t * const base = ctx->vtcm_base; + + factx.vtcm_q_tiles = (__fp16 *) (base + L.off_q_tiles); + factx.vtcm_o_tiles[0] = (__fp16 *) (base + L.off_o_tiles[0]); + factx.vtcm_o_tiles[1] = (__fp16 *) (base + L.off_o_tiles[1]); + factx.vtcm_k_fp16[0] = (__fp16 *) (base + L.off_k_fp16[0]); + factx.vtcm_k_fp16[1] = (__fp16 *) (base + L.off_k_fp16[1]); + factx.vtcm_v_fp16[0] = (__fp16 *) (base + L.off_v_fp16[0]); + factx.vtcm_v_fp16[1] = (__fp16 *) (base + L.off_v_fp16[1]); + factx.vtcm_k_tiles = (__fp16 *) (base + L.off_k_tiles); + factx.vtcm_v_tiles[0] = (__fp16 *) (base + L.off_v_tiles[0]); + factx.vtcm_v_tiles[1] = pipeline ? (__fp16 *) (base + L.off_v_tiles[1]) : NULL; + factx.vtcm_s_tiles = (__fp16 *) (base + L.off_s_tiles); + factx.vtcm_p_tiles = (__fp16 *) (base + L.off_p_tiles); + factx.vtcm_d_tiles = (__fp16 *) (base + L.off_d_tiles); + factx.vtcm_m_vec = (HVX_Vector *) (base + L.off_m_vec); + factx.vtcm_l_vec = (HVX_Vector *) (base + L.off_l_vec); + factx.vtcm_s_rowmax = (HVX_Vector *) (base + L.off_s_rowmax); + factx.vtcm_p_rowsum = (HVX_Vector *) (base + L.off_p_rowsum); + factx.vtcm_row_bufs = (HVX_Vector *) (base + L.off_row_bufs); + factx.row_buf_stride = L.row_buf_stride; + factx.vtcm_hmx_scales_id = base + L.off_hmx_scales_id; + factx.vtcm_hmx_scales_qk = base + L.off_hmx_scales_qk; + factx.vtcm_mask_buf = (__fp16 *) (base + L.off_mask_buf); + factx.mask_buf_row_stride = L.mask_buf_row_stride; + factx.vtcm_slopes = (__fp16 *) (base + L.off_slopes); + + const size_t m_line_bytes = L.m_line_bytes; // used by the mask DMAs in the KV loop + + dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE); + // ======== 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)); diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index fbae247e011c..a01731a3c061 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -81,39 +81,120 @@ struct htp_fa_kernel_params { 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) { +// VTCM region layout for the HMX flash-attention kernel. +// +// Single source of truth for both the host (which needs the total size to pick a +// (Br, Bc) tiling that fits the VTCM budget) and the device (which needs the actual +// byte offsets to place each scratch buffer). Building the layout once and reading +// offsets/total from it makes host estimate and device allocation impossible to +// desync -- previously they were duplicated formulas in two files and drifted. +// +// All fields are byte offsets / byte sizes -- no HVX_Vector type is named here so the +// header stays host-includable. The device casts (base + off_*) to the proper type. +// An offset of 0 marks a region that is not allocated for this configuration (only +// off_v_tiles[1], which exists only when pipelining); the device sets such pointers NULL. +struct hmx_fa_layout { + // Byte offsets from vtcm_base for each region. + size_t off_q_tiles; + size_t off_o_tiles[2]; + size_t off_k_fp16[2]; + size_t off_v_fp16[2]; + size_t off_k_tiles; + size_t off_v_tiles[2]; // [1] allocated only when pipeline, else 0 + size_t off_s_tiles; + size_t off_p_tiles; + size_t off_d_tiles; + size_t off_m_vec; + size_t off_l_vec; + size_t off_s_rowmax; + size_t off_p_rowsum; + size_t off_row_bufs; + size_t off_hmx_scales_id; + size_t off_hmx_scales_qk; + size_t off_mask_buf; + size_t off_slopes; + + // Region byte sizes reused by the device at runtime (not just for allocation). + size_t s_tile_bytes; // S and P tiles (same size) + size_t d_tile_bytes; + size_t m_line_bytes; // one mask row + size_t m_buf_slot_bytes; // one dma_cache slot = align_up(Br * m_line_bytes, 4096) + + // Derived strides. + size_t row_buf_stride; // HVX vectors (128B) per row buffer + size_t mask_buf_row_stride; // __fp16 elements per row in the mask buffer + + bool pipeline; + size_t total_bytes; +}; + +// HVX vector size in bytes. The device asserts sizeof(HVX_Vector) == this. +#define HMX_FA_HVX_VECTOR_BYTES 128 + +// Build the VTCM layout. Mirrors the device's vtcm_seq_alloc sequence exactly: +// a running offset bumped by each region size, with no inter-region alignment. +static inline void hmx_fa_layout_build(struct hmx_fa_layout * L, + 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_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 4096); // K DMA: [Bc, DK] x2 + const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 4096); // V DMA: [Bc, DV] x2 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 col_vec_size = hex_align_up(g_br * sizeof(float), 256); // m, l, s_rowmax, p_rowsum 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 m_buf_slot = hex_align_up(Br * m_line_size, 4096); + const size_t m_buf_size = m_buf_slot * 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) + size_t off = 0; + #define FA_ALLOC(field, sz) do { (L)->field = off; off += (sz); } while (0) + #define FA_ALLOC_OPTIONAL(field, sz, cond) do { if (cond) { FA_ALLOC(field, sz); } else { (L)->field = 0; } } while (0) + FA_ALLOC(off_q_tiles, q_tile_size); + FA_ALLOC(off_o_tiles[0], o_tile_size); + FA_ALLOC(off_o_tiles[1], o_tile_size); + FA_ALLOC(off_k_fp16[0], k_dma_size); + FA_ALLOC(off_k_fp16[1], k_dma_size); + FA_ALLOC(off_v_fp16[0], v_dma_size); + FA_ALLOC(off_v_fp16[1], v_dma_size); + FA_ALLOC(off_k_tiles, k_tile_size); + FA_ALLOC(off_v_tiles[0], v_tile_size); + FA_ALLOC_OPTIONAL(off_v_tiles[1], v_tile_size, pipeline); + FA_ALLOC(off_s_tiles, s_tile_size); + FA_ALLOC(off_p_tiles, s_tile_size); + FA_ALLOC(off_d_tiles, d_tile_size); + FA_ALLOC(off_m_vec, col_vec_size); + FA_ALLOC(off_l_vec, col_vec_size); + FA_ALLOC(off_s_rowmax, col_vec_size); + FA_ALLOC(off_p_rowsum, col_vec_size); + FA_ALLOC(off_row_bufs, row_vec_size * 2 * n_threads); + FA_ALLOC(off_hmx_scales_id, 256); + FA_ALLOC(off_hmx_scales_qk, 256); + FA_ALLOC(off_mask_buf, m_buf_size); + FA_ALLOC(off_slopes, slopes_size); + #undef FA_ALLOC_OPTIONAL + #undef FA_ALLOC + + L->s_tile_bytes = s_tile_size; + L->d_tile_bytes = d_tile_size; + L->m_line_bytes = m_line_size; + L->m_buf_slot_bytes = m_buf_slot; + L->row_buf_stride = row_vec_size / HMX_FA_HVX_VECTOR_BYTES; + L->mask_buf_row_stride = m_line_size / sizeof(__fp16); + L->pipeline = pipeline; + L->total_bytes = off; +} + +// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration. +static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) { + struct hmx_fa_layout L; + hmx_fa_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline); + return L.total_bytes; } #define FA_HVX_BLOCK_SIZE 64 From da33f881204c9a6810262b8841a8524147a48492 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Wed, 1 Jul 2026 12:26:19 -0700 Subject: [PATCH 02/32] hmx: add explicit compiler barriers to make hmx funcs more robust --- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 6 +++--- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 13 ++++++++++--- ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 15 +++++++++++++-- 3 files changed, 26 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 9f9471f69d9a..6ebb4b4888ef 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -20,6 +20,9 @@ add_library(${HTP_LIB} SHARED worker-pool.c hex-dma.c hmx-queue.c + matmul-ops.c + flash-attn-ops.c + gated-delta-net-ops.c binary-ops.c unary-ops.c sum-rows-ops.c @@ -37,10 +40,7 @@ add_library(${HTP_LIB} SHARED concat-ops.c diag-ops.c solve-tri-ops.c - gated-delta-net-ops.c pad-ops.c - matmul-ops.c - flash-attn-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index 1ab286e1c577..1ff95da7f782 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -47,19 +47,22 @@ static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) = }; // Inner HMX tile computation kernels -static inline void hmx_fa_qk_dot_tile( +static void hmx_fa_qk_dot_tile( const __fp16 * row_tiles, const __fp16 * col_tiles, __fp16 * out_tile, size_t n_dot_tiles ) { + asm volatile("" ::: "memory"); 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; } + asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(out_tile, 0); + asm volatile("" ::: "memory"); } static inline void hmx_fa_o_update_tile( @@ -70,17 +73,18 @@ static inline void hmx_fa_o_update_tile( __fp16 * o_tile_out, size_t n_col_tiles ) { + asm volatile("" ::: "memory"); 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; } - + asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(o_tile_out, 0); + asm volatile("" ::: "memory"); } static inline void hmx_fa_o_norm_tile( @@ -88,9 +92,12 @@ static inline void hmx_fa_o_norm_tile( const __fp16 * o_rc, __fp16 * o_out ) { + asm volatile("" ::: "memory"); Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047); Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047); + asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(o_out, 0); + asm volatile("" ::: "memory"); } #endif /* HMX_FA_KERNELS_H */ 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 4a0ca7885dcd..e2d1cb1a124d 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -614,26 +614,32 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict __builtin_assume(n_dot_tiles > 0); Q6_bias_mxmem2_A((void *)scales); + + asm volatile("" ::: "memory"); for (uint32_t r = 0; r < n_row_tiles; ++r) { for (size_t c = 0; c < n_col_tiles; ++c) { - Q6_mxclracc_hf(); - const __fp16 *row_tiles = activation + r * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; const __fp16 *col_tiles = weight + c * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + Q6_mxclracc_hf(); + asm volatile("" ::: "memory"); + for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { k_block = hex_smin(n_dot_tiles - k, 32); const uint32_t range = 2048u * (uint32_t)k_block - 1; Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); + asm volatile("" ::: "memory"); row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; } + asm volatile("" ::: "memory"); __fp16 *out_tile = output + (r * n_col_tiles + c) * HTP_MM_HMX_TILE_N_ELMS; Q6_mxmem_AR_after_hf(out_tile, 0); } } + asm volatile("" ::: "memory"); } // C += AB @@ -647,11 +653,14 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co Q6_bias_mxmem2_A((void *)col_scales); const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + + asm volatile("" ::: "memory"); for (size_t i = 0; i < n_row_tiles; ++i) { const __fp16 *row_base = a + i * dot_tile_stride; __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; for (size_t j = 0; j < n_col_tiles; ++j) { Q6_mxclracc_hf(); + asm volatile("" ::: "memory"); const __fp16 *col_tiles = b + j * dot_tile_stride; const __fp16 *row_tiles = row_base; @@ -670,9 +679,11 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; } + asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(accum_tile, 0); } } + asm volatile("" ::: "memory"); } // --- Async HMX matmul job (for pipeline overlap) --- From eaf9f0083ced29d18f21a8b6dd80c0ac30658191 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Wed, 1 Jul 2026 15:55:48 -0700 Subject: [PATCH 03/32] hex-vtcm: more generic vtcm layout builder for mm and flash-attn kernels --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 302 ++++++-------- ggml/src/ggml-hexagon/htp/concat-ops.c | 2 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 50 +-- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 57 ++- ggml/src/ggml-hexagon/htp/htp-vtcm.h | 19 + ggml/src/ggml-hexagon/htp/matmul-ops.c | 231 +++++------ ggml/src/ggml-hexagon/htp/matmul-ops.h | 436 +++++++++++++-------- ggml/src/ggml-hexagon/htp/vtcm-utils.h | 16 - 8 files changed, 571 insertions(+), 542 deletions(-) create mode 100644 ggml/src/ggml-hexagon/htp/htp-vtcm.h delete mode 100644 ggml/src/ggml-hexagon/htp/vtcm-utils.h diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 4f9dafed0f4c..502baf9397b6 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -2385,31 +2385,30 @@ 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, vtcm_dst_size = 0; + struct htp_mm_hvx_vtcm_layout L; 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_dst_size + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0->nb[1], 0, d, true, false, false ); - if (total_size <= vtcm_budget) { + if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; break; } } - 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_dst_size + if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0->nb[1], 0, 2, true, false, false ); } 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 = vtcm_dst_size; + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; } else { bool try_tiled = (k_align && opt_mm_select >= 2); if (try_tiled) { @@ -2420,37 +2419,36 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; } + struct htp_mm_hvx_vtcm_layout L; uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; uint32_t best_n_prefetch = 2; - size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0; - size_t total_size = 0; for (uint32_t d = max_prefetch; d >= 2; d /= 2) { - total_size = htp_mm_hvx_get_vtcm_sizes( - kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], d, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false ); - if (total_size <= vtcm_budget) { + if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; break; } } - if (best_n_prefetch == 2 && total_size > vtcm_budget) { - total_size = htp_mm_hvx_get_vtcm_sizes( - kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], 2, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false ); } kparams->n_prefetch = best_n_prefetch; - if (total_size <= vtcm_budget) { - kparams->vtcm_size = total_size; - kparams->vtcm_src0_size = vtcm_src0_size; - kparams->vtcm_src1_size = vtcm_src1_size; - kparams->vtcm_dst_size = vtcm_dst_size; + if (L.total_bytes <= vtcm_budget) { + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; goto done_quant; } - HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), total_size, vtcm_budget); + HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), L.total_bytes, vtcm_budget); } // Flat HVX fallback @@ -2458,17 +2456,17 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; - size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0; - size_t total_size = htp_mm_hvx_get_vtcm_sizes( - kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false ); kparams->n_prefetch = 16; - kparams->vtcm_size = total_size; - kparams->vtcm_src0_size = vtcm_src0_size; - kparams->vtcm_src1_size = vtcm_src1_size; - kparams->vtcm_dst_size = vtcm_dst_size; + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; } } @@ -2478,19 +2476,19 @@ static void ggml_hexagon_precompute_hvx_mm_params( const bool is_batched = (ne02 > 1) || (ne03 > 1); const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1); - size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0; - size_t vtcm_size = htp_mm_hvx_get_vtcm_sizes( - HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false ); - if (!is_batched && !is_permuted && vtcm_size <= vtcm_budget) { + if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_VTCM; kparams->src1_row_size = hex_round_up(ne10 * 2, 128); - kparams->vtcm_size = vtcm_size; - kparams->vtcm_src0_size = vtcm_src0_size; - kparams->vtcm_src1_size = vtcm_src1_size; - kparams->vtcm_dst_size = vtcm_dst_size; + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; kparams->n_prefetch = 16; } else { if (src1->type == GGML_TYPE_F32) { @@ -2499,14 +2497,14 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_DDR; } kparams->src1_row_size = src1->nb[1]; - size_t ddr_size = htp_mm_hvx_get_vtcm_sizes( - kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false ); - kparams->vtcm_size = ddr_size; - kparams->vtcm_src0_size = vtcm_src0_size; - kparams->vtcm_src1_size = vtcm_src1_size; - kparams->vtcm_dst_size = vtcm_dst_size; + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; kparams->n_prefetch = 16; } } else { @@ -2514,31 +2512,31 @@ static void ggml_hexagon_precompute_hvx_mm_params( const bool is_batched = (ne02 > 1) || (ne03 > 1); const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1); - size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0; - size_t vtcm_size = htp_mm_hvx_get_vtcm_sizes( - HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false ); - if (!is_batched && !is_permuted && vtcm_size <= vtcm_budget) { + if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_VTCM; kparams->src1_row_size = hex_round_up(ne10 * 4, 128); - kparams->vtcm_size = vtcm_size; - kparams->vtcm_src0_size = vtcm_src0_size; - kparams->vtcm_src1_size = vtcm_src1_size; - kparams->vtcm_dst_size = vtcm_dst_size; + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; kparams->n_prefetch = 16; } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_DDR; kparams->src1_row_size = src1->nb[1]; - size_t ddr_size = htp_mm_hvx_get_vtcm_sizes( - kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size + htp_mm_hvx_vtcm_layout_build( + &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false ); - kparams->vtcm_size = ddr_size; - kparams->vtcm_src0_size = vtcm_src0_size; - kparams->vtcm_src1_size = vtcm_src1_size; - kparams->vtcm_dst_size = vtcm_dst_size; + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; kparams->n_prefetch = 16; } } @@ -2608,80 +2606,57 @@ static void ggml_hexagon_precompute_fused_qkv_params( const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); const size_t src0_row_size = src0->nb[1]; - const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - size_t src0_sz_per_thread = 0; - size_t src2_sz_per_thread = 0; - 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_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128); - size_t src1_sz = src1_sz_per_thread; - const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; 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); - 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 + quant_scratch_size; - - if (tiled_vtcm_size <= sess->vtcm_size) { + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, src1_row_size, d, false, true, false + ); + if (L.total_bytes <= sess->vtcm_size) { best_n_prefetch = d; - src0_sz_per_thread = repacked_vtcm_size; - src2_sz_per_thread = hex_round_up(d * tile_row_size, 128); - src3_sz_per_thread = hex_round_up(d * tile_row_size, 128); break; } } - if (best_n_prefetch == 2 && src0_sz_per_thread == 0) { - size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128); - 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); - } - } else { - best_n_prefetch = 16; - src0_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128); - src2_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128); - src3_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128); } - size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128); - - size_t src0_sz = src0_sz_per_thread * sess->n_threads; - size_t src1_sz = src1_sz_per_thread; - 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 + quant_scratch_size; + struct htp_mm_hvx_vtcm_layout L; bool try_tiled = (opt_mm_select >= 2); - if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) { + + // Test tiled first + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false + ); + + if (try_tiled && L.total_bytes <= 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_src3_size = src3_sz; - kparams->vtcm_dst_size = quant_scratch_size; - kparams->vtcm_size = tiled_vtcm_size; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_src2_size = L.src2_bytes; + kparams->vtcm_src3_size = L.src3_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - size_t flat_src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128); - kparams->vtcm_src0_size = src0_sz; - kparams->vtcm_src1_size = flat_src1_sz; - kparams->vtcm_src2_size = src2_sz; - kparams->vtcm_src3_size = src3_sz; - kparams->vtcm_dst_size = quant_scratch_size; - kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + src3_sz + quant_scratch_size; + + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false + ); + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_src2_size = L.src2_bytes; + kparams->vtcm_src3_size = L.src3_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; } } @@ -2701,72 +2676,55 @@ static void ggml_hexagon_precompute_fused_ffn_params( const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); const size_t src0_row_size = src0->nb[1]; - const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - size_t src0_sz_per_thread = 0; - 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_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128); - size_t src1_sz = src1_sz_per_thread; - const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; 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); - 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 + quant_scratch_size; - - if (tiled_vtcm_size <= sess->vtcm_size) { + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, src1_row_size, d, false, false, true + ); + if (L.total_bytes <= sess->vtcm_size) { best_n_prefetch = d; - src0_sz_per_thread = repacked_vtcm_size; - src2_sz_per_thread = hex_round_up(d * tile_row_size, 128); break; } } - if (best_n_prefetch == 2 && src0_sz_per_thread == 0) { - size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128); - src0_sz_per_thread = repacked_vtcm_size; - src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128); - } - } else { - best_n_prefetch = 16; - src0_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128); - src2_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128); } - size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128); - - size_t src0_sz = src0_sz_per_thread * sess->n_threads; - 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 + quant_scratch_size; + struct htp_mm_hvx_vtcm_layout L; bool try_tiled = (opt_mm_select >= 2); - if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) { + + // Test tiled first + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true + ); + + if (try_tiled && L.total_bytes <= 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->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_src2_size = L.src2_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - size_t flat_src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128); - kparams->vtcm_src0_size = src0_sz; - kparams->vtcm_src1_size = flat_src1_sz; - kparams->vtcm_src2_size = src2_sz; - kparams->vtcm_dst_size = quant_scratch_size; - kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + quant_scratch_size; + + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true + ); + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_src2_size = L.src2_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; } } diff --git a/ggml/src/ggml-hexagon/htp/concat-ops.c b/ggml/src/ggml-hexagon/htp/concat-ops.c index f2a381313c5d..51d39e8d98f5 100644 --- a/ggml/src/ggml-hexagon/htp/concat-ops.c +++ b/ggml/src/ggml-hexagon/htp/concat-ops.c @@ -4,7 +4,7 @@ #include "hexagon_protos.h" #include "hvx_hexagon_protos.h" #include "hex-dma.h" -#include "vtcm-utils.h" +#include "htp-vtcm.h" #include "hvx-utils.h" #include "hex-fastdiv.h" #include diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index b873e606d345..ea77c6e32fec 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -22,7 +22,7 @@ #include "hvx-copy.h" #include "hvx-reduce.h" #include "hvx-flash-attn.h" -#include "vtcm-utils.h" +#include "htp-vtcm.h" #include "worker-pool.h" #define GGML_COMMON_DECL_C @@ -1593,8 +1593,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // scratch buffer at its computed offset. This is the single source of truth for // sizes/offsets/strides -- host and device can no longer disagree. _Static_assert(sizeof(HVX_Vector) == HMX_FA_HVX_VECTOR_BYTES, "HVX vector size mismatch"); - struct hmx_fa_layout L; - hmx_fa_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline); + struct hmx_fa_vtcm_layout L; + hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline); if (L.total_bytes > ctx->vtcm_size) { return HTP_STATUS_VTCM_TOO_SMALL; @@ -1602,30 +1602,30 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { uint8_t * const base = ctx->vtcm_base; - factx.vtcm_q_tiles = (__fp16 *) (base + L.off_q_tiles); - factx.vtcm_o_tiles[0] = (__fp16 *) (base + L.off_o_tiles[0]); - factx.vtcm_o_tiles[1] = (__fp16 *) (base + L.off_o_tiles[1]); - factx.vtcm_k_fp16[0] = (__fp16 *) (base + L.off_k_fp16[0]); - factx.vtcm_k_fp16[1] = (__fp16 *) (base + L.off_k_fp16[1]); - factx.vtcm_v_fp16[0] = (__fp16 *) (base + L.off_v_fp16[0]); - factx.vtcm_v_fp16[1] = (__fp16 *) (base + L.off_v_fp16[1]); - factx.vtcm_k_tiles = (__fp16 *) (base + L.off_k_tiles); - factx.vtcm_v_tiles[0] = (__fp16 *) (base + L.off_v_tiles[0]); - factx.vtcm_v_tiles[1] = pipeline ? (__fp16 *) (base + L.off_v_tiles[1]) : NULL; - factx.vtcm_s_tiles = (__fp16 *) (base + L.off_s_tiles); - factx.vtcm_p_tiles = (__fp16 *) (base + L.off_p_tiles); - factx.vtcm_d_tiles = (__fp16 *) (base + L.off_d_tiles); - factx.vtcm_m_vec = (HVX_Vector *) (base + L.off_m_vec); - factx.vtcm_l_vec = (HVX_Vector *) (base + L.off_l_vec); - factx.vtcm_s_rowmax = (HVX_Vector *) (base + L.off_s_rowmax); - factx.vtcm_p_rowsum = (HVX_Vector *) (base + L.off_p_rowsum); - factx.vtcm_row_bufs = (HVX_Vector *) (base + L.off_row_bufs); + factx.vtcm_q_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles); + factx.vtcm_o_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]); + factx.vtcm_o_tiles[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]); + factx.vtcm_k_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[0]); + factx.vtcm_k_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]); + factx.vtcm_v_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]); + factx.vtcm_v_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]); + factx.vtcm_k_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles); + factx.vtcm_v_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]); + factx.vtcm_v_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline); + factx.vtcm_s_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles); + factx.vtcm_p_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles); + factx.vtcm_d_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles); + factx.vtcm_m_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec); + factx.vtcm_l_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec); + factx.vtcm_s_rowmax = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax); + factx.vtcm_p_rowsum = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_p_rowsum); + factx.vtcm_row_bufs = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_row_bufs); factx.row_buf_stride = L.row_buf_stride; - factx.vtcm_hmx_scales_id = base + L.off_hmx_scales_id; - factx.vtcm_hmx_scales_qk = base + L.off_hmx_scales_qk; - factx.vtcm_mask_buf = (__fp16 *) (base + L.off_mask_buf); + factx.vtcm_hmx_scales_id = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_id); + factx.vtcm_hmx_scales_qk = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_qk); + factx.vtcm_mask_buf = VTCM_LAYOUT_PTR(__fp16, base, L.off_mask_buf); factx.mask_buf_row_stride = L.mask_buf_row_stride; - factx.vtcm_slopes = (__fp16 *) (base + L.off_slopes); + factx.vtcm_slopes = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes); const size_t m_line_bytes = L.m_line_bytes; // used by the mask DMAs in the KV loop diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index a01731a3c061..abf947f6ba47 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -7,6 +7,7 @@ #include "hex-fastdiv.h" #include "hex-common.h" +#include "htp-vtcm.h" #ifdef __cplusplus extern "C" { @@ -93,7 +94,7 @@ static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params // header stays host-includable. The device casts (base + off_*) to the proper type. // An offset of 0 marks a region that is not allocated for this configuration (only // off_v_tiles[1], which exists only when pipelining); the device sets such pointers NULL. -struct hmx_fa_layout { +struct hmx_fa_vtcm_layout { // Byte offsets from vtcm_base for each region. size_t off_q_tiles; size_t off_o_tiles[2]; @@ -133,7 +134,7 @@ struct hmx_fa_layout { // Build the VTCM layout. Mirrors the device's vtcm_seq_alloc sequence exactly: // a running offset bumped by each region size, with no inter-region alignment. -static inline void hmx_fa_layout_build(struct hmx_fa_layout * L, +static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) { const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS); @@ -153,32 +154,28 @@ static inline void hmx_fa_layout_build(struct hmx_fa_layout * L, const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128); size_t off = 0; - #define FA_ALLOC(field, sz) do { (L)->field = off; off += (sz); } while (0) - #define FA_ALLOC_OPTIONAL(field, sz, cond) do { if (cond) { FA_ALLOC(field, sz); } else { (L)->field = 0; } } while (0) - FA_ALLOC(off_q_tiles, q_tile_size); - FA_ALLOC(off_o_tiles[0], o_tile_size); - FA_ALLOC(off_o_tiles[1], o_tile_size); - FA_ALLOC(off_k_fp16[0], k_dma_size); - FA_ALLOC(off_k_fp16[1], k_dma_size); - FA_ALLOC(off_v_fp16[0], v_dma_size); - FA_ALLOC(off_v_fp16[1], v_dma_size); - FA_ALLOC(off_k_tiles, k_tile_size); - FA_ALLOC(off_v_tiles[0], v_tile_size); - FA_ALLOC_OPTIONAL(off_v_tiles[1], v_tile_size, pipeline); - FA_ALLOC(off_s_tiles, s_tile_size); - FA_ALLOC(off_p_tiles, s_tile_size); - FA_ALLOC(off_d_tiles, d_tile_size); - FA_ALLOC(off_m_vec, col_vec_size); - FA_ALLOC(off_l_vec, col_vec_size); - FA_ALLOC(off_s_rowmax, col_vec_size); - FA_ALLOC(off_p_rowsum, col_vec_size); - FA_ALLOC(off_row_bufs, row_vec_size * 2 * n_threads); - FA_ALLOC(off_hmx_scales_id, 256); - FA_ALLOC(off_hmx_scales_qk, 256); - FA_ALLOC(off_mask_buf, m_buf_size); - FA_ALLOC(off_slopes, slopes_size); - #undef FA_ALLOC_OPTIONAL - #undef FA_ALLOC + VTCM_LAYOUT_ALLOC(off_q_tiles, q_tile_size); + VTCM_LAYOUT_ALLOC(off_o_tiles[0], o_tile_size); + VTCM_LAYOUT_ALLOC(off_o_tiles[1], o_tile_size); + VTCM_LAYOUT_ALLOC(off_k_fp16[0], k_dma_size); + VTCM_LAYOUT_ALLOC(off_k_fp16[1], k_dma_size); + VTCM_LAYOUT_ALLOC(off_v_fp16[0], v_dma_size); + VTCM_LAYOUT_ALLOC(off_v_fp16[1], v_dma_size); + VTCM_LAYOUT_ALLOC(off_k_tiles, k_tile_size); + VTCM_LAYOUT_ALLOC(off_v_tiles[0], v_tile_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_v_tiles[1], v_tile_size, pipeline); + VTCM_LAYOUT_ALLOC(off_s_tiles, s_tile_size); + VTCM_LAYOUT_ALLOC(off_p_tiles, s_tile_size); + VTCM_LAYOUT_ALLOC(off_d_tiles, d_tile_size); + VTCM_LAYOUT_ALLOC(off_m_vec, col_vec_size); + VTCM_LAYOUT_ALLOC(off_l_vec, col_vec_size); + VTCM_LAYOUT_ALLOC(off_s_rowmax, col_vec_size); + VTCM_LAYOUT_ALLOC(off_p_rowsum, col_vec_size); + VTCM_LAYOUT_ALLOC(off_row_bufs, row_vec_size * 2 * n_threads); + VTCM_LAYOUT_ALLOC(off_hmx_scales_id, 256); + VTCM_LAYOUT_ALLOC(off_hmx_scales_qk, 256); + VTCM_LAYOUT_ALLOC(off_mask_buf, m_buf_size); + VTCM_LAYOUT_ALLOC(off_slopes, slopes_size); L->s_tile_bytes = s_tile_size; L->d_tile_bytes = d_tile_size; @@ -192,8 +189,8 @@ static inline void hmx_fa_layout_build(struct hmx_fa_layout * L, // Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration. static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) { - struct hmx_fa_layout L; - hmx_fa_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline); + struct hmx_fa_vtcm_layout L; + hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline); return L.total_bytes; } diff --git a/ggml/src/ggml-hexagon/htp/htp-vtcm.h b/ggml/src/ggml-hexagon/htp/htp-vtcm.h new file mode 100644 index 000000000000..f4b2239a2ebb --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/htp-vtcm.h @@ -0,0 +1,19 @@ +#ifndef HTP_VTCM_H +#define HTP_VTCM_H + +#include +#include + +static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) { + uint8_t *p = *vtcm_ptr; + *vtcm_ptr += size; + return p; +} + +#define VTCM_LAYOUT_ALLOC(field, sz) do { (L)->field = off; off += (sz); } while (0) +#define VTCM_LAYOUT_ALLOC_OPTIONAL(field, sz, cond) do { if (cond) { VTCM_LAYOUT_ALLOC(field, sz); } else { (L)->field = 0; } } while (0) + +#define VTCM_LAYOUT_PTR(type, base, offset) ((type *)((uint8_t *)(base) + (offset))) +#define VTCM_LAYOUT_PTR_OPTIONAL(type, base, offset, cond) ((cond) ? VTCM_LAYOUT_PTR(type, base, offset) : NULL) + +#endif // HTP_VTCM_H diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 79c292dc242f..cd8c5308dd1a 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -20,7 +20,7 @@ #include "htp-ctx.h" #include "htp-ops.h" #include "matmul-ops.h" -#include "vtcm-utils.h" +#include "htp-vtcm.h" static void hvx_tensor_add_f32_grid( const struct htp_tensor * restrict dst, @@ -1514,37 +1514,26 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { break; } - size_t src0_sz = 0, src1_sz = 0, dst_sz = 0; - if (kparams->vtcm_src0_size > 0 || kparams->vtcm_src1_size > 0 || kparams->vtcm_dst_size > 0) { - src0_sz = kparams->vtcm_src0_size; - src1_sz = kparams->vtcm_src1_size; - dst_sz = kparams->vtcm_dst_size; - } else { - const uint32_t n_prefetch = kparams->n_prefetch; - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); - htp_mm_hvx_get_vtcm_sizes( - kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, - dst_row_size, src0_row_size, src1_row_size, n_prefetch, - &src0_sz, &src1_sz, &dst_sz - ); - } + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, + dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false); if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { - mmctx->vtcm_src1_size_per_thread = src1_sz; + mmctx->vtcm_src1_size_per_thread = L.src1_bytes; } else { - mmctx->vtcm_src1_size_per_thread = src1_sz / octx->n_threads; + mmctx->vtcm_src1_size_per_thread = L.src1_bytes / octx->n_threads; } - mmctx->vtcm_src0_size_per_thread = src0_sz / octx->n_threads; - mmctx->vtcm_dst_size_per_thread = dst_sz / octx->n_threads; + mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; + mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; - size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : (src1_sz + src0_sz + dst_sz); + size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; FARF(HIGH, "matmul-%s : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu (%zu)\n", mmctx->type, - src0_sz, src1_sz, dst_sz, vtcm_size); + L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size); FARF(HIGH, "matmul-%s : %ux%ux%ux%u * %ux%ux%ux%u-> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], @@ -1556,10 +1545,10 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { return HTP_STATUS_VTCM_TOO_SMALL; } - uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz); - mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz); - mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz); + uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; @@ -2198,45 +2187,30 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; - const size_t act_f32_size = hex_align_up((size_t)act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE); + struct htp_mm_hmx_vtcm_layout L; + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size); - const size_t weight_area_size = is_quant - ? hex_align_up((n_chunk_n_cols / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE) - : hex_align_up(n_chunk_n_cols * row_stride, HTP_MM_HMX_TILE_SIZE); - const size_t act_area_size = hex_align_up(m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t output_area_size = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE); - - size_t scratch0_size, scratch1_size, scratch2_size; - scratch0_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE); // dequant buf 0 - scratch1_size = pipeline ? scratch0_size : 0; // dequant buf 1 - scratch2_size = pipeline ? output_area_size : 0; // output buf 1 - - uint8_t *vtcm_ptr = (uint8_t *) ctx->vtcm_base; - __fp16 *vtcm_weight_raw[2] = { NULL, NULL }; - if (weight_area_size) { - if (pipeline) { - vtcm_weight_raw[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size); - vtcm_weight_raw[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size); - } else { - 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_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; - void *vtcm_scratch2 = scratch2_size ? vtcm_seq_alloc(&vtcm_ptr, scratch2_size) : NULL; - __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256); - - vtcm_used = vtcm_ptr - (uint8_t *) ctx->vtcm_base; + vtcm_used = L.total_bytes; if (vtcm_used > vtcm_budget) { FARF(ERROR, "hmx-mm-2d-precomputed: VTCM overflow: used %zu budget %zu, m %d k %d n %d mc %zu nc %zu", vtcm_used, vtcm_budget, m, k, n, m_chunk_n_rows, n_chunk_n_cols); return -1; } + uint8_t * const base = (uint8_t *) ctx->vtcm_base; + __fp16 *vtcm_weight_raw[2] = { + VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]), + VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_weight[1], pipeline) + }; + + __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act); + float *vtcm_f32_act = VTCM_LAYOUT_PTR(float, base, L.off_act_f32); + __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]); + void *vtcm_scratch0 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]); + void *vtcm_scratch1 = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_scratch[1], pipeline); + void *vtcm_scratch2 = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_dst[1], pipeline); + __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales); + hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 FARF(HIGH, "hmx-mm-2d-precomputed: standard : m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu", @@ -2458,32 +2432,29 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ size_t n_chunk_n_cols = n_chunk; size_t vtcm_used = vtcm_size; - const size_t act_head_stride = m_chunk_n_rows * (size_t) params->k; // fp16 elements between heads - const size_t weight_area_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t activation_area_size = hex_align_up(group_size * m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t output_area_size = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE); - const size_t scratch_area_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 = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size); - __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, activation_area_size); - __fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size); - void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch_area_size); - void *vtcm_scratch1 = vtcm_seq_alloc(&vtcm_ptr, scratch_area_size); - __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256); - float *vtcm_f32_act = use_dma_activation ? (float *) vtcm_seq_alloc(&vtcm_ptr, f32_scratch_size) : NULL; - - if ((size_t) (vtcm_ptr - (uint8_t *) ctx->vtcm_base) > vtcm_budget) { + struct htp_mm_hmx_vtcm_layout L; + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0); + + if (L.total_bytes > vtcm_budget) { FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__); return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size); } + uint8_t * const base = (uint8_t *) ctx->vtcm_base; + __fp16 *vtcm_weight = VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]); + __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act); + __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]); + void *vtcm_scratch0 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]); + void *vtcm_scratch1 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[1]); + __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales); + float *vtcm_f32_act = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_act_f32, use_dma_activation); + hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 FARF(HIGH, "%s: grouped path m=%d k=%d n=%d group=%d streams=%d mc=%zu nc=%zu vtcm=%zu/%zu", __func__, params->m, params->k, params->n, group_size, params->ne13, m_chunk_n_rows, n_chunk_n_cols, - (size_t) (vtcm_ptr - (uint8_t *) ctx->vtcm_base), vtcm_budget); + L.total_bytes, vtcm_budget); const size_t fp16_row_bytes = (size_t) params->k * sizeof(__fp16); const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16); @@ -2505,7 +2476,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ // thrashing from HVX loads at large strides. 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; + __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; transfer_activation_chunk_threaded(ctx, vtcm_act_g, activation_chunk, (int) n_rows, params->k, params->act_stride, act_threads, params->k, vtcm_f32_act); @@ -2545,7 +2516,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS]; htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, g); { - const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * act_head_stride; + const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; core_dot_chunk_fp16(vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32); } @@ -2960,22 +2931,14 @@ static int hvx_mm_matmul_id( } 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); - // Scratchpad sizes are computed on the host (htp_mm_hvx_id_get_vtcm_sizes) and passed in. - // The ID layout is routing-independent, so the host has exact visibility -- consume it here - // rather than recomputing, to keep host budgeting and device allocation in lockstep. - 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 = kparams->vtcm_dst_size; - size_t vtcm_size = kparams->vtcm_size; + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, + 0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false); - 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 = dst_sz / octx->n_threads; + size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; - 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); + FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size 0 dst-spad-size %zu (%zu)\n", mmctx->type, + L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size); FARF(HIGH, "matmul-id-%s : %ux%ux%ux%u * %ux%ux%ux%u (%ux%ux%ux%u) -> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], @@ -2989,11 +2952,11 @@ static int hvx_mm_matmul_id( return HTP_STATUS_VTCM_TOO_SMALL; } - uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base; - 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); + uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src2 = NULL; + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; @@ -3003,10 +2966,10 @@ static int hvx_mm_matmul_id( mmctx->vtcm_src0_stride = src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; - 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; + mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; + mmctx->vtcm_src1_size_per_thread = L.src1_bytes; + mmctx->vtcm_src2_size_per_thread = 0; + mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; @@ -3181,19 +3144,11 @@ int op_matmul_qkv(struct htp_ops_context * octx) { src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]); } - // Set up scratchpads using precomputed sizes from the host - 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 src3_sz = kparams->vtcm_src3_size; - size_t dst_sz = kparams->vtcm_dst_size; - size_t vtcm_size = kparams->vtcm_size; + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, + 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false); - 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; + size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; if (octx->ctx->vtcm_size < vtcm_size) { FARF(ERROR, "matmul-qkv: current VTCM reservation %zu is too small, needed %zu\n", @@ -3201,12 +3156,12 @@ int op_matmul_qkv(struct htp_ops_context * octx) { return HTP_STATUS_VTCM_TOO_SMALL; } - uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base; - 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_src3 = vtcm_seq_alloc(&vtcm_ptr, src3_sz); - mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz); + uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); + mmctx->vtcm_src3 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src3); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; @@ -3219,11 +3174,11 @@ int op_matmul_qkv(struct htp_ops_context * octx) { mmctx->vtcm_src3_stride = is_repacked ? 0 : src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; - 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_src3_size_per_thread = src3_sz_per_thread; - mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread; + mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; + mmctx->vtcm_src1_size_per_thread = L.src1_bytes; + mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads; + mmctx->vtcm_src3_size_per_thread = L.src3_bytes / octx->n_threads; + mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) return HTP_STATUS_OK; @@ -3331,28 +3286,22 @@ int op_matmul_ffn(struct htp_ops_context * octx) { src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]); } - // Set up scratchpads using precomputed sizes from the host - 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; + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, + 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true); - 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; + size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; 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); return HTP_STATUS_VTCM_TOO_SMALL; } - uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base; - 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); + uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; @@ -3363,10 +3312,10 @@ int op_matmul_ffn(struct htp_ops_context * octx) { mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; - 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; + mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; + mmctx->vtcm_src1_size_per_thread = L.src1_bytes; + mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads; + mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) return HTP_STATUS_OK; diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index 96369825a93e..c3752d4dba62 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -6,6 +6,7 @@ #include "htp-ops.h" #include "hex-fastdiv.h" #include "hex-common.h" +#include "htp-vtcm.h" #ifdef __cplusplus extern "C" { @@ -295,197 +296,318 @@ static inline void htp_mm_hmx_get_batched_chunk_costs( *size_per_mn_out = sizeof(uint16_t); } -static inline size_t htp_mm_hmx_get_2d_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size -) { - const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; - const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); - const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k); - const size_t vec_dot_size = k * sizeof(uint16_t); - - const size_t act_f32_size = htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE); - size_t weight_area_size = is_quant - ? htp_mm_round_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE) - : htp_mm_round_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE); - if (pipeline) { - weight_area_size *= 2; - } - const size_t act_area_size = htp_mm_round_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t output_area_size = htp_mm_round_up(mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); +struct htp_mm_hmx_vtcm_layout { + // Byte offsets from vtcm_base for each region + size_t off_weight[2]; // [1] is only used when pipelined + size_t off_act; + size_t off_act_f32; // fp32 activation conversion scratch + size_t off_dst[2]; // [1] is only used when pipelined + size_t off_scratch[2]; // dequantization scratch pads + size_t off_scales; // HMX scales (256 bytes) + + // Cached sizes of regions for HMX kernel use + size_t weight_area_bytes; + size_t act_area_bytes; + size_t act_f32_bytes; + size_t output_area_bytes; + size_t scratch_bytes[2]; + size_t act_head_stride; + + size_t total_bytes; +}; - size_t scratch0_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - size_t scratch1_size = pipeline ? scratch0_size : 0; - size_t scratch2_size = pipeline ? output_area_size : 0; +struct htp_mm_hvx_vtcm_layout { + // Byte offsets from vtcm_base for each region + size_t off_src1; // vtcm_src1 (activation) + size_t off_src0; // vtcm_src0 (weight/Wk) + size_t off_src2; // vtcm_src2 (Wq / fused only) + size_t off_src3; // vtcm_src3 (Wv / fused only) + size_t off_dst; // vtcm_dst (output scratch) + + // Cached sizes + size_t src0_bytes; + size_t src1_bytes; + size_t src2_bytes; + size_t src3_bytes; + size_t dst_bytes; + + size_t total_bytes; +}; - return weight_area_size + act_area_size + act_f32_size + output_area_size + - scratch0_size + scratch1_size + scratch2_size + 256; -} +static inline void htp_mm_hmx_vtcm_layout_build( + struct htp_mm_hmx_vtcm_layout * L, + int kernel_type, + int wtype, + uint32_t k, + size_t mc, + size_t nc, + uint32_t group_size, + bool use_dma_activation, + bool pipeline, + uint32_t act_threads, + uint32_t aligned_tile_size +) { + size_t off = 0; + + if (kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) { + const size_t vec_dot_size = k * sizeof(uint16_t); + const size_t act_head_stride = mc * k; + const size_t weight_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); + const size_t activation_area_size = hex_align_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); + const size_t output_area_size = hex_align_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); + const size_t scratch_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); + const size_t f32_scratch_size = use_dma_activation + ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0; + + VTCM_LAYOUT_ALLOC(off_weight[0], weight_area_size); + L->off_weight[1] = 0; + VTCM_LAYOUT_ALLOC(off_act, activation_area_size); + VTCM_LAYOUT_ALLOC(off_dst[0], output_area_size); + L->off_dst[1] = 0; + VTCM_LAYOUT_ALLOC(off_scratch[0], scratch_area_size); + VTCM_LAYOUT_ALLOC(off_scratch[1], scratch_area_size); + VTCM_LAYOUT_ALLOC(off_scales, 256); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_act_f32, f32_scratch_size, use_dma_activation); + + L->weight_area_bytes = weight_area_size; + L->act_area_bytes = activation_area_size; + L->act_f32_bytes = f32_scratch_size; + L->output_area_bytes = output_area_size; + L->scratch_bytes[0] = scratch_area_size; + L->scratch_bytes[1] = scratch_area_size; + L->act_head_stride = act_head_stride; + } else { + // HTP_MM_KERNEL_HMX_2D + const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); + const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k); + const size_t vec_dot_size = k * sizeof(uint16_t); + const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; + + const size_t act_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE); + const size_t weight_area_size = is_quant + ? hex_align_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE) + : hex_align_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE); + const size_t act_area_size = hex_align_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); + const size_t output_area_size = hex_align_up(mc * nc * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE); + + const size_t scratch0_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); + const size_t scratch1_size = pipeline ? scratch0_size : 0; + + VTCM_LAYOUT_ALLOC(off_weight[0], weight_area_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_weight[1], weight_area_size, pipeline); + VTCM_LAYOUT_ALLOC(off_act, act_area_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_act_f32, act_f32_size, true); + VTCM_LAYOUT_ALLOC(off_dst[0], output_area_size); + VTCM_LAYOUT_ALLOC(off_scratch[0], scratch0_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_scratch[1], scratch0_size, pipeline); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_dst[1], output_area_size, pipeline); + VTCM_LAYOUT_ALLOC(off_scales, 256); + + L->weight_area_bytes = weight_area_size; + L->act_area_bytes = act_area_size; + L->act_f32_bytes = act_f32_size; + L->output_area_bytes = output_area_size; + L->scratch_bytes[0] = scratch0_size; + L->scratch_bytes[1] = scratch1_size; + L->act_head_stride = 0; + } -static inline size_t htp_mm_hmx_get_batched_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) { - (void)wtype; - (void)pipeline; - const size_t vec_dot_size = k * sizeof(uint16_t); - const size_t f32_scratch_size = use_dma_activation - ? htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0; - - const size_t act_head_stride = mc * k; - const size_t weight_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t act_area_size = htp_mm_round_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); - const size_t output_area_size = htp_mm_round_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); - const size_t scratch_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - - return weight_area_size + act_area_size + output_area_size + - 2 * scratch_area_size + 256 + f32_scratch_size; + L->total_bytes = off; } -static inline size_t htp_mm_hvx_get_vtcm_sizes( +static inline void htp_mm_hvx_vtcm_layout_build( + struct htp_mm_hvx_vtcm_layout * L, int kernel_type, int wtype, uint32_t ne10, // k - uint32_t src1_nrows, // m_total (or act_nrows) + uint32_t src1_nrows, // m_total uint32_t n_threads, size_t dst_row_size, size_t src0_row_size, size_t src1_row_size, uint32_t n_prefetch, - size_t * vtcm_src0_size_out, - size_t * vtcm_src1_size_out, - size_t * vtcm_dst_size_out + bool is_matmul_id, + bool is_fused_qkv, + bool is_fused_ffn ) { - size_t vtcm_src0_size = 0; - size_t vtcm_src1_size = 0; - size_t vtcm_dst_size = 0; + size_t src0_sz = 0; + size_t src1_sz = 0; + size_t src2_sz = 0; + size_t src3_sz = 0; + size_t dst_sz = 0; const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || wtype == HTP_TYPE_MXFP4); - const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); - const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1; - - switch (kernel_type) { - case HTP_MM_KERNEL_HVX_F16_F16_VTCM: { - size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128); - vtcm_src1_size = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256); - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; - vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } - case HTP_MM_KERNEL_HVX_F16_F32_DDR: - case HTP_MM_KERNEL_HVX_F16_F16_DDR: - case HTP_MM_KERNEL_HVX_F32_F32_DDR: - case HTP_MM_KERNEL_HVX_F32_F16_DDR: { - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads; - vtcm_src1_size = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads; - vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } - case HTP_MM_KERNEL_HVX_F32_F32_VTCM: { - size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128); - vtcm_src1_size = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256); - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; - vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } - case HTP_MM_KERNEL_HVX_QUANT_BLOCK: - 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_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); - - vtcm_src0_size = vtcm_src0_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); - vtcm_src0_size = repacked_vtcm_size * n_threads; + if (is_fused_qkv || is_fused_ffn) { + const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); + const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; + + size_t src0_sz_per_thread = 0; + size_t src2_sz_per_thread = 0; + size_t src3_sz_per_thread = 0; + + 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; + + src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); + src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); + if (is_fused_qkv) { + src3_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); } - - 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; + } else { + src0_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); + src2_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); + if (is_fused_qkv) { + src3_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); } - 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_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 flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); + size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + + if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { + src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128); + } else { + src1_sz = hex_round_up(tiled_src1_row_size * src1_nrows, 128); + } - vtcm_src0_size = vtcm_src0_size * n_threads; + src0_sz = src0_sz_per_thread * n_threads; + src2_sz = src2_sz_per_thread * n_threads; + src3_sz = src3_sz_per_thread * n_threads; + dst_sz = quant_scratch_size; + } else if (is_matmul_id) { + const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); + const size_t src1_row_size_tiled = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) + : htp_mm_q8_0_tiled_row_size(ne10); + + size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); + src1_sz = htp_mm_round_up(src1_row_size_tiled * src1_nrows, 256); + + 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); + src0_sz_per_thread = repacked_vtcm_size; + } - 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); - vtcm_src0_size = repacked_vtcm_size * n_threads; + src0_sz = src0_sz_per_thread * n_threads; + dst_sz = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; + } else { + const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); + const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1; + + switch (kernel_type) { + case HTP_MM_KERNEL_HVX_F16_F16_VTCM: { + size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128); + src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256); + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; + dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + break; } - - 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; + case HTP_MM_KERNEL_HVX_F16_F32_DDR: + case HTP_MM_KERNEL_HVX_F16_F16_DDR: + case HTP_MM_KERNEL_HVX_F32_F32_DDR: + case HTP_MM_KERNEL_HVX_F32_F16_DDR: { + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads; + src1_sz = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads; + dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + break; + } + case HTP_MM_KERNEL_HVX_F32_F32_VTCM: { + size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128); + src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256); + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; + dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + break; } - vtcm_dst_size = dst_size_per_thread * n_threads; - break; + case HTP_MM_KERNEL_HVX_QUANT_BLOCK: + 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); + + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); + src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256); + + src0_sz = src0_sz * 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); + src0_sz = 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; + } + dst_sz = 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); + + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); + src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256); + + src0_sz = src0_sz * 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); + src0_sz = 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; + } + dst_sz = dst_size_per_thread * n_threads; + break; + } + default: + break; } - default: - break; } - *vtcm_src0_size_out = vtcm_src0_size; - *vtcm_src1_size_out = vtcm_src1_size; - *vtcm_dst_size_out = vtcm_dst_size; - - return vtcm_src0_size + vtcm_src1_size + vtcm_dst_size; + size_t off = 0; + VTCM_LAYOUT_ALLOC(off_src1, src1_sz); + VTCM_LAYOUT_ALLOC(off_src0, src0_sz); + VTCM_LAYOUT_ALLOC(off_src2, src2_sz); + VTCM_LAYOUT_ALLOC(off_src3, src3_sz); + VTCM_LAYOUT_ALLOC(off_dst, dst_sz); + + L->src0_bytes = src0_sz; + L->src1_bytes = src1_sz; + L->src2_bytes = src2_sz; + L->src3_bytes = src3_sz; + L->dst_bytes = dst_sz; + L->total_bytes = off; } -static inline size_t htp_mm_hvx_id_get_vtcm_sizes( - int wtype, - uint32_t ne10, // k - uint32_t src1_nrows, - uint32_t n_threads, - 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_dst_size_out +static inline size_t htp_mm_hmx_get_2d_vtcm_size( + int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size ) { - const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 || - wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || - wtype == HTP_TYPE_MXFP4); - - const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); - const size_t src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) - : htp_mm_q8_0_tiled_row_size(ne10); - - 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); - - 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); - 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; + struct htp_mm_hmx_vtcm_layout L; + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size); + return L.total_bytes; +} - return vtcm_src0_size + src1_sz + vtcm_dst_size; +static inline size_t htp_mm_hmx_get_batched_vtcm_size( + int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) { + (void)pipeline; + struct htp_mm_hmx_vtcm_layout L; + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0); + return L.total_bytes; } #ifdef __cplusplus diff --git a/ggml/src/ggml-hexagon/htp/vtcm-utils.h b/ggml/src/ggml-hexagon/htp/vtcm-utils.h deleted file mode 100644 index b129fb74e315..000000000000 --- a/ggml/src/ggml-hexagon/htp/vtcm-utils.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef VTCM_UTILS_H -#define VTCM_UTILS_H - -#include "hex-utils.h" - -#include -#include -#include - -static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) { - uint8_t *p = *vtcm_ptr; - *vtcm_ptr += size; - return p; -} - -#endif // VTCM_UTILS_H From 8fc4a9434ed64b008fac222cb4a4fb358d26e4ec Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Wed, 1 Jul 2026 17:40:35 -0700 Subject: [PATCH 04/32] hex-hmx: unroll inner kernels --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 140 ++++++++++-------- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 78 +++++++++- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 58 ++++++-- 3 files changed, 194 insertions(+), 82 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index ea77c6e32fec..8c9f4889cd1c 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -702,9 +702,9 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { 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; + __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; *pv_out0 = Q6_V_lo_W(vp); *pv_out1 = Q6_V_hi_W(vp); @@ -798,11 +798,10 @@ 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 n_rows_g = args->n_rows_g; const size_t rows_per_t = args->rows_per_t; + const size_t G = factx->G; + const size_t DV = factx->DV; const size_t start = (size_t) i * rows_per_t; const size_t end = hex_smin(start + rows_per_t, n_rows_g); @@ -831,10 +830,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) { 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); + const __fp16 * in_dtile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2; + const HVX_Vector * pv_in0 = ((const HVX_Vector *) in_dtile) + 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 { @@ -957,14 +956,14 @@ static inline void fa_softmax_impl( 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; + size_t ci = c / 64; + const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2; + const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2; + 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_VectorPair vp_s_drow = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2); + HVX_Vector v_s_row0 = Q6_V_lo_W(vp_s_drow); + HVX_Vector v_s_row1 = Q6_V_hi_W(vp_s_drow); HVX_Vector t0 = hvx_vec_tanh_f16(v_s_row0); my_row_buf0[ci] = hvx_vec_mul_f16_f16(t0, v_cap); @@ -974,14 +973,14 @@ static inline void fa_softmax_impl( } } 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); + size_t ci = c / 64; + const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2; + const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2; + const HVX_Vector * pv_s_in1 = pv_s_in0 + 16; + + HVX_VectorPair vp_s_drow = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2); + my_row_buf0[ci] = Q6_V_lo_W(vp_s_drow); + my_row_buf1[ci] = Q6_V_hi_W(vp_s_drow); } } @@ -1118,9 +1117,9 @@ static inline void fa_softmax_impl( 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; + __fp16 * out_dtile = p_st_base + ci * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_p_out0 = ((HVX_Vector *) out_dtile) + 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); @@ -1150,7 +1149,7 @@ static inline void fa_softmax_impl( } // 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_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); @@ -1160,7 +1159,7 @@ static inline void fa_softmax_impl( 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 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); @@ -1332,13 +1331,16 @@ static void hmx_fa_qk_dot_worker(void * data) { __builtin_assume(n_dot_tiles > 0); Q6_bias_mxmem2_A((void *) job->hmx_scales); + const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; 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; + const __fp16 * row_tiles = q_tiles + r * dot_stride; + const __fp16 * col_tiles = k_tiles; + __fp16 * out_tile = s_tiles + r * n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; + for (size_t c = 0; c < n_col_tiles; ++c) { hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles); + col_tiles += dot_stride; + out_tile += HMX_FP16_TILE_N_ELMS; } } } @@ -1374,16 +1376,20 @@ static void hmx_fa_o_update_worker(void * data) { __builtin_assume(DV_tiles > 0); Q6_bias_mxmem2_A((void *) job->hmx_scales); + const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; + const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; 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; + const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; + const __fp16 * p_tile_in = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS; + const __fp16 * o_rc = o_prev + r * HMX_FP16_TILE_N_ELMS; + const __fp16 * v_tile_in = v_tiles; + __fp16 * o_tile_out = o_curr + r * HMX_FP16_TILE_N_ELMS; + for (size_t c = 0; c < DV_tiles; ++c) { hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles); + o_rc += o_stride; + v_tile_in += v_stride; + o_tile_out += o_stride; } } } @@ -1410,13 +1416,16 @@ static void hmx_fa_o_norm_worker(void * data) { __builtin_assume(DV_tiles > 0); Q6_bias_mxmem2_A((void *) job->hmx_scales); + const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; 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; + const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; + const __fp16 * o_rc = o_prev + r * HMX_FP16_TILE_N_ELMS; + __fp16 * o_out = o_curr + r * DV_tiles * HMX_FP16_TILE_N_ELMS; + for (size_t c = 0; c < DV_tiles; ++c) { hmx_fa_o_norm_tile(d_diag, o_rc, o_out); + o_rc += o_stride; + o_out += HMX_FP16_TILE_N_ELMS; } } } @@ -1871,13 +1880,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); + const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; 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; + const __fp16 * row_tiles = q_base + r * dot_stride; + const __fp16 * col_tiles = k_base; + __fp16 * out_tile = s_base + r * n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; + for (size_t c = 0; c < n_col_tiles; ++c) { hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles); + col_tiles += dot_stride; + out_tile += HMX_FP16_TILE_N_ELMS; } } htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); @@ -1930,15 +1942,20 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); + const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; + const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; 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; + const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; + const __fp16 * p_tile_in = p_base + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS; + const __fp16 * o_rc = op_base + r * HMX_FP16_TILE_N_ELMS; + const __fp16 * v_tile_in = v_base; + __fp16 * o_tile_out = oc_base + r * HMX_FP16_TILE_N_ELMS; + for (size_t c = 0; c < DV_tiles; ++c) { hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles); + o_rc += o_stride; + v_tile_in += v_stride; + o_tile_out += o_stride; } } htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); @@ -1975,13 +1992,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); + const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; 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; + const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; + const __fp16 * o_rc = op_base + r * HMX_FP16_TILE_N_ELMS; + __fp16 * o_out = oc_base + r * DV_tiles * HMX_FP16_TILE_N_ELMS; + for (size_t c = 0; c < DV_tiles; ++c) { hmx_fa_o_norm_tile(d_diag, o_rc, o_out); + o_rc += o_stride; + o_out += HMX_FP16_TILE_N_ELMS; } } htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index 1ff95da7f782..439b594700ab 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -54,11 +54,44 @@ static void hmx_fa_qk_dot_tile( size_t n_dot_tiles ) { asm volatile("" ::: "memory"); - for (size_t k = 0; k < n_dot_tiles; ++k) { + if (n_dot_tiles == 2) { 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_activation_hf_mxmem_RR((unsigned int) (row_tiles + HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + HMX_FP16_TILE_N_ELMS), 2047); + } else if (n_dot_tiles == 4) { + Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047); + Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); + } else if (n_dot_tiles == 8) { + Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047); + Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 4 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 4 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 5 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 5 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 6 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 6 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 7 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 7 * HMX_FP16_TILE_N_ELMS), 2047); + } else { + 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; + } } asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(out_tile, 0); @@ -76,11 +109,44 @@ static inline void hmx_fa_o_update_tile( asm volatile("" ::: "memory"); 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) { + if (n_col_tiles == 2) { + Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047); + Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + } else if (n_col_tiles == 4) { + Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047); + Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); + } else if (n_col_tiles == 8) { 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_activation_hf_mxmem_RR((unsigned int) (p_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 4 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 4 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 5 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 5 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 6 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 6 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 7 * HMX_FP16_TILE_N_ELMS), 2047); + Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 7 * HMX_FP16_TILE_N_ELMS), 2047); + } else { + 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; + } } asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(o_tile_out, 0); 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 e2d1cb1a124d..21682a1b3d19 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -615,28 +615,42 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict Q6_bias_mxmem2_A((void *)scales); + const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + asm volatile("" ::: "memory"); for (uint32_t r = 0; r < n_row_tiles; ++r) { + const __fp16 *row_base = activation + r * dot_stride; + const __fp16 *col_base = weight; + __fp16 *out_tile = output + r * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; + for (size_t c = 0; c < n_col_tiles; ++c) { - const __fp16 *row_tiles = activation + r * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; - const __fp16 *col_tiles = weight + c * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + const __fp16 *row_tiles = row_base; + const __fp16 *col_tiles = col_base; Q6_mxclracc_hf(); asm volatile("" ::: "memory"); - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * (uint32_t)k_block - 1; + if (n_dot_tiles <= 32) { + const uint32_t range = 2048u * n_dot_tiles - 1; Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); - asm volatile("" ::: "memory"); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + } else { + for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { + k_block = hex_smin(n_dot_tiles - k, 32); + const uint32_t range = 2048u * (uint32_t)k_block - 1; + Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); + Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); + asm volatile("" ::: "memory"); + row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + } } asm volatile("" ::: "memory"); - __fp16 *out_tile = output + (r * n_col_tiles + c) * HTP_MM_HMX_TILE_N_ELMS; Q6_mxmem_AR_after_hf(out_tile, 0); + + col_base += dot_stride; + out_tile += HTP_MM_HMX_TILE_N_ELMS; } } asm volatile("" ::: "memory"); @@ -658,29 +672,41 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co for (size_t i = 0; i < n_row_tiles; ++i) { const __fp16 *row_base = a + i * dot_tile_stride; __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; + const __fp16 *col_base = b; + __fp16 *accum_tile = res_base; + for (size_t j = 0; j < n_col_tiles; ++j) { Q6_mxclracc_hf(); asm volatile("" ::: "memory"); - const __fp16 *col_tiles = b + j * dot_tile_stride; + const __fp16 *col_tiles = col_base; const __fp16 *row_tiles = row_base; - __fp16 *accum_tile = res_base + j * HTP_MM_HMX_TILE_N_ELMS; + if (!zero_init) { Q6_activation_hf_mxmem_RR((unsigned int)accum_tile, 2047); Q6_weight_hf_mxmem_RR((unsigned int)eye_tile, 2047); } - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * k_block - 1; + if (n_dot_tiles <= 32) { + const uint32_t range = 2048u * n_dot_tiles - 1; Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + } else { + for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { + k_block = hex_smin(n_dot_tiles - k, 32); + const uint32_t range = 2048u * k_block - 1; + Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); + Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); + row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + } } asm volatile("" ::: "memory"); Q6_mxmem_AR_after_hf(accum_tile, 0); + + col_base += dot_tile_stride; + accum_tile += HTP_MM_HMX_TILE_N_ELMS; } } asm volatile("" ::: "memory"); From 69fecbc6ae8aa1dadd1e32a492efef0a7ab882af Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Wed, 1 Jul 2026 21:26:34 -0700 Subject: [PATCH 05/32] hex-hmx: use inline asm instead of intrinsics to avoid compiler issues --- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 284 +++++++++++++++------ 1 file changed, 208 insertions(+), 76 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index 439b594700ab..c9fe637b2b6d 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -53,49 +53,108 @@ static void hmx_fa_qk_dot_tile( __fp16 * out_tile, size_t n_dot_tiles ) { - asm volatile("" ::: "memory"); if (n_dot_tiles == 2) { - Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + HMX_FP16_TILE_N_ELMS), 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%3, %2)\n" + " weight.hf = mxmem(%4, %2)\n" + "}\n" + : + : "r"(row_tiles), "r"(col_tiles), "r"(2047), + "r"(row_tiles + HMX_FP16_TILE_N_ELMS), "r"(col_tiles + HMX_FP16_TILE_N_ELMS) + ); } else if (n_dot_tiles == 4) { - Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%3, %2)\n" + " weight.hf = mxmem(%4, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%5, %2)\n" + " weight.hf = mxmem(%6, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%7, %2)\n" + " weight.hf = mxmem(%8, %2)\n" + "}\n" + : + : "r"(row_tiles), "r"(col_tiles), "r"(2047), + "r"(row_tiles + HMX_FP16_TILE_N_ELMS), "r"(col_tiles + HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS) + ); } else if (n_dot_tiles == 8) { - Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 3 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 4 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 4 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 5 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 5 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 6 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 6 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (row_tiles + 7 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (col_tiles + 7 * HMX_FP16_TILE_N_ELMS), 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%3, %2)\n" + " weight.hf = mxmem(%4, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%5, %2)\n" + " weight.hf = mxmem(%6, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%7, %2)\n" + " weight.hf = mxmem(%8, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%9, %2)\n" + " weight.hf = mxmem(%10, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%11, %2)\n" + " weight.hf = mxmem(%12, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%13, %2)\n" + " weight.hf = mxmem(%14, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%15, %2)\n" + " weight.hf = mxmem(%16, %2)\n" + "}\n" + : + : "r"(row_tiles), "r"(col_tiles), "r"(2047), + "r"(row_tiles + HMX_FP16_TILE_N_ELMS), "r"(col_tiles + HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 4 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 4 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 5 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 5 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 6 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 6 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 7 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 7 * HMX_FP16_TILE_N_ELMS) + ); } else { 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); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + : + : "r"(row_tiles), "r"(col_tiles), "r"(2047) + ); row_tiles += HMX_FP16_TILE_N_ELMS; col_tiles += HMX_FP16_TILE_N_ELMS; } } - asm volatile("" ::: "memory"); - Q6_mxmem_AR_after_hf(out_tile, 0); - asm volatile("" ::: "memory"); + asm volatile( + "mxmem(%0, %1):after.hf = acc\n" + : + : "r"(out_tile), "r"(0) + : "memory" + ); } static inline void hmx_fa_o_update_tile( @@ -106,51 +165,116 @@ static inline void hmx_fa_o_update_tile( __fp16 * o_tile_out, size_t n_col_tiles ) { - asm volatile("" ::: "memory"); - Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + : + : "r"(d_diag), "r"(o_rc), "r"(2047) + ); if (n_col_tiles == 2) { - Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + HMX_FP16_TILE_N_ELMS), 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%3, %2)\n" + " weight.hf = mxmem(%4, %2)\n" + "}\n" + : + : "r"(p_tile_in), "r"(v_tile_in), "r"(2047), + "r"(p_tile_in + HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + HMX_FP16_TILE_N_ELMS) + ); } else if (n_col_tiles == 4) { - Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%3, %2)\n" + " weight.hf = mxmem(%4, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%5, %2)\n" + " weight.hf = mxmem(%6, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%7, %2)\n" + " weight.hf = mxmem(%8, %2)\n" + "}\n" + : + : "r"(p_tile_in), "r"(v_tile_in), "r"(2047), + "r"(p_tile_in + HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS) + ); } else if (n_col_tiles == 8) { - Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 3 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 4 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 4 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 5 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 5 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 6 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 6 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_activation_hf_mxmem_RR((unsigned int) (p_tile_in + 7 * HMX_FP16_TILE_N_ELMS), 2047); - Q6_weight_hf_mxmem_RR((unsigned int) (v_tile_in + 7 * HMX_FP16_TILE_N_ELMS), 2047); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%3, %2)\n" + " weight.hf = mxmem(%4, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%5, %2)\n" + " weight.hf = mxmem(%6, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%7, %2)\n" + " weight.hf = mxmem(%8, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%9, %2)\n" + " weight.hf = mxmem(%10, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%11, %2)\n" + " weight.hf = mxmem(%12, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%13, %2)\n" + " weight.hf = mxmem(%14, %2)\n" + "}\n" + "{\n" + " activation.hf = mxmem(%15, %2)\n" + " weight.hf = mxmem(%16, %2)\n" + "}\n" + : + : "r"(p_tile_in), "r"(v_tile_in), "r"(2047), + "r"(p_tile_in + HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 4 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 4 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 5 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 5 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 6 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 6 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 7 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 7 * HMX_FP16_TILE_N_ELMS) + ); } else { 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); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + : + : "r"(p_tile_in), "r"(v_tile_in), "r"(2047) + ); p_tile_in += HMX_FP16_TILE_N_ELMS; v_tile_in += HMX_FP16_TILE_N_ELMS; } } - asm volatile("" ::: "memory"); - Q6_mxmem_AR_after_hf(o_tile_out, 0); - asm volatile("" ::: "memory"); + asm volatile( + "mxmem(%0, %1):after.hf = acc\n" + : + : "r"(o_tile_out), "r"(0) + : "memory" + ); } static inline void hmx_fa_o_norm_tile( @@ -158,12 +282,20 @@ static inline void hmx_fa_o_norm_tile( const __fp16 * o_rc, __fp16 * o_out ) { - asm volatile("" ::: "memory"); - Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047); - Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047); - asm volatile("" ::: "memory"); - Q6_mxmem_AR_after_hf(o_out, 0); - asm volatile("" ::: "memory"); + asm volatile( + "{\n" + " activation.hf = mxmem(%0, %2)\n" + " weight.hf = mxmem(%1, %2)\n" + "}\n" + : + : "r"(d_diag), "r"(o_rc), "r"(2047) + ); + asm volatile( + "mxmem(%0, %1):after.hf = acc\n" + : + : "r"(o_out), "r"(0) + : "memory" + ); } #endif /* HMX_FA_KERNELS_H */ From c71f12785b7e7ec8d64c04e0245052eb7fb77445 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Wed, 1 Jul 2026 21:50:38 -0700 Subject: [PATCH 06/32] hex-hmx: define inline asm macros and simplify code --- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 204 +++++------------- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 66 +++--- ggml/src/ggml-hexagon/htp/hmx-utils.h | 22 ++ 3 files changed, 120 insertions(+), 172 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index c9fe637b2b6d..7977e04b8cb8 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -55,79 +55,40 @@ static void hmx_fa_qk_dot_tile( ) { if (n_dot_tiles == 2) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%3, %2)\n" - " weight.hf = mxmem(%4, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") + HMX_LOAD_MPY_F16("%3", "%4", "%0") : - : "r"(row_tiles), "r"(col_tiles), "r"(2047), - "r"(row_tiles + HMX_FP16_TILE_N_ELMS), "r"(col_tiles + HMX_FP16_TILE_N_ELMS) + : "r"(2047), + "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS) ); } else if (n_dot_tiles == 4) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%3, %2)\n" - " weight.hf = mxmem(%4, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%5, %2)\n" - " weight.hf = mxmem(%6, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%7, %2)\n" - " weight.hf = mxmem(%8, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") + HMX_LOAD_MPY_F16("%3", "%4", "%0") + HMX_LOAD_MPY_F16("%5", "%6", "%0") + HMX_LOAD_MPY_F16("%7", "%8", "%0") : - : "r"(row_tiles), "r"(col_tiles), "r"(2047), - "r"(row_tiles + HMX_FP16_TILE_N_ELMS), "r"(col_tiles + HMX_FP16_TILE_N_ELMS), + : "r"(2047), + "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS) ); } else if (n_dot_tiles == 8) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%3, %2)\n" - " weight.hf = mxmem(%4, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%5, %2)\n" - " weight.hf = mxmem(%6, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%7, %2)\n" - " weight.hf = mxmem(%8, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%9, %2)\n" - " weight.hf = mxmem(%10, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%11, %2)\n" - " weight.hf = mxmem(%12, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%13, %2)\n" - " weight.hf = mxmem(%14, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%15, %2)\n" - " weight.hf = mxmem(%16, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") + HMX_LOAD_MPY_F16("%3", "%4", "%0") + HMX_LOAD_MPY_F16("%5", "%6", "%0") + HMX_LOAD_MPY_F16("%7", "%8", "%0") + HMX_LOAD_MPY_F16("%9", "%10", "%0") + HMX_LOAD_MPY_F16("%11", "%12", "%0") + HMX_LOAD_MPY_F16("%13", "%14", "%0") + HMX_LOAD_MPY_F16("%15", "%16", "%0") : - : "r"(row_tiles), "r"(col_tiles), "r"(2047), - "r"(row_tiles + HMX_FP16_TILE_N_ELMS), "r"(col_tiles + HMX_FP16_TILE_N_ELMS), + : "r"(2047), + "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS), + "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(row_tiles + 4 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 4 * HMX_FP16_TILE_N_ELMS), @@ -138,19 +99,16 @@ static void hmx_fa_qk_dot_tile( } else { for (size_t k = 0; k < n_dot_tiles; ++k) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") : - : "r"(row_tiles), "r"(col_tiles), "r"(2047) + : "r"(2047), "r"(row_tiles), "r"(col_tiles) ); row_tiles += HMX_FP16_TILE_N_ELMS; col_tiles += HMX_FP16_TILE_N_ELMS; } } asm volatile( - "mxmem(%0, %1):after.hf = acc\n" + HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory" @@ -166,88 +124,46 @@ static inline void hmx_fa_o_update_tile( size_t n_col_tiles ) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") : - : "r"(d_diag), "r"(o_rc), "r"(2047) + : "r"(2047), "r"(d_diag), "r"(o_rc) ); if (n_col_tiles == 2) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%3, %2)\n" - " weight.hf = mxmem(%4, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") + HMX_LOAD_MPY_F16("%3", "%4", "%0") : - : "r"(p_tile_in), "r"(v_tile_in), "r"(2047), - "r"(p_tile_in + HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + HMX_FP16_TILE_N_ELMS) + : "r"(2047), + "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS) ); } else if (n_col_tiles == 4) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%3, %2)\n" - " weight.hf = mxmem(%4, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%5, %2)\n" - " weight.hf = mxmem(%6, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%7, %2)\n" - " weight.hf = mxmem(%8, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") + HMX_LOAD_MPY_F16("%3", "%4", "%0") + HMX_LOAD_MPY_F16("%5", "%6", "%0") + HMX_LOAD_MPY_F16("%7", "%8", "%0") : - : "r"(p_tile_in), "r"(v_tile_in), "r"(2047), - "r"(p_tile_in + HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + HMX_FP16_TILE_N_ELMS), + : "r"(2047), + "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS) ); } else if (n_col_tiles == 8) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%3, %2)\n" - " weight.hf = mxmem(%4, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%5, %2)\n" - " weight.hf = mxmem(%6, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%7, %2)\n" - " weight.hf = mxmem(%8, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%9, %2)\n" - " weight.hf = mxmem(%10, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%11, %2)\n" - " weight.hf = mxmem(%12, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%13, %2)\n" - " weight.hf = mxmem(%14, %2)\n" - "}\n" - "{\n" - " activation.hf = mxmem(%15, %2)\n" - " weight.hf = mxmem(%16, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") + HMX_LOAD_MPY_F16("%3", "%4", "%0") + HMX_LOAD_MPY_F16("%5", "%6", "%0") + HMX_LOAD_MPY_F16("%7", "%8", "%0") + HMX_LOAD_MPY_F16("%9", "%10", "%0") + HMX_LOAD_MPY_F16("%11", "%12", "%0") + HMX_LOAD_MPY_F16("%13", "%14", "%0") + HMX_LOAD_MPY_F16("%15", "%16", "%0") : - : "r"(p_tile_in), "r"(v_tile_in), "r"(2047), - "r"(p_tile_in + HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + HMX_FP16_TILE_N_ELMS), + : "r"(2047), + "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS), + "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(p_tile_in + 4 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 4 * HMX_FP16_TILE_N_ELMS), @@ -258,19 +174,16 @@ static inline void hmx_fa_o_update_tile( } else { for (size_t k = 0; k < n_col_tiles; ++k) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") : - : "r"(p_tile_in), "r"(v_tile_in), "r"(2047) + : "r"(2047), "r"(p_tile_in), "r"(v_tile_in) ); p_tile_in += HMX_FP16_TILE_N_ELMS; v_tile_in += HMX_FP16_TILE_N_ELMS; } } asm volatile( - "mxmem(%0, %1):after.hf = acc\n" + HMX_STORE_AFTER_F16("%0", "%1") : : "r"(o_tile_out), "r"(0) : "memory" @@ -283,15 +196,12 @@ static inline void hmx_fa_o_norm_tile( __fp16 * o_out ) { asm volatile( - "{\n" - " activation.hf = mxmem(%0, %2)\n" - " weight.hf = mxmem(%1, %2)\n" - "}\n" + HMX_LOAD_MPY_F16("%1", "%2", "%0") : - : "r"(d_diag), "r"(o_rc), "r"(2047) + : "r"(2047), "r"(d_diag), "r"(o_rc) ); asm volatile( - "mxmem(%0, %1):after.hf = acc\n" + HMX_STORE_AFTER_F16("%0", "%1") : : "r"(o_out), "r"(0) : "memory" 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 21682a1b3d19..c814623034fe 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -613,11 +613,10 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); - Q6_bias_mxmem2_A((void *)scales); + asm volatile(HMX_SET_BIAS("%0") :: "r"(scales)); const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; - asm volatile("" ::: "memory"); for (uint32_t r = 0; r < n_row_tiles; ++r) { const __fp16 *row_base = activation + r * dot_stride; const __fp16 *col_base = weight; @@ -627,33 +626,40 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict const __fp16 *row_tiles = row_base; const __fp16 *col_tiles = col_base; - Q6_mxclracc_hf(); - asm volatile("" ::: "memory"); + asm volatile(HMX_CLRACC_F16()); if (n_dot_tiles <= 32) { const uint32_t range = 2048u * n_dot_tiles - 1; - Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); - Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); + asm volatile( + HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") + : + : "r"(range), "r"(row_tiles), "r"(col_tiles) + ); } else { for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { k_block = hex_smin(n_dot_tiles - k, 32); const uint32_t range = 2048u * (uint32_t)k_block - 1; - Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); - Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); - asm volatile("" ::: "memory"); + asm volatile( + HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") + : + : "r"(range), "r"(row_tiles), "r"(col_tiles) + ); row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; } } - asm volatile("" ::: "memory"); - Q6_mxmem_AR_after_hf(out_tile, 0); + asm volatile( + HMX_STORE_AFTER_F16("%0", "%1") + : + : "r"(out_tile), "r"(0) + : "memory" + ); col_base += dot_stride; out_tile += HTP_MM_HMX_TILE_N_ELMS; } } - asm volatile("" ::: "memory"); } // C += AB @@ -664,11 +670,10 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); - Q6_bias_mxmem2_A((void *)col_scales); + asm volatile(HMX_SET_BIAS("%0") :: "r"(col_scales)); const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; - asm volatile("" ::: "memory"); for (size_t i = 0; i < n_row_tiles; ++i) { const __fp16 *row_base = a + i * dot_tile_stride; __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; @@ -676,40 +681,51 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co __fp16 *accum_tile = res_base; for (size_t j = 0; j < n_col_tiles; ++j) { - Q6_mxclracc_hf(); - asm volatile("" ::: "memory"); + asm volatile(HMX_CLRACC_F16()); const __fp16 *col_tiles = col_base; const __fp16 *row_tiles = row_base; if (!zero_init) { - Q6_activation_hf_mxmem_RR((unsigned int)accum_tile, 2047); - Q6_weight_hf_mxmem_RR((unsigned int)eye_tile, 2047); + asm volatile( + HMX_LOAD_MPY_F16("%1", "%2", "%0") + : + : "r"(2047), "r"(accum_tile), "r"(eye_tile) + ); } if (n_dot_tiles <= 32) { const uint32_t range = 2048u * n_dot_tiles - 1; - Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); - Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); + asm volatile( + HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") + : + : "r"(range), "r"(row_tiles), "r"(col_tiles) + ); } else { for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { k_block = hex_smin(n_dot_tiles - k, 32); const uint32_t range = 2048u * k_block - 1; - Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); - Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); + asm volatile( + HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") + : + : "r"(range), "r"(row_tiles), "r"(col_tiles) + ); row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; } } - asm volatile("" ::: "memory"); - Q6_mxmem_AR_after_hf(accum_tile, 0); + asm volatile( + HMX_STORE_AFTER_F16("%0", "%1") + : + : "r"(accum_tile), "r"(0) + : "memory" + ); col_base += dot_tile_stride; accum_tile += HTP_MM_HMX_TILE_N_ELMS; } } - asm volatile("" ::: "memory"); } // --- Async HMX matmul job (for pipeline overlap) --- diff --git a/ggml/src/ggml-hexagon/htp/hmx-utils.h b/ggml/src/ggml-hexagon/htp/hmx-utils.h index 561d26944f14..2a61ca7349df 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-utils.h +++ b/ggml/src/ggml-hexagon/htp/hmx-utils.h @@ -197,4 +197,26 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out, } } +// --- HMX inline asm macros for load-store packetization --- +#define HMX_LOAD_MPY_F16(act, wt, range) \ + "{\n" \ + " activation.hf = mxmem(" act ", " range ")\n" \ + " weight.hf = mxmem(" wt ", " range ")\n" \ + "}\n" + +#define HMX_LOAD_MPY_DEEP_F16(act, wt, range) \ + "{\n" \ + " activation.hf = mxmem(" act ", " range "):deep\n" \ + " weight.hf = mxmem(" wt ", " range ")\n" \ + "}\n" + +#define HMX_STORE_AFTER_F16(out, scale_reg) \ + "mxmem(" out ", " scale_reg "):after.hf = acc\n" + +#define HMX_SET_BIAS(scales) \ + "bias = mxmem2(" scales ")\n" + +#define HMX_CLRACC_F16() \ + "mxclracc.hf\n" + #endif // HMX_UTILS_H From 6290eeb2eb9f4a22ad15e50195c49734cdd818e2 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 14:12:55 -0700 Subject: [PATCH 07/32] hex-hmx: replace leftover intrinsics --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 12 ++++++------ ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 10 ++++++++-- ggml/src/ggml-hexagon/htp/matmul-ops.c | 6 +++--- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 8c9f4889cd1c..1ad1ef04d2d9 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1330,7 +1330,7 @@ static void hmx_fa_qk_dot_worker(void * data) { __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); - Q6_bias_mxmem2_A((void *) job->hmx_scales); + asm volatile(HMX_SET_BIAS("%0") :: "r"(job->hmx_scales)); const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * row_tiles = q_tiles + r * dot_stride; @@ -1375,7 +1375,7 @@ static void hmx_fa_o_update_worker(void * data) { __builtin_assume(n_col_tiles > 0); __builtin_assume(DV_tiles > 0); - Q6_bias_mxmem2_A((void *) job->hmx_scales); + asm volatile(HMX_SET_BIAS("%0") :: "r"(job->hmx_scales)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { @@ -1415,7 +1415,7 @@ static void hmx_fa_o_norm_worker(void * data) { __builtin_assume(n_row_tiles > 0); __builtin_assume(DV_tiles > 0); - Q6_bias_mxmem2_A((void *) job->hmx_scales); + asm volatile(HMX_SET_BIAS("%0") :: "r"(job->hmx_scales)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; @@ -1879,7 +1879,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { __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); + asm volatile(HMX_SET_BIAS("%0") :: "r"(factx.vtcm_hmx_scales_qk)); const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * row_tiles = q_base + r * dot_stride; @@ -1941,7 +1941,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { __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); + asm volatile(HMX_SET_BIAS("%0") :: "r"(factx.vtcm_hmx_scales_id)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { @@ -1991,7 +1991,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { __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); + asm volatile(HMX_SET_BIAS("%0") :: "r"(factx.vtcm_hmx_scales_id)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; 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 c814623034fe..16cc36aa0a14 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -607,8 +607,8 @@ static void quantize_f32_weight_to_fp16_tiles_task( // --- End tiled dequantizers --- // requires external HMX lock -static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { +static inline void core_dot_chunk_fp16_inline(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { __builtin_assume(n_row_tiles > 0); __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); @@ -662,6 +662,12 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict } } +static __attribute__((noinline)) +void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { + core_dot_chunk_fp16_inline(output, activation, weight, scales, n_row_tiles, n_col_tiles, n_dot_tiles); +} + // C += AB static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b, const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile, diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index cd8c5308dd1a..f945ad70e051 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2517,8 +2517,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, g); { const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; - core_dot_chunk_fp16(vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, - params->k / 32); + core_dot_chunk_fp16(vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, g); @@ -2527,7 +2526,8 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ 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, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, params->src2_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); } } } From 71f04e1a678d82e3c407a3e60fa3b9f539f229a7 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 15:19:53 -0700 Subject: [PATCH 08/32] hmx-fa: minor cleanup for hmx asm --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 12 ++--- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 2 +- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 49 ++----------------- ggml/src/ggml-hexagon/htp/matmul-ops.c | 34 +++++++++++++ 4 files changed, 45 insertions(+), 52 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 1ad1ef04d2d9..3adbd8b4d579 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1330,7 +1330,7 @@ static void hmx_fa_qk_dot_worker(void * data) { __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); - asm volatile(HMX_SET_BIAS("%0") :: "r"(job->hmx_scales)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales)); const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * row_tiles = q_tiles + r * dot_stride; @@ -1375,7 +1375,7 @@ static void hmx_fa_o_update_worker(void * data) { __builtin_assume(n_col_tiles > 0); __builtin_assume(DV_tiles > 0); - asm volatile(HMX_SET_BIAS("%0") :: "r"(job->hmx_scales)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { @@ -1415,7 +1415,7 @@ static void hmx_fa_o_norm_worker(void * data) { __builtin_assume(n_row_tiles > 0); __builtin_assume(DV_tiles > 0); - asm volatile(HMX_SET_BIAS("%0") :: "r"(job->hmx_scales)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; @@ -1879,7 +1879,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { __builtin_assume(n_dot_tiles > 0); htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); - asm volatile(HMX_SET_BIAS("%0") :: "r"(factx.vtcm_hmx_scales_qk)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)factx.vtcm_hmx_scales_qk)); const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * row_tiles = q_base + r * dot_stride; @@ -1941,7 +1941,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { __builtin_assume(DV_tiles > 0); htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); - asm volatile(HMX_SET_BIAS("%0") :: "r"(factx.vtcm_hmx_scales_id)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)factx.vtcm_hmx_scales_id)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { @@ -1991,7 +1991,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { __builtin_assume(DV_tiles > 0); htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); - asm volatile(HMX_SET_BIAS("%0") :: "r"(factx.vtcm_hmx_scales_id)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)factx.vtcm_hmx_scales_id)); const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; for (size_t r = 0; r < n_row_tiles; ++r) { const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index 7977e04b8cb8..f866dc1a166f 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -115,7 +115,7 @@ static void hmx_fa_qk_dot_tile( ); } -static inline void hmx_fa_o_update_tile( +static void hmx_fa_o_update_tile( const __fp16 * d_diag, const __fp16 * o_rc, const __fp16 * p_tile_in, 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 16cc36aa0a14..c995bdce8117 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -607,13 +607,13 @@ static void quantize_f32_weight_to_fp16_tiles_task( // --- End tiled dequantizers --- // requires external HMX lock -static inline void core_dot_chunk_fp16_inline(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { +static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { __builtin_assume(n_row_tiles > 0); __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); - asm volatile(HMX_SET_BIAS("%0") :: "r"(scales)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales)); const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; @@ -662,12 +662,6 @@ static inline void core_dot_chunk_fp16_inline(__fp16 *restrict output, const __f } } -static __attribute__((noinline)) -void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { - core_dot_chunk_fp16_inline(output, activation, weight, scales, n_row_tiles, n_col_tiles, n_dot_tiles); -} - // C += AB static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b, const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile, @@ -676,7 +670,7 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); - asm volatile(HMX_SET_BIAS("%0") :: "r"(col_scales)); + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales)); const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; @@ -734,41 +728,6 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co } } -// --- Async HMX matmul job (for pipeline overlap) --- - -typedef struct { - __fp16 * output; - const __fp16 * activation; - const __fp16 * weight; - const __fp16 * scales; - uint32_t n_row_tiles; - uint32_t n_col_tiles; - uint32_t n_dot_tiles; -} hmx_matmul_job_t; - -static void hmx_matmul_worker_fn(void * data) { - hmx_matmul_job_t * job = (hmx_matmul_job_t *) data; - FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles); - core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles); -} - -static inline void hmx_matmul_job_init(hmx_matmul_job_t * job, - __fp16 * output, - const __fp16 * activation, - const __fp16 * weight, - const __fp16 * scales, - uint32_t n_row_tiles, - uint32_t n_col_tiles, - uint32_t n_dot_tiles) { - job->output = output; - job->activation = activation; - job->weight = weight; - job->scales = scales; - job->n_row_tiles = n_row_tiles; - job->n_col_tiles = n_col_tiles; - job->n_dot_tiles = n_dot_tiles; -} - // output : fp16 -> f32p static void transfer_output_chunk_fp16_to_fp32( diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index f945ad70e051..37ea2e2068ff 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2127,6 +2127,40 @@ static void transfer_activation_chunk_threaded( worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, n_tasks); } } +// --- Async HMX matmul job (for pipeline overlap) --- + +typedef struct { + __fp16 * output; + const __fp16 * activation; + const __fp16 * weight; + const __fp16 * scales; + uint32_t n_row_tiles; + uint32_t n_col_tiles; + uint32_t n_dot_tiles; +} hmx_matmul_job_t; + +static void hmx_matmul_worker_fn(void * data) { + hmx_matmul_job_t * job = (hmx_matmul_job_t *) data; + FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles); + core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles); +} + +static inline void hmx_matmul_job_init(hmx_matmul_job_t * job, + __fp16 * output, + const __fp16 * activation, + const __fp16 * weight, + const __fp16 * scales, + uint32_t n_row_tiles, + uint32_t n_col_tiles, + uint32_t n_dot_tiles) { + job->output = output; + job->activation = activation; + job->weight = weight; + job->scales = scales; + job->n_row_tiles = n_row_tiles; + job->n_col_tiles = n_col_tiles; + job->n_dot_tiles = n_dot_tiles; +} static int hmx_mm_2d_f32(struct htp_context *ctx, float *restrict dst, From 1c181947e1d99533e10ceddea7803a9b13ae6f30 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 15:26:53 -0700 Subject: [PATCH 09/32] hmx-mm: move per-task stucts out of the kernels header --- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 35 ------------------ ggml/src/ggml-hexagon/htp/matmul-ops.c | 37 +++++++++++++++++++ 2 files changed, 37 insertions(+), 35 deletions(-) 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 c995bdce8117..a5825f170781 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -1024,42 +1024,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( } } -typedef struct { - const struct mmid_row_mapping *matrix_rows; - __fp16 *dst; - const float *src; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t k_block; - uint32_t cur_a; - uint32_t mapping_stride; - uint32_t ne11; - struct fastdiv_values ne11_div; - size_t nb11; - size_t nb12; - uint32_t start_row; - uint32_t cne1; - uint32_t k_valid; - struct htp_thread_trace *traces; -} activation_transfer_gathered_task_state_t; -typedef struct { - const struct mmid_row_mapping *matrix_rows; - const __fp16 *vtcm_src; - float *dst; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t n_cols; - uint32_t cur_a; - uint32_t mapping_stride; - size_t dst_nb1; - size_t dst_nb2; - uint32_t start_row; - uint32_t cne1; - struct htp_thread_trace *traces; -} output_transfer_scattered_task_state_t; static void transfer_activation_chunk_fp32_to_fp16_gathered( __fp16 *restrict vtcm_dst, diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 37ea2e2068ff..46633a55f93b 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -1965,6 +1965,43 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx); } +typedef struct { + const struct mmid_row_mapping *matrix_rows; + __fp16 *dst; + const float *src; + uint32_t n_tasks; + uint32_t n_tot_chunks; + uint32_t n_chunks_per_task; + uint32_t k_block; + uint32_t cur_a; + uint32_t mapping_stride; + uint32_t ne11; + struct fastdiv_values ne11_div; + size_t nb11; + size_t nb12; + uint32_t start_row; + uint32_t cne1; + uint32_t k_valid; + struct htp_thread_trace *traces; +} activation_transfer_gathered_task_state_t; + +typedef struct { + const struct mmid_row_mapping *matrix_rows; + const __fp16 *vtcm_src; + float *dst; + uint32_t n_tasks; + uint32_t n_tot_chunks; + uint32_t n_chunks_per_task; + uint32_t n_cols; + uint32_t cur_a; + uint32_t mapping_stride; + size_t dst_nb1; + size_t dst_nb2; + uint32_t start_row; + uint32_t cne1; + struct htp_thread_trace *traces; +} output_transfer_scattered_task_state_t; + static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_gathered_task_state_t *st = data; struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL; From 6491fef3accc6cfaa02cb632dec00fe0d74fb8c5 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 15:56:28 -0700 Subject: [PATCH 10/32] hmx-mm: simplify core_dot_chunk --- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 68 +++++++++++-------- 1 file changed, 40 insertions(+), 28 deletions(-) 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 a5825f170781..43771f4fad0f 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -606,12 +606,42 @@ static void quantize_f32_weight_to_fp16_tiles_task( // --- End tiled dequantizers --- -// requires external HMX lock -static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { +static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { __builtin_assume(n_row_tiles > 0); __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); + __builtin_assume(n_dot_tiles <= 32); + + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales)); + + const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + const uint32_t range = 2048u * n_dot_tiles - 1; + + for (uint32_t r = 0; r < n_row_tiles; ++r) { + const __fp16 *row_base = activation + r * dot_stride; + const __fp16 *col_base = weight; + __fp16 *out_tile = output + r * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; + + for (size_t c = 0; c < n_col_tiles; ++c) { + asm volatile(HMX_CLRACC_F16()); + asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_base), "r"(col_base)); + asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory"); + col_base += dot_stride; + out_tile += HTP_MM_HMX_TILE_N_ELMS; + } + } +} + +static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { + if (n_dot_tiles <= 32) { + core_dot_chunk_fp16_short(output, activation, weight, scales, n_row_tiles, n_col_tiles, n_dot_tiles); + return; + } + __builtin_assume(n_row_tiles > 0); + __builtin_assume(n_col_tiles > 0); + __builtin_assume(n_dot_tiles > 32); asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales)); @@ -628,33 +658,15 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict asm volatile(HMX_CLRACC_F16()); - if (n_dot_tiles <= 32) { - const uint32_t range = 2048u * n_dot_tiles - 1; - asm volatile( - HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") - : - : "r"(range), "r"(row_tiles), "r"(col_tiles) - ); - } else { - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * (uint32_t)k_block - 1; - asm volatile( - HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") - : - : "r"(range), "r"(row_tiles), "r"(col_tiles) - ); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - } + for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { + k_block = hex_smin(n_dot_tiles - k, 32); + const uint32_t range = 2048u * (uint32_t)k_block - 1; + asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles)); + row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; + col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; } - asm volatile( - HMX_STORE_AFTER_F16("%0", "%1") - : - : "r"(out_tile), "r"(0) - : "memory" - ); + asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory"); col_base += dot_stride; out_tile += HTP_MM_HMX_TILE_N_ELMS; From dd1820c970bb07c981013c101563bb386c25a408 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 16:16:19 -0700 Subject: [PATCH 11/32] hmx-mm: simplify inner loops that call hmx instructions --- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 131 +++++++++++------- 1 file changed, 78 insertions(+), 53 deletions(-) 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 43771f4fad0f..a91ea2d0d26a 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -506,7 +506,8 @@ static void dequantize_tiled_weight_to_fp16_task_q8_0( } } -static void convert_f16_weight_to_fp16_tiles_task( +static __attribute__((noinline)) +void convert_f16_weight_to_fp16_tiles_task( const tiled_dequantize_state_t *state, uint32_t start_tile, uint32_t end_tile) { @@ -543,17 +544,13 @@ static void convert_f16_weight_to_fp16_tiles_task( Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v1); v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); } - (void) *(volatile HVX_Vector *)(tile_base); } ++t; ++kt; } - - if (start_tile < end_tile) { - (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS); - } } -static void quantize_f32_weight_to_fp16_tiles_task( +static __attribute__((noinline)) +void quantize_f32_weight_to_fp16_tiles_task( const tiled_dequantize_state_t *state, uint32_t start_tile, uint32_t end_tile) { @@ -594,20 +591,18 @@ static void quantize_f32_weight_to_fp16_tiles_task( Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out_hi); v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); } - (void) *(volatile HVX_Vector *)(tile_base); } ++t; ++kt; } - - if (start_tile < end_tile) { - (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS); - } } // --- End tiled dequantizers --- -static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { +// dot-chunk functions require external HMX lock + +static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *restrict activation, + const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { __builtin_assume(n_row_tiles > 0); __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); @@ -633,8 +628,9 @@ static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *res } } -static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { +static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, + const __fp16 *restrict weight, const __fp16 *restrict scales, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { if (n_dot_tiles <= 32) { core_dot_chunk_fp16_short(output, activation, weight, scales, n_row_tiles, n_col_tiles, n_dot_tiles); return; @@ -658,12 +654,18 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict asm volatile(HMX_CLRACC_F16()); - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * (uint32_t)k_block - 1; + const uint32_t n_loops = n_dot_tiles / 32; + const uint32_t rem = n_dot_tiles % 32; + + for (uint32_t l = 0; l < n_loops; ++l) { + asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(65535), "r"(row_tiles), "r"(col_tiles)); + row_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS; + col_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS; + } + + if (rem > 0) { + const uint32_t range = 2048u * rem - 1; asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles)); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; } asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory"); @@ -674,17 +676,18 @@ static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict } } -// C += AB -static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b, +static void core_mma_chunk_fp16_short(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b, const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile, uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) { __builtin_assume(n_row_tiles > 0); __builtin_assume(n_col_tiles > 0); __builtin_assume(n_dot_tiles > 0); + __builtin_assume(n_dot_tiles <= 32); asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales)); const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + const uint32_t range = 2048u * n_dot_tiles - 1; for (size_t i = 0; i < n_row_tiles; ++i) { const __fp16 *row_base = a + i * dot_tile_stride; @@ -695,44 +698,66 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co for (size_t j = 0; j < n_col_tiles; ++j) { asm volatile(HMX_CLRACC_F16()); + if (!zero_init) { + asm volatile(HMX_LOAD_MPY_F16("%1", "%2", "%0") : : "r"(2047), "r"(accum_tile), "r"(eye_tile)); + } + + asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_base), "r"(col_base)); + + asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(accum_tile), "r"(0) : "memory"); + + col_base += dot_tile_stride; + accum_tile += HTP_MM_HMX_TILE_N_ELMS; + } + } +} + +static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b, + const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile, + uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) { + if (n_dot_tiles <= 32) { + core_mma_chunk_fp16_short(c, a, b, col_scales, eye_tile, n_row_tiles, n_col_tiles, n_dot_tiles, zero_init); + return; + } + __builtin_assume(n_row_tiles > 0); + __builtin_assume(n_col_tiles > 0); + __builtin_assume(n_dot_tiles > 32); + + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales)); + + const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; + + for (size_t i = 0; i < n_row_tiles; ++i) { + const __fp16 *row_base = a + i * dot_tile_stride; + __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; + const __fp16 *col_base = b; + __fp16 *accum_tile = res_base; + + for (size_t j = 0; j < n_col_tiles; ++j) { const __fp16 *col_tiles = col_base; const __fp16 *row_tiles = row_base; + asm volatile(HMX_CLRACC_F16()); + if (!zero_init) { - asm volatile( - HMX_LOAD_MPY_F16("%1", "%2", "%0") - : - : "r"(2047), "r"(accum_tile), "r"(eye_tile) - ); + asm volatile(HMX_LOAD_MPY_F16("%1", "%2", "%0") : : "r"(2047), "r"(accum_tile), "r"(eye_tile)); } - if (n_dot_tiles <= 32) { - const uint32_t range = 2048u * n_dot_tiles - 1; - asm volatile( - HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") - : - : "r"(range), "r"(row_tiles), "r"(col_tiles) - ); - } else { - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * k_block - 1; - asm volatile( - HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") - : - : "r"(range), "r"(row_tiles), "r"(col_tiles) - ); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - } + const uint32_t n_loops = n_dot_tiles / 32; + const uint32_t rem = n_dot_tiles % 32; + + for (uint32_t l = 0; l < n_loops; ++l) { + asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(65535), "r"(row_tiles), "r"(col_tiles)); + row_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS; + col_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS; + } + + if (rem > 0) { + const uint32_t range = 2048u * rem - 1; + asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles)); } - asm volatile( - HMX_STORE_AFTER_F16("%0", "%1") - : - : "r"(accum_tile), "r"(0) - : "memory" - ); + asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(accum_tile), "r"(0) : "memory"); col_base += dot_tile_stride; accum_tile += HTP_MM_HMX_TILE_N_ELMS; From abea2a9f11930d8a8dea9a21dabd3989efa8c47c Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 17:22:57 -0700 Subject: [PATCH 12/32] hmx-mm: proper instrumentation for activation prep work for dma pipelined version --- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 126 +++++------------- ggml/src/ggml-hexagon/htp/matmul-ops.c | 80 ++++++++++- 2 files changed, 108 insertions(+), 98 deletions(-) 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 a91ea2d0d26a..96371f2b89bf 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -956,113 +956,55 @@ static void transfer_activation_chunk_fp32_to_fp16(__fp16 *restrict vtcm_dst, co } } -typedef struct { - __fp16 *dst; - const float *src; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t k_block; - uint32_t k_stride; - uint32_t k_valid; - struct htp_thread_trace * traces; - struct htp_context * ctx; - float * vtcm_f32_act; -} activation_transfer_task_state_t; - -static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( - dma_queue *dma_q, +static void transfer_activation_row_pair_fp32_to_fp16( __fp16 *restrict vtcm_dst, - const float *restrict src, - uint32_t n_rows, + const float *restrict row0, + const float *restrict row1, + uint32_t r, uint32_t k_block, - uint32_t k_stride, uint32_t k_valid, - float *thread_f32_act) { - - const uint32_t R = HTP_MM_DMA_ACT_ROWS_PER_STEP; - const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); - - const uint32_t n_steps = n_rows_padded / R; - - // pre-fetch step 0 - if (n_steps > 0 && n_rows > 0) { - uint32_t nrows_to_fetch = hex_smin(n_rows, R); - dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src), - k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); - } + bool row0_valid, + bool row1_valid) { - for (uint32_t s = 0; s < n_steps; ++s) { - uint32_t r = R * s; - float *curr_buf = thread_f32_act + (s % 2) * R * k_block; + uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index + uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - if (r < n_rows) { - dma_queue_pop(dma_q); - } + uint32_t c = 0; + for (; c + 32 <= k_valid; c += 32) { + HVX_Vector v0 = Q6_V_vzero(); + HVX_Vector v1 = Q6_V_vzero(); + if (row0_valid) v0 = *(const HVX_Vector *)(row0 + c); + if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c); - uint32_t next_s = s + 1; - uint32_t next_r = R * next_s; - if (next_r < n_rows) { - uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); - const float *next_src = src + next_r * k_stride; - float *next_buf = thread_f32_act + (next_s % 2) * R * k_block; - dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), - k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); - } + HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - #pragma unroll - for (uint32_t i = 0; i < HTP_MM_DMA_ACT_ROWS_PER_STEP; i += 2) { - uint32_t curr_r = r + i; - const bool row0_valid = (curr_r < n_rows); - const bool row1_valid = (curr_r + 1) < n_rows; + uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index + uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - const float *ptr_in0 = curr_buf + i * k_block; - const float *ptr_in1 = curr_buf + (i + 1) * k_block; + HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); + tile[r1 / 2] = v_out; + } + if (c < k_block) { + HVX_Vector v0 = Q6_V_vzero(); + HVX_Vector v1 = Q6_V_vzero(); + if (row0_valid) v0 = *(const HVX_Vector *)(row0 + c); + if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c); - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); + uint32_t rem = k_valid - c; + HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); + v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); + v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); + HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - uint32_t r0 = curr_r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = curr_r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; + uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index + uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t r0 = curr_r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = curr_r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } + HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); + tile[r1 / 2] = v_out; } } - - static void transfer_activation_chunk_fp32_to_fp16_gathered( __fp16 *restrict vtcm_dst, const float *restrict src, diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 46633a55f93b..8af133aa57e2 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -1937,14 +1937,82 @@ static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx); } +typedef struct { + const struct mmid_row_mapping *matrix_rows; + __fp16 *dst; + const float *src; + uint32_t n_tasks; + uint32_t n_tot_chunks; + uint32_t n_chunks_per_task; + uint32_t k_block; + uint32_t k_stride; + uint32_t k_valid; + struct htp_thread_trace * traces; + struct htp_context * ctx; + float * vtcm_f32_act; +} activation_transfer_task_state_t; + +static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( + dma_queue *dma_q, + __fp16 *restrict vtcm_dst, + const float *restrict src, + uint32_t n_rows, + uint32_t k_block, + uint32_t k_stride, + uint32_t k_valid, + float *thread_f32_act, + struct htp_thread_trace *tr) { + + const uint32_t R = HTP_MM_DMA_ACT_ROWS_PER_STEP; + const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); + + const uint32_t n_steps = n_rows_padded / R; + + // pre-fetch step 0 + if (n_steps > 0 && n_rows > 0) { + uint32_t nrows_to_fetch = hex_smin(n_rows, R); + dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src), + k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); + } + + for (uint32_t s = 0; s < n_steps; ++s) { + uint32_t r = R * s; + float *curr_buf = thread_f32_act + (s % 2) * R * k_block; + + if (r < n_rows) { + dma_queue_pop(dma_q); + } + + uint32_t next_s = s + 1; + uint32_t next_r = R * next_s; + if (next_r < n_rows) { + uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); + const float *next_src = src + next_r * k_stride; + float *next_buf = thread_f32_act + (next_s % 2) * R * k_block; + dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), + k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); + } + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r); + transfer_activation_row_pair_fp32_to_fp16( + vtcm_dst, + curr_buf, + curr_buf + k_block, + r, + k_block, + k_valid, + (r < n_rows), + ((r + 1) < n_rows) + ); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r); + } +} + static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_task_state_t *st = (activation_transfer_task_state_t *) data; struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL; - int start_chunk_idx = i * st->n_chunks_per_task; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx); - for (unsigned int task_id = i; task_id < (unsigned int)st->n_tasks; task_id += n) { int chunk_idx = task_id * st->n_chunks_per_task; size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task); @@ -1955,14 +2023,14 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, if (st->vtcm_f32_act) { float *thread_f32_act = st->vtcm_f32_act + i * HTP_MM_DMA_ACT_MULTIPLIER * st->k_block; transfer_activation_chunk_fp32_to_fp16_dma_pipelined( - st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act + st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr ); } else { + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); transfer_activation_chunk_fp32_to_fp16(dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); } } - - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx); } typedef struct { From 44fa4d1901b32ec90afde30df8c3491e4c68bcc8 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Thu, 2 Jul 2026 17:56:40 -0700 Subject: [PATCH 13/32] hmx-mm: update a-prep loop for better prefetch --- ggml/src/ggml-hexagon/htp/matmul-ops.c | 37 +++++++++++++++++--------- ggml/src/ggml-hexagon/htp/matmul-ops.h | 2 +- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 8af133aa57e2..f56dfd33d987 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -1968,30 +1968,31 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( const uint32_t n_steps = n_rows_padded / R; - // pre-fetch step 0 + // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } - - for (uint32_t s = 0; s < n_steps; ++s) { - uint32_t r = R * s; - float *curr_buf = thread_f32_act + (s % 2) * R * k_block; - - if (r < n_rows) { - dma_queue_pop(dma_q); - } - - uint32_t next_s = s + 1; - uint32_t next_r = R * next_s; + // Push step 1 (if valid) + if (n_steps > 1) { + uint32_t next_r = R * 1; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride; - float *next_buf = thread_f32_act + (next_s % 2) * R * k_block; + float *next_buf = thread_f32_act + 1 * R * k_block; dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } + } + + for (uint32_t s = 0; s < n_steps; ++s) { + uint32_t r = R * s; + float *curr_buf = NULL; + + if (r < n_rows) { + curr_buf = (float *) dma_queue_pop(dma_q).dst; + } htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r); transfer_activation_row_pair_fp32_to_fp16( @@ -2005,6 +2006,16 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( ((r + 1) < n_rows) ); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r); + + // Push step s + 2 + uint32_t next_s = s + 2; + uint32_t next_r = R * next_s; + if (next_r < n_rows) { + uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); + const float *next_src = src + next_r * k_stride; + dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src), + k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); + } } } diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index c3752d4dba62..42dafd271af7 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -45,7 +45,7 @@ extern "C" { // --- DMA Activation Transfer Configuration --- #define HTP_MM_DMA_ACT_ROWS_PER_STEP 2 -#define HTP_MM_DMA_ACT_MULTIPLIER 4 +#define HTP_MM_DMA_ACT_MULTIPLIER (2 * HTP_MM_DMA_ACT_ROWS_PER_STEP) enum htp_mm_kernel_type { HTP_MM_KERNEL_UNSUPPORTED = 0, From e4dd7a8e665cf42d51994730d54da84883364e0a Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Sat, 4 Jul 2026 12:15:41 -0700 Subject: [PATCH 14/32] hex-vtcm: improved vtcm layout alloc for mm to support overlapping areas --- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 44 ++++----- ggml/src/ggml-hexagon/htp/htp-vtcm.h | 4 +- ggml/src/ggml-hexagon/htp/matmul-ops.c | 106 ++++++++++++++------- ggml/src/ggml-hexagon/htp/matmul-ops.h | 93 ++++++++++++------ 4 files changed, 157 insertions(+), 90 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index abf947f6ba47..b834cec19de1 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -154,28 +154,28 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128); size_t off = 0; - VTCM_LAYOUT_ALLOC(off_q_tiles, q_tile_size); - VTCM_LAYOUT_ALLOC(off_o_tiles[0], o_tile_size); - VTCM_LAYOUT_ALLOC(off_o_tiles[1], o_tile_size); - VTCM_LAYOUT_ALLOC(off_k_fp16[0], k_dma_size); - VTCM_LAYOUT_ALLOC(off_k_fp16[1], k_dma_size); - VTCM_LAYOUT_ALLOC(off_v_fp16[0], v_dma_size); - VTCM_LAYOUT_ALLOC(off_v_fp16[1], v_dma_size); - VTCM_LAYOUT_ALLOC(off_k_tiles, k_tile_size); - VTCM_LAYOUT_ALLOC(off_v_tiles[0], v_tile_size); - VTCM_LAYOUT_ALLOC_OPTIONAL(off_v_tiles[1], v_tile_size, pipeline); - VTCM_LAYOUT_ALLOC(off_s_tiles, s_tile_size); - VTCM_LAYOUT_ALLOC(off_p_tiles, s_tile_size); - VTCM_LAYOUT_ALLOC(off_d_tiles, d_tile_size); - VTCM_LAYOUT_ALLOC(off_m_vec, col_vec_size); - VTCM_LAYOUT_ALLOC(off_l_vec, col_vec_size); - VTCM_LAYOUT_ALLOC(off_s_rowmax, col_vec_size); - VTCM_LAYOUT_ALLOC(off_p_rowsum, col_vec_size); - VTCM_LAYOUT_ALLOC(off_row_bufs, row_vec_size * 2 * n_threads); - VTCM_LAYOUT_ALLOC(off_hmx_scales_id, 256); - VTCM_LAYOUT_ALLOC(off_hmx_scales_qk, 256); - VTCM_LAYOUT_ALLOC(off_mask_buf, m_buf_size); - VTCM_LAYOUT_ALLOC(off_slopes, slopes_size); + VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size); + VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size); + VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size); + VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size); + VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size); + VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size); + VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size); + VTCM_LAYOUT_ALLOC(off, off_k_tiles, k_tile_size); + VTCM_LAYOUT_ALLOC(off, off_v_tiles[0], v_tile_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline); + VTCM_LAYOUT_ALLOC(off, off_s_tiles, s_tile_size); + VTCM_LAYOUT_ALLOC(off, off_p_tiles, s_tile_size); + VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size); + VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size); + VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size); + VTCM_LAYOUT_ALLOC(off, off_s_rowmax, col_vec_size); + VTCM_LAYOUT_ALLOC(off, off_p_rowsum, col_vec_size); + VTCM_LAYOUT_ALLOC(off, off_row_bufs, row_vec_size * 2 * n_threads); + VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256); + VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256); + VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size); + VTCM_LAYOUT_ALLOC(off, off_slopes, slopes_size); L->s_tile_bytes = s_tile_size; L->d_tile_bytes = d_tile_size; diff --git a/ggml/src/ggml-hexagon/htp/htp-vtcm.h b/ggml/src/ggml-hexagon/htp/htp-vtcm.h index f4b2239a2ebb..a1e90575fb1c 100644 --- a/ggml/src/ggml-hexagon/htp/htp-vtcm.h +++ b/ggml/src/ggml-hexagon/htp/htp-vtcm.h @@ -10,8 +10,8 @@ static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) { return p; } -#define VTCM_LAYOUT_ALLOC(field, sz) do { (L)->field = off; off += (sz); } while (0) -#define VTCM_LAYOUT_ALLOC_OPTIONAL(field, sz, cond) do { if (cond) { VTCM_LAYOUT_ALLOC(field, sz); } else { (L)->field = 0; } } while (0) +#define VTCM_LAYOUT_ALLOC(off, field, sz) do { (L)->field = (off); (off) += (sz); } while (0) +#define VTCM_LAYOUT_ALLOC_OPTIONAL(off, field, sz, cond) do { if (cond) { VTCM_LAYOUT_ALLOC(off, field, sz); } else { (L)->field = 0; } } while (0) #define VTCM_LAYOUT_PTR(type, base, offset) ((type *)((uint8_t *)(base) + (offset))) #define VTCM_LAYOUT_PTR_OPTIONAL(type, base, offset, cond) ((cond) ? VTCM_LAYOUT_PTR(type, base, offset) : NULL) diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index f56dfd33d987..55d6364f216b 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -1950,6 +1950,9 @@ typedef struct { struct htp_thread_trace * traces; struct htp_context * ctx; float * vtcm_f32_act; + size_t vtcm_f32_act_bytes_per_thread; + uint32_t dma_step_rows; + uint32_t dma_step_rows_shift; } activation_transfer_task_state_t; static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( @@ -1961,12 +1964,14 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( uint32_t k_stride, uint32_t k_valid, float *thread_f32_act, - struct htp_thread_trace *tr) { + struct htp_thread_trace *tr, + uint32_t dma_step_rows, + uint32_t dma_step_rows_shift) { - const uint32_t R = HTP_MM_DMA_ACT_ROWS_PER_STEP; + const uint32_t R = dma_step_rows; const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); - const uint32_t n_steps = n_rows_padded / R; + const uint32_t n_steps = n_rows_padded >> dma_step_rows_shift; // Push step 0 if (n_steps > 0 && n_rows > 0) { @@ -1985,31 +1990,28 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } - for (uint32_t s = 0; s < n_steps; ++s) { - uint32_t r = R * s; - float *curr_buf = NULL; + uint32_t r = s << dma_step_rows_shift; + float *curr_buf = thread_f32_act; if (r < n_rows) { curr_buf = (float *) dma_queue_pop(dma_q).dst; } htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r); - transfer_activation_row_pair_fp32_to_fp16( - vtcm_dst, - curr_buf, - curr_buf + k_block, - r, - k_block, - k_valid, - (r < n_rows), - ((r + 1) < n_rows) - ); + for (uint32_t p = 0; p < (R >> 1); ++p) { + uint32_t row_idx = r + (p << 1); + float *pair_buf = curr_buf + (p << 1) * k_block; + bool r0_valid = ((row_idx + 0) < n_rows); + bool r1_valid = ((row_idx + 1) < n_rows); + + transfer_activation_row_pair_fp32_to_fp16(vtcm_dst, pair_buf, pair_buf + k_block, row_idx, k_block, k_valid, r0_valid, r1_valid); + } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r); // Push step s + 2 uint32_t next_s = s + 2; - uint32_t next_r = R * next_s; + uint32_t next_r = next_s << dma_step_rows_shift; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride; @@ -2032,9 +2034,9 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, const float *src = st->src + chunk_idx * st->k_stride; if (st->vtcm_f32_act) { - float *thread_f32_act = st->vtcm_f32_act + i * HTP_MM_DMA_ACT_MULTIPLIER * st->k_block; + float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * st->vtcm_f32_act_bytes_per_thread); transfer_activation_chunk_fp32_to_fp16_dma_pipelined( - st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr + st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift ); } else { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); @@ -2217,30 +2219,62 @@ static void transfer_activation_chunk_threaded( int k_stride, int n_threads, int k_valid, - float *vtcm_f32_act) { + float *vtcm_f32_act, + size_t vtcm_f32_act_bytes) { + if (n_rows <= 0) { + return; + } + assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0 && k_stride % HTP_MM_HMX_TILE_N_COLS == 0); size_t n_tot_chunks = n_rows; size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : 32; // must be multiple of 32 to ensure correct destination address + uint32_t dma_step_rows = 2; + uint32_t dma_step_rows_shift = 1; + if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) { + size_t thread_scratch_elements = vtcm_f32_act_bytes / (n_threads * sizeof(float)); + size_t dma_step_rows_max = (thread_scratch_elements / 2) / k_block; + if (dma_step_rows_max >= 32) { + dma_step_rows = 32; + dma_step_rows_shift = 5; + } else if (dma_step_rows_max >= 16) { + dma_step_rows = 16; + dma_step_rows_shift = 4; + } else if (dma_step_rows_max >= 8) { + dma_step_rows = 8; + dma_step_rows_shift = 3; + } else if (dma_step_rows_max >= 4) { + dma_step_rows = 4; + dma_step_rows_shift = 2; + } else { + dma_step_rows = 2; + dma_step_rows_shift = 1; + } + } + activation_transfer_task_state_t state; - state.n_tasks = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task; - state.n_tot_chunks = n_tot_chunks; - state.n_chunks_per_task = n_chunks_per_task; - state.dst = dst; - state.src = src; - state.k_block = k_block; - state.k_stride = k_stride; - state.k_valid = k_valid; - state.traces = ctx->trace; - state.ctx = ctx; - state.vtcm_f32_act = vtcm_f32_act; + state.n_tasks = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task; + state.n_tot_chunks = n_tot_chunks; + state.n_chunks_per_task = n_chunks_per_task; + state.dst = dst; + state.src = src; + state.k_block = k_block; + state.k_stride = k_stride; + state.k_valid = k_valid; + state.traces = ctx->trace; + state.ctx = ctx; + state.vtcm_f32_act = vtcm_f32_act; + + int active_threads = hex_smin(n_threads, (int)state.n_tasks); + state.vtcm_f32_act_bytes_per_thread = (vtcm_f32_act_bytes / active_threads) & ~127u; + state.dma_step_rows = dma_step_rows; + state.dma_step_rows_shift = dma_step_rows_shift; if (state.n_tasks == 1 || n_threads == 1) { transfer_activation_chunk_worker_fn(1, 0, &state); } else { - int n_tasks = hex_smin((int) state.n_tasks, n_threads); - worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, n_tasks); + worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, active_threads); } } // --- Async HMX matmul job (for pipeline overlap) --- @@ -2378,7 +2412,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 }; void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 }; - transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act); + transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes); // 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); @@ -2458,7 +2492,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); - transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act); + transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes); for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); @@ -2629,7 +2663,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; transfer_activation_chunk_threaded(ctx, vtcm_act_g, activation_chunk, (int) n_rows, - params->k, params->act_stride, act_threads, params->k, vtcm_f32_act); + params->k, params->act_stride, act_threads, params->k, vtcm_f32_act, L.act_f32_bytes); } void *buf_curr = vtcm_scratch0; diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index 42dafd271af7..2e131bc3d025 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -356,26 +356,43 @@ static inline void htp_mm_hmx_vtcm_layout_build( const size_t activation_area_size = hex_align_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); const size_t output_area_size = hex_align_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); const size_t scratch_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t f32_scratch_size = use_dma_activation - ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0; - - VTCM_LAYOUT_ALLOC(off_weight[0], weight_area_size); - L->off_weight[1] = 0; - VTCM_LAYOUT_ALLOC(off_act, activation_area_size); - VTCM_LAYOUT_ALLOC(off_dst[0], output_area_size); - L->off_dst[1] = 0; - VTCM_LAYOUT_ALLOC(off_scratch[0], scratch_area_size); - VTCM_LAYOUT_ALLOC(off_scratch[1], scratch_area_size); - VTCM_LAYOUT_ALLOC(off_scales, 256); - VTCM_LAYOUT_ALLOC_OPTIONAL(off_act_f32, f32_scratch_size, use_dma_activation); + const size_t min_f32_size = use_dma_activation + ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128) : 0; + + // Group A: Permanent activation tiles and scales + size_t off_group_a = 0; + VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size); + VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data + + // Group B: Compute-only buffers (starts at off_group_a) + size_t off_group_b = off_group_a; + VTCM_LAYOUT_ALLOC(off_group_b, off_weight[0], weight_area_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_weight[1], weight_area_size, false); + VTCM_LAYOUT_ALLOC(off_group_b, off_dst[0], output_area_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst[1], output_area_size, false); + VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[0], scratch_area_size); + VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[1], scratch_area_size); + + const size_t group_b_size = off_group_b - off_group_a; + + // Group C: Activation prep temporary buffer (overlaps Group B, starting at off_group_a) + const size_t max_f32_size = act_threads * 64 * k * sizeof(float); + const size_t act_f32_size = use_dma_activation + ? hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128) : 0; + size_t off_group_c = off_group_a; + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_f32, act_f32_size, use_dma_activation); + + const size_t group_c_size = off_group_c - off_group_a; L->weight_area_bytes = weight_area_size; L->act_area_bytes = activation_area_size; - L->act_f32_bytes = f32_scratch_size; + L->act_f32_bytes = act_f32_size; L->output_area_bytes = output_area_size; L->scratch_bytes[0] = scratch_area_size; L->scratch_bytes[1] = scratch_area_size; L->act_head_stride = act_head_stride; + + off = off_group_a + hex_smax(group_b_size, group_c_size); } else { // HTP_MM_KERNEL_HMX_2D const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); @@ -383,7 +400,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( const size_t vec_dot_size = k * sizeof(uint16_t); const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; - const size_t act_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE); + const size_t min_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128); const size_t weight_area_size = is_quant ? hex_align_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE) : hex_align_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE); @@ -393,15 +410,29 @@ static inline void htp_mm_hmx_vtcm_layout_build( const size_t scratch0_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); const size_t scratch1_size = pipeline ? scratch0_size : 0; - VTCM_LAYOUT_ALLOC(off_weight[0], weight_area_size); - VTCM_LAYOUT_ALLOC_OPTIONAL(off_weight[1], weight_area_size, pipeline); - VTCM_LAYOUT_ALLOC(off_act, act_area_size); - VTCM_LAYOUT_ALLOC_OPTIONAL(off_act_f32, act_f32_size, true); - VTCM_LAYOUT_ALLOC(off_dst[0], output_area_size); - VTCM_LAYOUT_ALLOC(off_scratch[0], scratch0_size); - VTCM_LAYOUT_ALLOC_OPTIONAL(off_scratch[1], scratch0_size, pipeline); - VTCM_LAYOUT_ALLOC_OPTIONAL(off_dst[1], output_area_size, pipeline); - VTCM_LAYOUT_ALLOC(off_scales, 256); + // Group A: Scales and activation tiles (must not overlap with Group B or C) + size_t off_group_a = 0; + VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data + VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size); + + // Group B: Compute-only buffers (starts at off_group_a) + size_t off_group_b = off_group_a; + VTCM_LAYOUT_ALLOC(off_group_b, off_weight[0], weight_area_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_weight[1], weight_area_size, pipeline); + VTCM_LAYOUT_ALLOC(off_group_b, off_dst[0], output_area_size); + VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[0], scratch0_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_scratch[1], scratch0_size, pipeline); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst[1], output_area_size, pipeline); + + const size_t group_b_size = off_group_b - off_group_a; + + // Group C: Activation prep temporary buffer (overlaps Group B, starting at off_group_a) + const size_t max_f32_size = act_threads * 64 * k * sizeof(float); + const size_t act_f32_size = hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128); + size_t off_group_c = off_group_a; + VTCM_LAYOUT_ALLOC(off_group_c, off_act_f32, act_f32_size); + + const size_t group_c_size = off_group_c - off_group_a; L->weight_area_bytes = weight_area_size; L->act_area_bytes = act_area_size; @@ -410,6 +441,8 @@ static inline void htp_mm_hmx_vtcm_layout_build( L->scratch_bytes[0] = scratch0_size; L->scratch_bytes[1] = scratch1_size; L->act_head_stride = 0; + + off = off_group_a + hex_smax(group_b_size, group_c_size); } L->total_bytes = off; @@ -452,7 +485,7 @@ static inline void htp_mm_hvx_vtcm_layout_build( 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; - + src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); if (is_fused_qkv) { @@ -468,7 +501,7 @@ static inline void htp_mm_hvx_vtcm_layout_build( size_t flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - + if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128); } else { @@ -580,11 +613,11 @@ static inline void htp_mm_hvx_vtcm_layout_build( } size_t off = 0; - VTCM_LAYOUT_ALLOC(off_src1, src1_sz); - VTCM_LAYOUT_ALLOC(off_src0, src0_sz); - VTCM_LAYOUT_ALLOC(off_src2, src2_sz); - VTCM_LAYOUT_ALLOC(off_src3, src3_sz); - VTCM_LAYOUT_ALLOC(off_dst, dst_sz); + VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz); + VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz); + VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz); + VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz); + VTCM_LAYOUT_ALLOC(off, off_dst, dst_sz); L->src0_bytes = src0_sz; L->src1_bytes = src1_sz; From b74d67160bf6188027a22e04faf53ee5a0b41669 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Sat, 4 Jul 2026 13:16:48 -0700 Subject: [PATCH 15/32] hmx-mm: reduce the number of act fetch tows to 4 for now, going larger doesnt help here --- ggml/src/ggml-hexagon/htp/matmul-ops.c | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 55d6364f216b..8c17a1db42e4 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2235,16 +2235,7 @@ static void transfer_activation_chunk_threaded( if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) { size_t thread_scratch_elements = vtcm_f32_act_bytes / (n_threads * sizeof(float)); size_t dma_step_rows_max = (thread_scratch_elements / 2) / k_block; - if (dma_step_rows_max >= 32) { - dma_step_rows = 32; - dma_step_rows_shift = 5; - } else if (dma_step_rows_max >= 16) { - dma_step_rows = 16; - dma_step_rows_shift = 4; - } else if (dma_step_rows_max >= 8) { - dma_step_rows = 8; - dma_step_rows_shift = 3; - } else if (dma_step_rows_max >= 4) { + if (dma_step_rows_max >= 4) { dma_step_rows = 4; dma_step_rows_shift = 2; } else { From fc3088c567a884c87f185db0968ab8c65498d462 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 01:12:08 -0700 Subject: [PATCH 16/32] hex-hmx: always use hmx-queue in all modes --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 131 ++++++--------------- ggml/src/ggml-hexagon/htp/matmul-ops.c | 35 +++--- 2 files changed, 53 insertions(+), 113 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 3adbd8b4d579..9cb0292be9f1 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1657,10 +1657,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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; @@ -1870,29 +1866,17 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); - asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)factx.vtcm_hmx_scales_qk)); - const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; - for (size_t r = 0; r < n_row_tiles; ++r) { - const __fp16 * row_tiles = q_base + r * dot_stride; - const __fp16 * col_tiles = k_base; - __fp16 * out_tile = s_base + r * n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; - - for (size_t c = 0; c < n_col_tiles; ++c) { - hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles); - col_tiles += dot_stride; - out_tile += HMX_FP16_TILE_N_ELMS; - } - } - htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); + 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 = (size_t) (DK / 32); + qk_job.n_tiles_per_bc = n_tiles_per_bc; + qk_job.hmx_scales = factx.vtcm_hmx_scales_qk; + + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job)); + hmx_queue_pop(ctx->hmx_queue); } // ---- Phase 3: softmax + build_D ---- @@ -1930,35 +1914,21 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); - asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)factx.vtcm_hmx_scales_id)); - const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; - const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS; - for (size_t r = 0; r < n_row_tiles; ++r) { - const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; - const __fp16 * p_tile_in = p_base + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS; - const __fp16 * o_rc = op_base + r * HMX_FP16_TILE_N_ELMS; - const __fp16 * v_tile_in = v_base; - __fp16 * o_tile_out = oc_base + r * HMX_FP16_TILE_N_ELMS; - - for (size_t c = 0; c < DV_tiles; ++c) { - hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles); - o_rc += o_stride; - v_tile_in += v_stride; - o_tile_out += o_stride; - } - } - htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); + 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[0]; + 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 = n_col_tiles; + 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(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job)); + hmx_queue_pop(ctx->hmx_queue); + hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev); } @@ -1972,40 +1942,15 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); - asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)factx.vtcm_hmx_scales_id)); - const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS; - for (size_t r = 0; r < n_row_tiles; ++r) { - const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS; - const __fp16 * o_rc = op_base + r * HMX_FP16_TILE_N_ELMS; - __fp16 * o_out = oc_base + r * DV_tiles * HMX_FP16_TILE_N_ELMS; - - for (size_t c = 0; c < DV_tiles; ++c) { - hmx_fa_o_norm_tile(d_diag, o_rc, o_out); - o_rc += o_stride; - o_out += HMX_FP16_TILE_N_ELMS; - } - } - htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start); - } + on_job.o_curr = o_tile_curr; + on_job.o_prev = o_tile_prev; + on_job.d_tiles = factx.vtcm_d_tiles; + on_job.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); } // ---- Store O block ---- @@ -2014,11 +1959,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { } } - if (factx.pipeline) { - hmx_queue_suspend(ctx->hmx_queue); - } else { - HAP_compute_res_hmx_unlock(ctx->vtcm_rctx); - } + hmx_queue_suspend(ctx->hmx_queue); return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 8c17a1db42e4..f16601da535b 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2476,10 +2476,9 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, } } } - hmx_queue_suspend(ctx->hmx_queue); } else { // --- Synchronous Un-pipelined loop (m <= 32 or fallback) --- - HAP_compute_res_hmx_lock(ctx->vtcm_rctx); + hmx_matmul_job_t job; for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); @@ -2502,8 +2501,10 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); - // C: HMX Compute (Synchronous) - core_dot_chunk_fp16(vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + // C: HMX Compute (Queue-based) + hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); + hmx_queue_pop(ctx->hmx_queue); // D: Output Store float *output_chunk = dst + (mr * dst_stride + nc); @@ -2514,9 +2515,10 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, } } } - HAP_compute_res_hmx_unlock(ctx->vtcm_rctx); } + hmx_queue_suspend(ctx->hmx_queue); + return 0; } @@ -2634,7 +2636,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ const size_t fp16_row_bytes = (size_t) params->k * sizeof(__fp16); const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16); - HAP_compute_res_hmx_lock(ctx->vtcm_rctx); + hmx_matmul_job_t job; for (int b3 = 0; b3 < params->ne13; ++b3) { for (int b2_base = 0; b2_base < params->ne12; b2_base += group_size) { @@ -2688,13 +2690,12 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ // Reuse the interleaved weight for every q_head in this GQA group for (int g = 0; g < group_size; ++g) { - struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, g); { const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; - core_dot_chunk_fp16(vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32); + hmx_matmul_job_init(&job, vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32); + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); + hmx_queue_pop(ctx->hmx_queue); } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, g); { float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc; @@ -2711,8 +2712,6 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ } } - HAP_compute_res_hmx_unlock(ctx->vtcm_rctx); - return 0; } @@ -2904,7 +2903,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); - HAP_compute_res_hmx_lock(ctx->vtcm_rctx); + hmx_matmul_job_t job; for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(m_padded - mr, m_chunk_n_rows); @@ -2929,10 +2928,9 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads ); - struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, nc); - core_dot_chunk_fp16(vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); - htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, nc); + hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); + hmx_queue_pop(ctx->hmx_queue); transfer_output_chunk_scattered_threaded( ctx, dst + nc, vtcm_output, (int) mr, (int) n_rows, (int) n_cols, @@ -2940,7 +2938,8 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, } } - HAP_compute_res_hmx_unlock(ctx->vtcm_rctx); + hmx_queue_suspend(ctx->hmx_queue); + return 0; } From 29700a84255d77df2e7056e2fbf525e7d8c3838b Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 09:36:20 -0700 Subject: [PATCH 17/32] hmx-mm: update comments and minor formatting --- ggml/src/ggml-hexagon/htp/matmul-ops.c | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index f16601da535b..b9f65bfd8bf2 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2477,7 +2477,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, } } } else { - // --- Synchronous Un-pipelined loop (m <= 32 or fallback) --- + // --- Synchronous loop (m <= 32 or fallback) --- hmx_matmul_job_t job; for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); @@ -2491,8 +2491,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // A: Weight DMA (Synchronous) 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_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) From 5e62db6c65e8879ac2f09142d9492bf102cb0219 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 10:14:07 -0700 Subject: [PATCH 18/32] hmx-mm: further improve synchro fallback path to prefetch the weights earlier --- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 2 +- ggml/src/ggml-hexagon/htp/matmul-ops.c | 43 ++++++++++++++++++++---- 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 6ebb4b4888ef..ca22d8f719c3 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -21,7 +21,6 @@ add_library(${HTP_LIB} SHARED hex-dma.c hmx-queue.c matmul-ops.c - flash-attn-ops.c gated-delta-net-ops.c binary-ops.c unary-ops.c @@ -41,6 +40,7 @@ add_library(${HTP_LIB} SHARED diag-ops.c solve-tri-ops.c pad-ops.c + flash-attn-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index b9f65bfd8bf2..06d0f04bbd6a 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2484,14 +2484,19 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes); + // A0: Pre-fetch the first weight chunk (nc = 0) + if (n > 0) { + const size_t n_cols = hex_smin(n, n_chunk_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), dma_dst_stride, dma_src_stride, dma_width_bytes, height); + } + for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); - // A: Weight DMA (Synchronous) - 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); + // A: Wait for weight DMA dma_queue_pop(ctx->dma[0]); // B: Weight Dequantize (Threaded) @@ -2500,6 +2505,14 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); + // Start weight DMA for the next chunk early + const size_t nc_next = nc + n_chunk_n_cols; + if (nc_next < n) { + const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); + const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + } + // C: HMX Compute (Queue-based) hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); @@ -2912,25 +2925,43 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, ctx, vtcm_f16_act, activation, (int) mr, (int) n_rows, k, matrix_rows, cur_a, mapping_stride, ne11, act_nb1, act_nb2, cne1, n_threads, k_valid); + // A0: Pre-fetch the first weight chunk (nc = 0) + if (n > 0) { + const size_t n_cols = hex_smin((size_t) n, n_chunk_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), + dma_dst_stride, dma_src_stride, dma_width_bytes, height); + } + for (size_t nc = 0; nc < (size_t) n; nc += n_chunk_n_cols) { 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); - 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); + // A: Wait for weight DMA dma_queue_pop(ctx->dma[0]); + // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( ctx, vtcm_scratch0, vtcm_weight, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads ); + // Start weight DMA for the next chunk early + const size_t nc_next = nc + n_chunk_n_cols; + if (nc_next < (size_t) n) { + const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols); + const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc_next * weight_stride), + dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + } + + // C: HMX Compute (Queue-based) hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); hmx_queue_pop(ctx->hmx_queue); + // D: Output Store transfer_output_chunk_scattered_threaded( ctx, dst + nc, vtcm_output, (int) mr, (int) n_rows, (int) n_cols, matrix_rows, cur_a, mapping_stride, dst_nb1, dst_nb2, cne1, n_threads); From 71a55d8a3cda76c1a3b7f2b57eb42da00f8f8d2e Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 10:46:40 -0700 Subject: [PATCH 19/32] hex-fa: further pipeline improvements (earlier prefetch) --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 28 ++++++++++++---------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 9cb0292be9f1..83e92b48f6d0 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1728,6 +1728,17 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { } } + // Prefetch next KV block early + if (kv_blk + 1 < factx.n_kv_blocks) { + const uint32_t prefetch_start = (kv_blk + 1) * Bc; + const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start); + const size_t prefetch_buf = 1 - buf_idx; + const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3]; + dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows); + const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3]; + dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows); + } + // Wait for current KV DMA dma_queue_pop(dma); // K dma_queue_pop(dma); // V @@ -1759,16 +1770,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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) { @@ -1843,8 +1844,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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)) { @@ -1854,6 +1853,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); @@ -1863,6 +1863,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); } + + dma_queue_pop(dma); // K + dma_queue_pop(dma); // V + fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start); { From c8ee54f5a8fea89bab6018dd96b81e55d7e9885c Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 12:10:36 -0700 Subject: [PATCH 20/32] hmx-mm: cleanup dma pipelines to use dst cached in the queue --- ggml/src/ggml-hexagon/htp/matmul-ops.c | 122 ++++++++++++------------- 1 file changed, 61 insertions(+), 61 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 06d0f04bbd6a..3453f457c792 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2388,7 +2388,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 - FARF(HIGH, "hmx-mm-2d-precomputed: standard : m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu", + FARF(HIGH, "hmx-mm-2d: m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu", m, k, n, weight_type, m_chunk_n_rows, n_chunk_n_cols, vtcm_used, vtcm_budget); int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols); @@ -2406,74 +2406,72 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes); // 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); + const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); 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); + const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols); 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 - dma_queue_pop(ctx->dma[0]); - dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_weight_bufs[0], vtcm_weight_raw[0], - n_cols_A0, k, row_stride, weight_type, - n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); - - hmx_matmul_job_init(&job_slots[0], (__fp16 *) vtcm_output_bufs[0], (__fp16 *) vtcm_f16_act, - (__fp16 *) vtcm_weight_bufs[0], vtcm_scales, - hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS), - hmx_ceil_div(n_cols_A0, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS); - hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[0])); - - // Main loop: pop/dequantize A_{i+1} -> push A_{i+2} -> submit C_{i+1} -> wait C_i and store D_i + // Main loop: pop A_i -> dequantize A_i -> push A_{i+2} -> submit C_i -> wait C_{i-1} and store D_{i-1} for (int i = 0; i < n_chunk_cnt; ++i) { const size_t nc = i * n_chunk_n_cols; - const size_t nc_p1 = nc + 1 * n_chunk_n_cols; const size_t nc_p2 = nc + 2 * n_chunk_n_cols; const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); - const size_t n_cols_p1 = hex_smin(n - nc_p1, n_chunk_n_cols); const size_t n_cols_p2 = hex_smin(n - nc_p2, n_chunk_n_cols); - // 1. pop A_{i+1} and dequantize it (if i+1 < n_chunk_cnt) - if (i + 1 < n_chunk_cnt) { - dma_queue_pop(ctx->dma[0]); - dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_weight_bufs[(i + 1) % 2], vtcm_weight_raw[(i + 1) % 2], - n_cols_p1, k, row_stride, weight_type, - n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); - } + // 1. pop A_i + void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; - // 2. push A_{i+2} (if i+2 < n_chunk_cnt) + // 2. dequantize A_i + dequantize_tiled_weight_chunk_to_fp16_tiles( + ctx, vtcm_weight_bufs[i % 2], curr_raw, + n_cols, k, row_stride, weight_type, + n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); + + // 3. push A_{i+2} (if i+2 < n_chunk_cnt) if (i + 2 < n_chunk_cnt) { 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_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, 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) - if (i + 1 < n_chunk_cnt) { - hmx_matmul_job_init(&job_slots[(i + 1) % 2], (__fp16 *) vtcm_output_bufs[(i + 1) % 2], - (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[(i + 1) % 2], - vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS), - hmx_ceil_div(n_cols_p1, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS); - hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[(i + 1) % 2])); + // 4. submit C_i + hmx_matmul_job_init(&job_slots[i % 2], (__fp16 *) vtcm_output_bufs[i % 2], + (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], + vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS), + hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS); + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[i % 2])); + + // 5. wait C_{i-1} and store D_{i-1} (multi-thread HVX, parallel with C_i) + if (i > 0) { + hmx_queue_pop(ctx->hmx_queue); + const size_t nc_prev = (i - 1) * n_chunk_n_cols; + const size_t n_cols_prev = hex_smin(n - nc_prev, n_chunk_n_cols); + float *output_chunk = dst + (mr * dst_stride + nc_prev); + const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_prev) : NULL; + int chunk_dst_cols = dst_cols - (int)nc_prev; + if (chunk_dst_cols > 0) { + transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(i - 1) % 2], n_rows, n_cols_prev, dst_stride, src2_stride, chunk_dst_cols, n_threads); + } } + } - // 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, src2_chunk, vtcm_output_bufs[i % 2], n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads); - } + // Epilogue: wait C_{last} and store D_{last} + hmx_queue_pop(ctx->hmx_queue); + const size_t nc_last = (n_chunk_cnt - 1) * n_chunk_n_cols; + const size_t n_cols_last = hex_smin(n - nc_last, n_chunk_n_cols); + float *output_chunk = dst + (mr * dst_stride + nc_last); + const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_last) : NULL; + int chunk_dst_cols = dst_cols - (int)nc_last; + if (chunk_dst_cols > 0) { + transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(n_chunk_cnt - 1) % 2], n_rows, n_cols_last, dst_stride, src2_stride, chunk_dst_cols, n_threads); } } } else { @@ -2497,11 +2495,11 @@ 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: Wait for weight DMA - dma_queue_pop(ctx->dma[0]); + void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_scratch0, vtcm_weight_raw[0], + ctx, vtcm_scratch0, curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); @@ -2510,7 +2508,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, if (nc_next < n) { const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } // C: HMX Compute (Queue-based) @@ -2671,33 +2669,35 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ params->k, params->act_stride, act_threads, params->k, vtcm_f32_act, L.act_f32_bytes); } - void *buf_curr = vtcm_scratch0; - void *buf_next = vtcm_scratch1; - + // Prologue: Push A0 and A1 (if exists) { const size_t n_cols_first = hex_smin((size_t) params->n, n_chunk_n_cols); - dma_queue_push(ctx->dma[0], dma_make_ptr(buf_curr, weight_group), + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch0, weight_group), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_first); } + if (n_chunk_n_cols < (size_t) params->n) { + const size_t n_cols_second = hex_smin((size_t) params->n - n_chunk_n_cols, n_chunk_n_cols); + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch1, weight_group + params->weight_stride), + fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_second); + } for (size_t nc = 0; nc < (size_t) params->n; nc += n_chunk_n_cols) { - const size_t n_cols = hex_smin((size_t) params->n - nc, n_chunk_n_cols); + const size_t n_cols = hex_smin((size_t) params->n - nc, n_chunk_n_cols); const size_t n_col_tiles = hmx_ceil_div((int) n_cols, HTP_MM_HMX_TILE_N_COLS); { - dma_queue_pop(ctx->dma[0]); + void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; - const size_t nc_next = nc + n_chunk_n_cols; + hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) curr_raw, n_cols, params->k, params->k, 0, n_cols); + + const size_t nc_next = nc + n_chunk_n_cols * 2; if (nc_next < (size_t) params->n) { const size_t n_cols_next = hex_smin((size_t) params->n - nc_next, n_chunk_n_cols); const __fp16 *next_weight_chunk = weight_group + nc_next * params->weight_stride; - dma_queue_push(ctx->dma[0], dma_make_ptr(buf_next, next_weight_chunk), + dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, next_weight_chunk), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_next); } - - hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) buf_curr, n_cols, params->k, params->k, 0, n_cols); - hex_swap_ptr(&buf_curr, &buf_next); } // Reuse the interleaved weight for every q_head in this GQA group @@ -2938,11 +2938,11 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); // A: Wait for weight DMA - dma_queue_pop(ctx->dma[0]); + void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_scratch0, vtcm_weight, + ctx, vtcm_scratch0, curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads ); @@ -2952,7 +2952,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, if (nc_next < (size_t) n) { const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc_next * weight_stride), + dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } From a35aad116793685fb0ce64f00fe44f6eb419f9aa Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 13:54:47 -0700 Subject: [PATCH 21/32] hmx-fa: minor cleanup and opts for fa dma pipelines --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 40 +++++++++++++--------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 83e92b48f6d0..28f8758bc0cd 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -463,7 +463,7 @@ typedef struct { struct hmx_fa_context * factx; uint32_t kv_rows; size_t src_stride; - size_t buf_idx; + void * curr_k; uint32_t kv_start; uint32_t rows_per_t; } fa_k_int_args_t; @@ -483,19 +483,19 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) 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, + hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK, args->src_stride, start, end); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); } -static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, size_t buf_idx, uint32_t kv_start) { +static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) { worker_pool_context_t wp = factx->octx->ctx->worker_pool; 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 }; + fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t }; if (n > 1) { worker_pool_run_func(wp, fa_k_interleave_thread, &args, n); } else { @@ -507,6 +507,7 @@ typedef struct { struct hmx_fa_context * factx; uint32_t kv_rows; size_t src_stride; + void * curr_v; size_t buf_idx; size_t n_col_tiles; uint32_t kv_start; @@ -530,7 +531,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) 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, + hmx_interleave_cols_to_tiles(v_tiles_dest, (const __fp16 *) args->curr_v, 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)); } @@ -538,6 +539,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) static void fa_phase_v_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, + void * curr_v, size_t buf_idx, size_t n_col_tiles, uint32_t kv_start) { @@ -547,7 +549,7 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx, 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 }; + fa_v_int_args_t args = { factx, kv_rows, src_stride, curr_v, 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 { @@ -1739,10 +1741,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); } - // 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; @@ -1758,7 +1756,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); + + // Wait for current K DMA and interleave + void * curr_k = dma_queue_pop(dma).dst; + fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start); // ---- Phase 2: qk_dot ---- qk_job.q_tiles = factx.vtcm_q_tiles; @@ -1770,7 +1771,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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)); - fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc, kv_start); + + // Wait for current V DMA and interleave + void * curr_v = dma_queue_pop(dma).dst; + fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, buf_idx, n_tiles_per_bc, kv_start); if (kv_blk > 0) { hmx_queue_pop(hmx_q); @@ -1864,10 +1868,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); } - dma_queue_pop(dma); // K - dma_queue_pop(dma); // V - - fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start); + // Wait for current K DMA and interleave + void * curr_k = dma_queue_pop(dma).dst; + fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start); { qk_job.q_tiles = factx.vtcm_q_tiles; @@ -1883,6 +1886,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { hmx_queue_pop(ctx->hmx_queue); } + // Wait for current V DMA and interleave + void * curr_v = dma_queue_pop(dma).dst; + fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, buf_idx, n_tiles_per_bc, kv_start); + // ---- Phase 3: softmax + build_D ---- __fp16 * current_mask_vtcm = NULL; if (mask) { @@ -1915,7 +1922,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); { ou_job.o_curr = o_tile_curr; From b0011164739f99f84a8081a7d663d65a80a8fb02 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 17:26:36 -0700 Subject: [PATCH 22/32] hmx-fa: optimize q-prep stage with dma and unrolling --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 170 +++++----- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 344 +++++++++++++++++++++ 2 files changed, 443 insertions(+), 71 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 28f8758bc0cd..ae4a39c57a47 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -565,6 +566,9 @@ typedef struct { uint32_t ib3; size_t n_rows_g; size_t rows_per_t; + size_t n_rows_q; + bool q_transposed; + atomic_uint barrier; } fa_q_load_args_t; static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { @@ -591,7 +595,6 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { 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); @@ -645,72 +648,70 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { 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; + if (start < factx->g_br) { + 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); + assert(factx->DK == factx->DV); - 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; + const size_t o_tile_bytes = hex_align_up(factx->g_br * factx->DV * sizeof(__fp16), 4096); + const bool use_vtcm_temp = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2)); - 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 (use_vtcm_temp) { + if (factx->is_q_fp32) { + char * temp_q_vtcm = (char *) factx->vtcm_o_tiles[0]; + const uint32_t d_limit = DK / 32; + const size_t valid_end = hex_smin(end, n_rows_g); - 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; + if (d_limit == 2) { + hmx_fa_q_prep_fp32_d2(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + } else if (d_limit == 4) { + hmx_fa_q_prep_fp32_d4(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + } else { + hmx_fa_q_prep_fp32_generic(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); + } + } else { + char * temp_q_vtcm = (char *) factx->vtcm_o_tiles[0]; + const uint32_t d_limit = DK / 64; + const size_t valid_end = hex_smin(end, n_rows_g); + + if (d_limit == 1) { + hmx_fa_q_prep_fp16_d1(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + } else if (d_limit == 2) { + hmx_fa_q_prep_fp16_d2(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + } else { + hmx_fa_q_prep_fp16_generic(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); + } } } 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); + // Fallback: Original direct-from-DDR/L2 path + hmx_fa_q_prep_fallback(factx->vtcm_q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3], + q_start, kv_head, ib3, start, end, n_rows_g, G, DK, factx->is_q_fp32, &factx->div_G); + } + } - __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; - HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; - HVX_Vector * pv_out1 = pv_out0 + 16; + // Synchronize threads before zeroing out vtcm_o_tiles[0] to prevent race condition + if (n > 1) { + atomic_fetch_sub(&args->barrier, 1); + while (atomic_load(&args->barrier) > 0) { + // spin wait + } + } - *pv_out0 = Q6_V_lo_W(vp); - *pv_out1 = Q6_V_hi_W(vp); - } + // Zero out vtcm_o_tiles[0] as it was used as temp_q_vtcm + { + const uint32_t g_br = factx->g_br; + const uint32_t DV = factx->DV; + const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096); + 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 *) factx->vtcm_o_tiles[0] + o_start, 0, o_end - o_start); } } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start)); @@ -728,7 +729,18 @@ static void fa_phase_q_load(struct hmx_fa_context * factx, 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 }; + const uint32_t n_rows_q = hex_smin(factx->Br, factx->neq1 - q_start); + fa_q_load_args_t args; + args.factx = factx; + args.q = q; + args.q_start = q_start; + args.kv_head = kv_head; + args.ib3 = ib3; + args.n_rows_g = n_rows_g; + args.rows_per_t = rows_per_t; + args.n_rows_q = n_rows_q; + args.q_transposed = q->nb[1] < q->nb[2]; + atomic_init(&args.barrier, n); if (n > 1) { worker_pool_run_func(wp, fa_q_load_thread, &args, n); } else { @@ -1486,7 +1498,7 @@ static void fa_push_mask_dma_gqa( uint32_t G, uint32_t m_line_bytes, uint32_t kv_rows, - uint32_t n_q_rows, + uint32_t n_rows_q, struct hmx_fa_context * factx ) { for (uint32_t g = 0; g < G; ++g) { @@ -1495,7 +1507,7 @@ static void fa_push_mask_dma_gqa( 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); + dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); } } @@ -1669,8 +1681,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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 uint32_t n_rows_q = hex_smin(Br, neq1 - q_start); + const size_t n_rows_g = n_rows_q * 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; @@ -1680,17 +1692,33 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const uint32_t iv2 = kv_head; const uint32_t iv3 = ib3 / (neq3 / v->ne[3]); - // Prefetch first KV block + // 1. Push Q DMA (if vtcm_temp is used) + const size_t o_tile_bytes = hex_align_up(factx.g_br * factx.DV * sizeof(__fp16), 4096); + const bool use_vtcm_temp = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2)); + if (use_vtcm_temp) { + const bool q_transposed = q->nb[1] < q->nb[2]; + const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3]; + const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16); + const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size; + const size_t src_stride = q_transposed ? q->nb[2] : q->nb[1]; + const size_t n_rows = q_transposed ? factx.G : n_rows_q; + dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows); + } + + // 2. 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); + 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); + 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); + } + + // 3. Pop Q DMA (blocks until Q is loaded) + if (use_vtcm_temp) { + dma_queue_pop(dma); } // ---- Load Q block & Initialize per-block state ---- @@ -1724,9 +1752,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); + dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); } else { - fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx); + fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx); } } @@ -1852,9 +1880,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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); + dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); } else { - fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx); + fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx); } } diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index f866dc1a166f..b26d2abbb460 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -6,6 +6,7 @@ #include #include "hvx-utils.h" #include "hmx-utils.h" +#include "hex-fastdiv.h" // HMX-specific parameters, offsets and inner kernels for Flash Attention @@ -208,4 +209,347 @@ static inline void hmx_fa_o_norm_tile( ); } +static inline void hmx_fa_q_prep_fp32_d2( + __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, + size_t start, size_t end, size_t valid_end, + size_t DK, size_t G, size_t n_rows_q, + const struct fastdiv_values * div_G, bool q_transposed +) { + for (size_t r = start; r < end; r += 2) { + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (r >= valid_end) { + ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); + ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); + continue; + } + + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0); + const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float)); + const HVX_Vector * pv_in1 = (r + 1 < valid_end) + ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float)) + : NULL; + + { + HVX_Vector v0 = pv_in0[0]; + HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero(); + HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1); + ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf; + } + { + HVX_Vector v0 = pv_in0[1]; + HVX_Vector v1 = pv_in1 ? pv_in1[1] : Q6_V_vzero(); + HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1); + ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf; + } + } +} + +static inline void hmx_fa_q_prep_fp32_d4( + __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, + size_t start, size_t end, size_t valid_end, + size_t DK, size_t G, size_t n_rows_q, + const struct fastdiv_values * div_G, bool q_transposed +) { + for (size_t r = start; r < end; r += 2) { + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (r >= valid_end) { + for (uint32_t d = 0; d < 4; ++d) { + ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); + } + continue; + } + + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0); + const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float)); + const HVX_Vector * pv_in1 = (r + 1 < valid_end) + ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float)) + : NULL; + + for (uint32_t d = 0; d < 4; ++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_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf; + } + } +} + +static inline void hmx_fa_q_prep_fp32_generic( + __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, + size_t start, size_t end, size_t valid_end, + size_t DK, size_t G, size_t n_rows_q, + const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed +) { + for (size_t r = start; r < end; r += 2) { + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (r >= valid_end) { + for (uint32_t d = 0; d < d_limit; ++d) { + ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); + } + continue; + } + + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0); + const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float)); + const HVX_Vector * pv_in1 = (r + 1 < valid_end) + ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float)) + : NULL; + + for (uint32_t d = 0; d < d_limit; ++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; + } + } +} + +static inline void hmx_fa_q_prep_fp16_d1( + __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, + size_t start, size_t end, size_t valid_end, + size_t DK, size_t G, size_t n_rows_q, + const struct fastdiv_values * div_G, bool q_transposed +) { + for (size_t r = start; r < end; r += 2) { + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (r >= valid_end) { + __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + *pv_out0 = Q6_V_vzero(); + *pv_out1 = Q6_V_vzero(); + continue; + } + + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0); + const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16)); + const HVX_Vector * pv_in1 = (r + 1 < valid_end) + ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16)) + : NULL; + + HVX_Vector v0 = pv_in0[0]; + HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero(); + HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2); + + __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + + *pv_out0 = Q6_V_lo_W(vp); + *pv_out1 = Q6_V_hi_W(vp); + } +} + +static inline void hmx_fa_q_prep_fp16_d2( + __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, + size_t start, size_t end, size_t valid_end, + size_t DK, size_t G, size_t n_rows_q, + const struct fastdiv_values * div_G, bool q_transposed +) { + for (size_t r = start; r < end; r += 2) { + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (r >= valid_end) { + for (uint32_t d = 0; d < 2; ++d) { + __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + *pv_out0 = Q6_V_vzero(); + *pv_out1 = Q6_V_vzero(); + } + continue; + } + + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0); + const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16)); + const HVX_Vector * pv_in1 = (r + 1 < valid_end) + ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16)) + : NULL; + + { + HVX_Vector v0 = pv_in0[0]; + HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero(); + HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2); + + __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + + *pv_out0 = Q6_V_lo_W(vp); + *pv_out1 = Q6_V_hi_W(vp); + } + { + HVX_Vector v0 = pv_in0[1]; + HVX_Vector v1 = pv_in1 ? pv_in1[1] : Q6_V_vzero(); + HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2); + + __fp16 * out_dtile = out_base + 1 * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + + *pv_out0 = Q6_V_lo_W(vp); + *pv_out1 = Q6_V_hi_W(vp); + } + } +} + +static inline void hmx_fa_q_prep_fp16_generic( + __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, + size_t start, size_t end, size_t valid_end, + size_t DK, size_t G, size_t n_rows_q, + const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed +) { + for (size_t r = start; r < end; r += 2) { + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (r >= valid_end) { + for (uint32_t d = 0; d < d_limit; ++d) { + __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + *pv_out0 = Q6_V_vzero(); + *pv_out1 = Q6_V_vzero(); + } + continue; + } + + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0); + const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16)); + const HVX_Vector * pv_in1 = (r + 1 < valid_end) + ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16)) + : NULL; + + for (uint32_t d = 0; d < d_limit; ++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_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + + *pv_out0 = Q6_V_lo_W(vp); + *pv_out1 = Q6_V_hi_W(vp); + } + } +} + + +static inline void hmx_fa_q_prep_fallback( + __fp16 * vtcm_q_tiles, uintptr_t q_data, + size_t q_nb1, size_t q_nb2, size_t q_nb3, + uint32_t q_start, uint32_t kv_head, uint32_t ib3, + size_t start, size_t end, size_t n_rows_g, + size_t G, size_t DK, bool is_q_fp32, + const struct fastdiv_values * div_G +) { + for (size_t r = start; r < end; r += 2) { + const size_t q_idx0 = fastdiv(r + 0, div_G); + const size_t h_idx0 = fastmodulo(r + 0, G, div_G); + const size_t q_idx1 = fastdiv(r + 1, div_G); + const size_t h_idx1 = fastmodulo(r + 1, G, div_G); + + const uint8_t * q_ptr0 = (r + 0 < n_rows_g) ? ((const uint8_t *) q_data + (q_start + q_idx0) * q_nb1 + + (kv_head * G + h_idx0) * q_nb2 + ib3 * q_nb3) : + NULL; + const uint8_t * q_ptr1 = (r + 1 < n_rows_g) ? ((const uint8_t *) q_data + (q_start + q_idx1) * q_nb1 + + (kv_head * G + h_idx1) * q_nb2 + ib3 * q_nb3) : + NULL; + + size_t r0 = r / HMX_FP16_TILE_N_ROWS; + size_t r1 = r % HMX_FP16_TILE_N_ROWS; + __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; + + if (is_q_fp32) { + const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL; + const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) 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_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL; + const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) 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_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; + HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; + HVX_Vector * pv_out1 = pv_out0 + 16; + + *pv_out0 = Q6_V_lo_W(vp); + *pv_out1 = Q6_V_hi_W(vp); + } + } + } +} + #endif /* HMX_FA_KERNELS_H */ From fb7d78cd2fcc18dc9c6352c741ce60685a10fecd Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 17:47:06 -0700 Subject: [PATCH 23/32] hmx-fa: use o_tile size from layout instead of computing it --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 22 +++++++++++++--------- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 4 ++++ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index ae4a39c57a47..97c77de50c70 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -143,6 +143,8 @@ struct hmx_fa_context { __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 + size_t q_tile_bytes; + size_t o_tile_bytes; bool mask_broadcast; // true when mask->ne[2] == 1 (head-independent, single 2D DMA) dma_cache m_cache; }; @@ -658,10 +660,10 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { assert(factx->DK == factx->DV); - const size_t o_tile_bytes = hex_align_up(factx->g_br * factx->DV * sizeof(__fp16), 4096); - const bool use_vtcm_temp = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2)); + const size_t o_tile_bytes = factx->o_tile_bytes; + const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2)); - if (use_vtcm_temp) { + if (use_q_dma) { if (factx->is_q_fp32) { char * temp_q_vtcm = (char *) factx->vtcm_o_tiles[0]; const uint32_t d_limit = DK / 32; @@ -706,7 +708,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { { const uint32_t g_br = factx->g_br; const uint32_t DV = factx->DV; - const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096); + const size_t o_tile_bytes = factx->o_tile_bytes; 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); @@ -1648,6 +1650,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { factx.vtcm_hmx_scales_qk = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_qk); factx.vtcm_mask_buf = VTCM_LAYOUT_PTR(__fp16, base, L.off_mask_buf); factx.mask_buf_row_stride = L.mask_buf_row_stride; + factx.q_tile_bytes = L.q_tile_bytes; + factx.o_tile_bytes = L.o_tile_bytes; factx.vtcm_slopes = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes); const size_t m_line_bytes = L.m_line_bytes; // used by the mask DMAs in the KV loop @@ -1692,10 +1696,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const uint32_t iv2 = kv_head; const uint32_t iv3 = ib3 / (neq3 / v->ne[3]); - // 1. Push Q DMA (if vtcm_temp is used) - const size_t o_tile_bytes = hex_align_up(factx.g_br * factx.DV * sizeof(__fp16), 4096); - const bool use_vtcm_temp = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2)); - if (use_vtcm_temp) { + // 1. Push Q DMA (if Q DMA is used) + const size_t o_tile_bytes = factx.o_tile_bytes; + const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2)); + if (use_q_dma) { const bool q_transposed = q->nb[1] < q->nb[2]; const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3]; const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16); @@ -1717,7 +1721,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { } // 3. Pop Q DMA (blocks until Q is loaded) - if (use_vtcm_temp) { + if (use_q_dma) { dma_queue_pop(dma); } diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index b834cec19de1..ce032caa0cf5 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -116,6 +116,8 @@ struct hmx_fa_vtcm_layout { size_t off_slopes; // Region byte sizes reused by the device at runtime (not just for allocation). + size_t q_tile_bytes; + size_t o_tile_bytes; size_t s_tile_bytes; // S and P tiles (same size) size_t d_tile_bytes; size_t m_line_bytes; // one mask row @@ -177,6 +179,8 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size); VTCM_LAYOUT_ALLOC(off, off_slopes, slopes_size); + L->q_tile_bytes = q_tile_size; + L->o_tile_bytes = o_tile_size; L->s_tile_bytes = s_tile_size; L->d_tile_bytes = d_tile_size; L->m_line_bytes = m_line_size; From 89698e1085340a404e78175436b6b547319e7051 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 18:11:15 -0700 Subject: [PATCH 24/32] hmx-mm: cleanup types and size handling --- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 2 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 47 +++++++++++---------- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 48 +++++++++++----------- 3 files changed, 51 insertions(+), 46 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index ca22d8f719c3..ceb16080bec1 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -20,6 +20,7 @@ add_library(${HTP_LIB} SHARED worker-pool.c hex-dma.c hmx-queue.c + flash-attn-ops.c matmul-ops.c gated-delta-net-ops.c binary-ops.c @@ -40,7 +41,6 @@ add_library(${HTP_LIB} SHARED diag-ops.c solve-tri-ops.c pad-ops.c - flash-attn-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 97c77de50c70..216e540009c3 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -663,35 +663,40 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { const size_t o_tile_bytes = factx->o_tile_bytes; const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2)); + __fp16 * q_tiles = factx->vtcm_q_tiles; if (use_q_dma) { + const size_t g_rows_end = hex_smin(end, n_rows_g); + const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64; + + uint8_t * q_flat = (uint8_t *) factx->vtcm_o_tiles[0]; if (factx->is_q_fp32) { - char * temp_q_vtcm = (char *) factx->vtcm_o_tiles[0]; - const uint32_t d_limit = DK / 32; - const size_t valid_end = hex_smin(end, n_rows_g); - - if (d_limit == 2) { - hmx_fa_q_prep_fp32_d2(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - } else if (d_limit == 4) { - hmx_fa_q_prep_fp32_d4(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - } else { - hmx_fa_q_prep_fp32_generic(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); + switch (d_limit) { + case 2: + hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + break; + case 4: + hmx_fa_q_prep_fp32_d4(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + break; + default: + hmx_fa_q_prep_fp32_generic(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); + break; } } else { - char * temp_q_vtcm = (char *) factx->vtcm_o_tiles[0]; - const uint32_t d_limit = DK / 64; - const size_t valid_end = hex_smin(end, n_rows_g); - - if (d_limit == 1) { - hmx_fa_q_prep_fp16_d1(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - } else if (d_limit == 2) { - hmx_fa_q_prep_fp16_d2(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - } else { - hmx_fa_q_prep_fp16_generic(factx->vtcm_q_tiles, temp_q_vtcm, start, end, valid_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); + switch (d_limit) { + case 1: + hmx_fa_q_prep_fp16_d1(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + break; + case 2: + hmx_fa_q_prep_fp16_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); + break; + default: + hmx_fa_q_prep_fp16_generic(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); + break; } } } else { // Fallback: Original direct-from-DDR/L2 path - hmx_fa_q_prep_fallback(factx->vtcm_q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3], + hmx_fa_q_prep_fallback(q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3], q_start, kv_head, ib3, start, end, n_rows_g, G, DK, factx->is_q_fp32, &factx->div_G); } } diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index b26d2abbb460..040c66613797 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -210,8 +210,8 @@ static inline void hmx_fa_o_norm_tile( } static inline void hmx_fa_q_prep_fp32_d2( - __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, - size_t start, size_t end, size_t valid_end, + __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, + size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, const struct fastdiv_values * div_G, bool q_transposed ) { @@ -220,7 +220,7 @@ static inline void hmx_fa_q_prep_fp32_d2( size_t r1 = r % HMX_FP16_TILE_N_ROWS; __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; - if (r >= valid_end) { + if (r >= g_rows_end) { ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); continue; @@ -235,7 +235,7 @@ static inline void hmx_fa_q_prep_fp32_d2( const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float)); - const HVX_Vector * pv_in1 = (r + 1 < valid_end) + const HVX_Vector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float)) : NULL; @@ -255,8 +255,8 @@ static inline void hmx_fa_q_prep_fp32_d2( } static inline void hmx_fa_q_prep_fp32_d4( - __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, - size_t start, size_t end, size_t valid_end, + __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, + size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, const struct fastdiv_values * div_G, bool q_transposed ) { @@ -265,7 +265,7 @@ static inline void hmx_fa_q_prep_fp32_d4( size_t r1 = r % HMX_FP16_TILE_N_ROWS; __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; - if (r >= valid_end) { + if (r >= g_rows_end) { for (uint32_t d = 0; d < 4; ++d) { ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); } @@ -281,7 +281,7 @@ static inline void hmx_fa_q_prep_fp32_d4( const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float)); - const HVX_Vector * pv_in1 = (r + 1 < valid_end) + const HVX_Vector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float)) : NULL; @@ -295,8 +295,8 @@ static inline void hmx_fa_q_prep_fp32_d4( } static inline void hmx_fa_q_prep_fp32_generic( - __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, - size_t start, size_t end, size_t valid_end, + __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, + size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed ) { @@ -305,7 +305,7 @@ static inline void hmx_fa_q_prep_fp32_generic( size_t r1 = r % HMX_FP16_TILE_N_ROWS; __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; - if (r >= valid_end) { + if (r >= g_rows_end) { for (uint32_t d = 0; d < d_limit; ++d) { ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero(); } @@ -321,7 +321,7 @@ static inline void hmx_fa_q_prep_fp32_generic( const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float)); - const HVX_Vector * pv_in1 = (r + 1 < valid_end) + const HVX_Vector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float)) : NULL; @@ -337,8 +337,8 @@ static inline void hmx_fa_q_prep_fp32_generic( } static inline void hmx_fa_q_prep_fp16_d1( - __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, - size_t start, size_t end, size_t valid_end, + __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, + size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, const struct fastdiv_values * div_G, bool q_transposed ) { @@ -347,7 +347,7 @@ static inline void hmx_fa_q_prep_fp16_d1( size_t r1 = r % HMX_FP16_TILE_N_ROWS; __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; - if (r >= valid_end) { + if (r >= g_rows_end) { __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2; HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; HVX_Vector * pv_out1 = pv_out0 + 16; @@ -365,7 +365,7 @@ static inline void hmx_fa_q_prep_fp16_d1( const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16)); - const HVX_Vector * pv_in1 = (r + 1 < valid_end) + const HVX_Vector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16)) : NULL; @@ -383,8 +383,8 @@ static inline void hmx_fa_q_prep_fp16_d1( } static inline void hmx_fa_q_prep_fp16_d2( - __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, - size_t start, size_t end, size_t valid_end, + __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, + size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, const struct fastdiv_values * div_G, bool q_transposed ) { @@ -393,7 +393,7 @@ static inline void hmx_fa_q_prep_fp16_d2( size_t r1 = r % HMX_FP16_TILE_N_ROWS; __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; - if (r >= valid_end) { + if (r >= g_rows_end) { for (uint32_t d = 0; d < 2; ++d) { __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; @@ -413,7 +413,7 @@ static inline void hmx_fa_q_prep_fp16_d2( const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16)); - const HVX_Vector * pv_in1 = (r + 1 < valid_end) + const HVX_Vector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16)) : NULL; @@ -445,8 +445,8 @@ static inline void hmx_fa_q_prep_fp16_d2( } static inline void hmx_fa_q_prep_fp16_generic( - __fp16 * vtcm_q_tiles, const char * temp_q_vtcm, - size_t start, size_t end, size_t valid_end, + __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, + size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed ) { @@ -455,7 +455,7 @@ static inline void hmx_fa_q_prep_fp16_generic( size_t r1 = r % HMX_FP16_TILE_N_ROWS; __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK; - if (r >= valid_end) { + if (r >= g_rows_end) { for (uint32_t d = 0; d < d_limit; ++d) { __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2; HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2; @@ -475,7 +475,7 @@ static inline void hmx_fa_q_prep_fp16_generic( const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1); const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16)); - const HVX_Vector * pv_in1 = (r + 1 < valid_end) + const HVX_Vector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16)) : NULL; From 84466632388da4a853456a87afa810d501214bf1 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 18:21:27 -0700 Subject: [PATCH 25/32] hmx-mm: replace divs with fastdiv in qprep loops --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 8 ++++---- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 13 ++++++------- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 9 +++++---- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 502baf9397b6..2d13ff6dc4a4 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -2028,10 +2028,10 @@ static bool ggml_hexagon_precompute_flash_attn_params( 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]); + kparams->broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]); + kparams->broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]); + kparams->broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]); + kparams->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]); diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 216e540009c3..328c911e0aaa 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1680,7 +1680,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16); - // ======== Reusable job descriptors for pipeline ======== hmx_fa_qk_job_t qk_job; hmx_fa_o_update_job_t ou_job; @@ -1697,9 +1696,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { 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 ik3 = fastdiv(ib3, &kparams->broadcast_rk3); const uint32_t iv2 = kv_head; - const uint32_t iv3 = ib3 / (neq3 / v->ne[3]); + const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3); // 1. Push Q DMA (if Q DMA is used) const size_t o_tile_bytes = factx.o_tile_bytes; @@ -2041,10 +2040,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { factx.src0_div21 = kparams->u.hvx.src0_div21; factx.src0_div1 = kparams->u.hvx.src0_div1; - 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; + factx.broadcast_rk2 = kparams->broadcast_rk2; + factx.broadcast_rk3 = kparams->broadcast_rk3; + factx.broadcast_rv2 = kparams->broadcast_rv2; + factx.broadcast_rv3 = kparams->broadcast_rv3; if (mask) { factx.src3_div2 = kparams->src3_div2; diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index ce032caa0cf5..cda58730ffcc 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -55,6 +55,11 @@ struct htp_fa_kernel_params { struct fastdiv_values src3_div2; struct fastdiv_values src3_div3; + struct fastdiv_values broadcast_rk2; + struct fastdiv_values broadcast_rk3; + struct fastdiv_values broadcast_rv2; + struct fastdiv_values broadcast_rv3; + union { struct { uint32_t g_br; @@ -70,10 +75,6 @@ struct htp_fa_kernel_params { 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; }; From 5af9aa0085477e2c83638a4094c1f42650291d62 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 21:56:45 -0700 Subject: [PATCH 26/32] hmx-fa: minor update/formatting to q_tile handling --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 44 ++++++++-------------- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 4 +- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 328c911e0aaa..aef99081225e 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -510,8 +510,8 @@ typedef struct { struct hmx_fa_context * factx; uint32_t kv_rows; size_t src_stride; - void * curr_v; - size_t buf_idx; + void * v_src; + void * v_tiles_dst; size_t n_col_tiles; uint32_t kv_start; uint32_t rows_per_t; @@ -530,11 +530,11 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) return; } - __fp16 * v_tiles_dest = factx->pipeline ? factx->vtcm_v_tiles[args->buf_idx] : factx->vtcm_v_tiles[0]; + __fp16 * v_tiles_dst = (__fp16 *) args->v_tiles_dst; 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, (const __fp16 *) args->curr_v, total_rows, factx->DV, + hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV, args->src_stride, (uint32_t) args->n_col_tiles, start, end); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start)); } @@ -542,8 +542,8 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) static void fa_phase_v_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, - void * curr_v, - size_t buf_idx, + void * v_src, + void * v_tiles_dst, size_t n_col_tiles, uint32_t kv_start) { worker_pool_context_t wp = factx->octx->ctx->worker_pool; @@ -552,7 +552,7 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx, 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, curr_v, buf_idx, n_col_tiles, kv_start, rows_per_t }; + fa_v_int_args_t args = { factx, kv_rows, src_stride, v_src, v_tiles_dst, n_col_tiles, kv_start, rows_per_t }; if (n > 1) { worker_pool_run_func(wp, fa_v_interleave_thread, &args, n); } else { @@ -671,31 +671,19 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { uint8_t * q_flat = (uint8_t *) factx->vtcm_o_tiles[0]; if (factx->is_q_fp32) { switch (d_limit) { - case 2: - hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - break; - case 4: - hmx_fa_q_prep_fp32_d4(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - break; - default: - hmx_fa_q_prep_fp32_generic(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); - break; + case 2: hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break; + case 4: hmx_fa_q_prep_fp32_d4(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break; + default: hmx_fa_q_prep_fp32( q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); break; } } else { switch (d_limit) { - case 1: - hmx_fa_q_prep_fp16_d1(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - break; - case 2: - hmx_fa_q_prep_fp16_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); - break; - default: - hmx_fa_q_prep_fp16_generic(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); - break; + case 1: hmx_fa_q_prep_fp16_d1(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break; + case 2: hmx_fa_q_prep_fp16_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break; + default: hmx_fa_q_prep_fp16( q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); break; } } } else { - // Fallback: Original direct-from-DDR/L2 path + // Fallback: direct-from-DDR/L2 path hmx_fa_q_prep_fallback(q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3], q_start, kv_head, ib3, start, end, n_rows_g, G, DK, factx->is_q_fp32, &factx->div_G); } @@ -1810,7 +1798,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Wait for current V DMA and interleave void * curr_v = dma_queue_pop(dma).dst; - fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, buf_idx, n_tiles_per_bc, kv_start); + fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start); if (kv_blk > 0) { hmx_queue_pop(hmx_q); @@ -1924,7 +1912,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Wait for current V DMA and interleave void * curr_v = dma_queue_pop(dma).dst; - fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, buf_idx, n_tiles_per_bc, kv_start); + fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[0], n_tiles_per_bc, kv_start); // ---- Phase 3: softmax + build_D ---- __fp16 * current_mask_vtcm = NULL; diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index 040c66613797..d6795bf0b57d 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -294,7 +294,7 @@ static inline void hmx_fa_q_prep_fp32_d4( } } -static inline void hmx_fa_q_prep_fp32_generic( +static inline void hmx_fa_q_prep_fp32( __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, @@ -444,7 +444,7 @@ static inline void hmx_fa_q_prep_fp16_d2( } } -static inline void hmx_fa_q_prep_fp16_generic( +static inline void hmx_fa_q_prep_fp16( __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm, size_t start, size_t end, size_t g_rows_end, size_t DK, size_t G, size_t n_rows_q, From 32c360e62257d52d79c50e0f8231587e0e15280a Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 22:58:20 -0700 Subject: [PATCH 27/32] hmx-fa: cleanup the layout to avoid overpadding --- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 2 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 12 +++-- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 54 ++++++++++++---------- 3 files changed, 38 insertions(+), 30 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index ceb16080bec1..ca22d8f719c3 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -20,7 +20,6 @@ add_library(${HTP_LIB} SHARED worker-pool.c hex-dma.c hmx-queue.c - flash-attn-ops.c matmul-ops.c gated-delta-net-ops.c binary-ops.c @@ -41,6 +40,7 @@ add_library(${HTP_LIB} SHARED diag-ops.c solve-tri-ops.c pad-ops.c + flash-attn-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index aef99081225e..72e84a046190 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -145,6 +145,8 @@ struct hmx_fa_context { size_t mask_buf_row_stride; // elements (__fp16) per row in mask buffer size_t q_tile_bytes; size_t o_tile_bytes; + size_t col_vec_bytes; + size_t d_tile_bytes; bool mask_broadcast; // true when mask->ne[2] == 1 (head-independent, single 2D DMA) dma_cache m_cache; }; @@ -595,8 +597,8 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { 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 col_vec_bytes = factx->col_vec_bytes; + const size_t d_tile_bytes = factx->d_tile_bytes; // Initialize vtcm_l_vec & vtcm_m_vec const size_t l_bytes_per_t = hex_align_up(col_vec_bytes / n, 128); @@ -1608,9 +1610,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const size_t size_v_row_padded = hex_round_up(size_v_row, 128); // Build the VTCM layout once (shared with the host estimator) and place every - // scratch buffer at its computed offset. This is the single source of truth for - // sizes/offsets/strides -- host and device can no longer disagree. - _Static_assert(sizeof(HVX_Vector) == HMX_FA_HVX_VECTOR_BYTES, "HVX vector size mismatch"); + // scratch buffer at its computed offset. struct hmx_fa_vtcm_layout L; hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline); @@ -1645,6 +1645,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { factx.mask_buf_row_stride = L.mask_buf_row_stride; factx.q_tile_bytes = L.q_tile_bytes; factx.o_tile_bytes = L.o_tile_bytes; + factx.col_vec_bytes = L.col_vec_bytes; + factx.d_tile_bytes = L.d_tile_bytes; factx.vtcm_slopes = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes); const size_t m_line_bytes = L.m_line_bytes; // used by the mask DMAs in the KV loop diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index cda58730ffcc..e778070aa963 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -14,13 +14,16 @@ extern "C" { #endif // Tile constants (mirrored from hmx-utils.h for use on host side if needed) +#define HTP_FA_HMX_TILE_SIZE 2048 +#define HMX_FP16_TILE_SIZE 2048 #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 { @@ -123,6 +126,7 @@ struct hmx_fa_vtcm_layout { size_t d_tile_bytes; size_t m_line_bytes; // one mask row size_t m_buf_slot_bytes; // one dma_cache slot = align_up(Br * m_line_bytes, 4096) + size_t col_vec_bytes; // Derived strides. size_t row_buf_stride; // HVX vectors (128B) per row buffer @@ -132,44 +136,46 @@ struct hmx_fa_vtcm_layout { size_t total_bytes; }; -// HVX vector size in bytes. The device asserts sizeof(HVX_Vector) == this. -#define HMX_FA_HVX_VECTOR_BYTES 128 +// Build the VTCM layout. -// Build the VTCM layout. Mirrors the device's vtcm_seq_alloc sequence exactly: -// a running offset bumped by each region size, with no inter-region alignment. static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) { 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 - const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 4096); // V DMA: [Bc, DV] x2 - 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, s_rowmax, p_rowsum - 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_slot = hex_align_up(Br * m_line_size, 4096); + const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); + const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); + const size_t k_tile_size = hex_align_up(Bc * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); + const size_t v_tile_size = hex_align_up(Bc * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); + const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); + const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); + + const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128); + const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128); + const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256); + 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_slot = hex_align_up(Br * m_line_size, 256); const size_t m_buf_size = m_buf_slot * HMX_FA_DMA_CACHE_SIZE; const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128); size_t off = 0; + + // Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned) VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size); VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size); VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size); - VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size); - VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size); - VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size); - VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size); VTCM_LAYOUT_ALLOC(off, off_k_tiles, k_tile_size); VTCM_LAYOUT_ALLOC(off, off_v_tiles[0], v_tile_size); VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline); VTCM_LAYOUT_ALLOC(off, off_s_tiles, s_tile_size); VTCM_LAYOUT_ALLOC(off, off_p_tiles, s_tile_size); VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size); + + // Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned) + VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size); + VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size); + VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size); + VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size); VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size); VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size); VTCM_LAYOUT_ALLOC(off, off_s_rowmax, col_vec_size); @@ -182,11 +188,12 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, L->q_tile_bytes = q_tile_size; L->o_tile_bytes = o_tile_size; + L->col_vec_bytes = col_vec_size; L->s_tile_bytes = s_tile_size; L->d_tile_bytes = d_tile_size; L->m_line_bytes = m_line_size; L->m_buf_slot_bytes = m_buf_slot; - L->row_buf_stride = row_vec_size / HMX_FA_HVX_VECTOR_BYTES; + L->row_buf_stride = row_vec_size / 128; L->mask_buf_row_stride = m_line_size / sizeof(__fp16); L->pipeline = pipeline; L->total_bytes = off; @@ -242,8 +249,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out, // 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_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; From e4ac61bd6f19df41c112603b6253e3e5849abd8a Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 6 Jul 2026 23:54:04 -0700 Subject: [PATCH 28/32] hmx-fa: simplified and improved cost mode for hmx fa solver that uses vtcm layout funcs --- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 77 ++++++---------------- 1 file changed, 19 insertions(+), 58 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index e778070aa963..16822f22bf6e 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -243,22 +243,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out, 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; @@ -274,51 +260,26 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out, 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) { + // Try all Bc candidates from Bc_limit down to bc_unit + for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) { + size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline); + if (vtcm_needed <= vtcm_budget) { + // This Bc fits for this Br! + const size_t q_blocks = (qo_len + Br - 1) / Br; + const size_t kv_blocks = (kv_len + Bc - 1) / Bc; + const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed); + const size_t mn = Br * Bc; + + if (cost < best_cost || (cost == best_cost && mn > best_mn)) { + best_cost = cost; + best_mn = mn; + best_Br = Br; + best_Bc = Bc; + } + // Since we iterate Bc from largest to smallest, this is the largest Bc that fits + // for this Br. We can break to the next Br. 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) { @@ -326,7 +287,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out, } } - if (best_Br == 0) { + if (best_Br == 0 || best_Bc == 0) { return -1; } From c295549196046b83121918523e1e365760d360cc Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Sun, 5 Jul 2026 23:47:25 -0700 Subject: [PATCH 29/32] hmx-queue: add support queue wakeup and make suspend async to avoid hmx-lock latency --- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 2 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 1 + ggml/src/ggml-hexagon/htp/hmx-queue.c | 1 + ggml/src/ggml-hexagon/htp/hmx-queue.h | 22 +++++++++++++++++++--- ggml/src/ggml-hexagon/htp/matmul-ops.c | 1 + 5 files changed, 23 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index ca22d8f719c3..6ebb4b4888ef 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -21,6 +21,7 @@ add_library(${HTP_LIB} SHARED hex-dma.c hmx-queue.c matmul-ops.c + flash-attn-ops.c gated-delta-net-ops.c binary-ops.c unary-ops.c @@ -40,7 +41,6 @@ add_library(${HTP_LIB} SHARED diag-ops.c solve-tri-ops.c pad-ops.c - flash-attn-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 72e84a046190..a669b9f7a9df 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1740,6 +1740,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Pipeline path // ================================================================== struct hmx_queue * hmx_q = ctx->hmx_queue; + hmx_queue_wakeup(hmx_q); for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) { const uint32_t kv_start = kv_blk * Bc; diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.c b/ggml/src/ggml-hexagon/htp/hmx-queue.c index a0007539c561..e2e507a3f502 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.c +++ b/ggml/src/ggml-hexagon/htp/hmx-queue.c @@ -42,6 +42,7 @@ static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) { case HMX_QUEUE_NOOP: /* noop */; break; case HMX_QUEUE_KILL: *killed = true; break; case HMX_QUEUE_SUSPEND: hmx_unlock(q); break; + case HMX_QUEUE_WAKEUP: hmx_lock(q); break; default: hmx_lock(q); htp_trace_event_start(q->trace, HTP_TRACE_EVT_HMX_COMP, ir); diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.h b/ggml/src/ggml-hexagon/htp/hmx-queue.h index 83135cd91d8b..e8fac2781373 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.h +++ b/ggml/src/ggml-hexagon/htp/hmx-queue.h @@ -25,6 +25,7 @@ typedef void (*hmx_queue_func)(void *); // Dummy funcs used as signals enum hmx_queue_signal { HMX_QUEUE_NOOP = 0, // aka NULL + HMX_QUEUE_WAKEUP, HMX_QUEUE_SUSPEND, HMX_QUEUE_KILL }; @@ -97,7 +98,7 @@ static inline uint32_t hmx_queue_capacity(struct hmx_queue * q) { return q->capacity; } -static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) { +static inline struct hmx_queue_desc hmx_queue_pop_one(struct hmx_queue * q) { unsigned int ip = q->idx_pop; unsigned int iw = q->idx_write; @@ -120,13 +121,28 @@ static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) { return rd; } +static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) { + while (1) { + struct hmx_queue_desc d = hmx_queue_pop_one(q); + + uint32_t sig = (uint32_t) d.func; + if (sig && sig <= HMX_QUEUE_KILL) + continue; + + return d; + } +} + static inline void hmx_queue_flush(struct hmx_queue * q) { - while (hmx_queue_pop(q).func != NULL) ; + while (hmx_queue_pop_one(q).func != NULL) ; +} + +static inline void hmx_queue_wakeup(struct hmx_queue * q) { + hmx_queue_signal(q, HMX_QUEUE_WAKEUP); } static inline void hmx_queue_suspend(struct hmx_queue *q) { hmx_queue_signal(q, HMX_QUEUE_SUSPEND); - hmx_queue_flush(q); } #ifdef __cplusplus diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 3453f457c792..409c30ce4b84 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2395,6 +2395,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, if (pipeline) { // --- Asynchronous Pipelined Loop --- + hmx_queue_wakeup(ctx->hmx_queue); hmx_matmul_job_t job_slots[2]; // persistent double-buffered job descriptors for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { From 71c65293f80d5a35b7bfec1590da395b0dfd3d47 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Tue, 7 Jul 2026 11:40:30 -0700 Subject: [PATCH 30/32] hex-hmx: move queue wakeup / suspend to the op-batch level --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 11 ++--------- ggml/src/ggml-hexagon/htp/hmx-queue.c | 2 +- ggml/src/ggml-hexagon/htp/main.c | 4 ++++ ggml/src/ggml-hexagon/htp/matmul-ops.c | 6 ------ 4 files changed, 7 insertions(+), 16 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index a669b9f7a9df..6f2a643e69da 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -1735,13 +1735,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const size_t k_src_stride = size_k_row_padded / sizeof(__fp16); const size_t v_src_stride = size_v_row_padded / sizeof(__fp16); + struct hmx_queue * hmx_q = ctx->hmx_queue; + if (factx.pipeline) { - // ================================================================== // Pipeline path - // ================================================================== - struct hmx_queue * hmx_q = ctx->hmx_queue; - hmx_queue_wakeup(hmx_q); - 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); @@ -1868,9 +1865,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { } } 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); @@ -1996,8 +1991,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { } } - hmx_queue_suspend(ctx->hmx_queue); - return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.c b/ggml/src/ggml-hexagon/htp/hmx-queue.c index e2e507a3f502..129fc54cc87c 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.c +++ b/ggml/src/ggml-hexagon/htp/hmx-queue.c @@ -66,7 +66,7 @@ static void hmx_queue_thread(void * arg) { bool killed = false; - unsigned int poll_cnt = HMX_QUEUE_POLL_COUNT; + unsigned int poll_cnt = 1; // HMX_QUEUE_POLL_COUNT; unsigned int prev_seqn = 0; while (!killed) { unsigned int seqn = atomic_load(&q->seqn); diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 684625e4275f..bfb3e3138ead 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -948,6 +948,8 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) { int op_status = HTP_STATUS_OK; uint32_t op_wakeup = n_ops / 2; // half-way throgh the batch + hmx_queue_wakeup(ctx->hmx_queue); + for (uint32_t i=0; i < n_ops; i++) { struct profile_data prof; @@ -976,6 +978,8 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) { } } + hmx_queue_suspend(ctx->hmx_queue); + struct htp_opbatch_rsp rsp; rsp.id = req.id; rsp.status = op_status; diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 409c30ce4b84..1683131a813a 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -2395,7 +2395,6 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, if (pipeline) { // --- Asynchronous Pipelined Loop --- - hmx_queue_wakeup(ctx->hmx_queue); hmx_matmul_job_t job_slots[2]; // persistent double-buffered job descriptors for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { @@ -2528,8 +2527,6 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, } } - hmx_queue_suspend(ctx->hmx_queue); - return 0; } @@ -2969,12 +2966,9 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, } } - hmx_queue_suspend(ctx->hmx_queue); - return 0; } - // --- Dispatchers and Public Entry Points --- static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) { From 074fedeef29da21eab449f485edc83b6284224c5 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Sun, 5 Jul 2026 17:38:31 -0700 Subject: [PATCH 31/32] hex-threads: add hybrid polling to workpool --- ggml/src/ggml-hexagon/htp/CMakeLists.txt | 4 ++-- ggml/src/ggml-hexagon/htp/hmx-queue.c | 8 +++++++- ggml/src/ggml-hexagon/htp/hmx-queue.h | 7 ++++++- ggml/src/ggml-hexagon/htp/worker-pool.c | 18 +++++++++++++++--- ggml/src/ggml-hexagon/htp/worker-pool.h | 8 ++++++++ 5 files changed, 38 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 6ebb4b4888ef..be575796d6d3 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -20,8 +20,6 @@ add_library(${HTP_LIB} SHARED worker-pool.c hex-dma.c hmx-queue.c - matmul-ops.c - flash-attn-ops.c gated-delta-net-ops.c binary-ops.c unary-ops.c @@ -41,6 +39,8 @@ add_library(${HTP_LIB} SHARED diag-ops.c solve-tri-ops.c pad-ops.c + flash-attn-ops.c + matmul-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.c b/ggml/src/ggml-hexagon/htp/hmx-queue.c index 129fc54cc87c..5f6a5e206bb9 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.c +++ b/ggml/src/ggml-hexagon/htp/hmx-queue.c @@ -6,6 +6,7 @@ #include #include +#include #include @@ -66,14 +67,19 @@ static void hmx_queue_thread(void * arg) { bool killed = false; - unsigned int poll_cnt = 1; // HMX_QUEUE_POLL_COUNT; + unsigned int poll_cnt = HMX_QUEUE_POLL_COUNT; unsigned int prev_seqn = 0; while (!killed) { unsigned int seqn = atomic_load(&q->seqn); if (seqn == prev_seqn) { + // drop HVX context while spinning + if (poll_cnt > 1 && poll_cnt == HMX_QUEUE_POLL_COUNT) { + qurt_hvx_unlock(); + } if (--poll_cnt) { hex_pause(); continue; } FARF(HIGH, "hmx-queue-thread: sleeping"); qurt_futex_wait(&q->seqn, prev_seqn); + poll_cnt = HMX_QUEUE_POLL_COUNT; continue; } prev_seqn = seqn; diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.h b/ggml/src/ggml-hexagon/htp/hmx-queue.h index e8fac2781373..b176fa179611 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.h +++ b/ggml/src/ggml-hexagon/htp/hmx-queue.h @@ -18,7 +18,12 @@ extern "C" { #endif #define HMX_QUEUE_THREAD_STACK_SIZE (16 * 1024) -#define HMX_QUEUE_POLL_COUNT 2000 + +#if __HVX_ARCH__ > 79 +#define HMX_QUEUE_POLL_COUNT 2000 +#else +#define HMX_QUEUE_POLL_COUNT 1 +#endif typedef void (*hmx_queue_func)(void *); diff --git a/ggml/src/ggml-hexagon/htp/worker-pool.c b/ggml/src/ggml-hexagon/htp/worker-pool.c index 172e28908ebb..50960d2c75d7 100644 --- a/ggml/src/ggml-hexagon/htp/worker-pool.c +++ b/ggml/src/ggml-hexagon/htp/worker-pool.c @@ -1,6 +1,9 @@ #include "worker-pool.h" +#include "hex-utils.h" #include +#include + #include #include #include @@ -9,7 +12,6 @@ #include "HAP_farf.h" -#define WORKER_THREAD_STACK_SZ (2 * 16384) #define LOWEST_USABLE_QURT_PRIO (254) struct worker_pool_s; @@ -42,17 +44,27 @@ static void worker_pool_main(void * context) { FARF(HIGH, "worker-pool: thread %u started", me->id); unsigned int prev_seqn = 0; + unsigned int poll_cnt = WORKER_POOL_POLL_COUNT; while (!atomic_load(&pool->killed)) { unsigned int seqn = atomic_load(&pool->seqn); if (seqn == prev_seqn) { - // Nothing to do + // drop HVX context while spinning + if (poll_cnt > 1 && poll_cnt == WORKER_POOL_POLL_COUNT) { + qurt_hvx_unlock(); + } + if (--poll_cnt) { + hex_pause(); + continue; + } qurt_futex_wait(&pool->seqn, prev_seqn); + poll_cnt = WORKER_POOL_POLL_COUNT; continue; } - // New job prev_seqn = seqn; + poll_cnt = WORKER_POOL_POLL_COUNT; + // New job unsigned int n = atomic_load(&pool->n_jobs); unsigned int i = atomic_fetch_add(&pool->next_job, 1); if (i >= n) { diff --git a/ggml/src/ggml-hexagon/htp/worker-pool.h b/ggml/src/ggml-hexagon/htp/worker-pool.h index 6f8c9056c4b4..cba692126ad0 100644 --- a/ggml/src/ggml-hexagon/htp/worker-pool.h +++ b/ggml/src/ggml-hexagon/htp/worker-pool.h @@ -24,9 +24,17 @@ typedef struct { void * data; } worker_pool_job_t; +#define WORKER_THREAD_STACK_SZ (2 * 16384) + /// Maximum supported number of worker threads. #define MAX_NUM_WORKERS 10 +#if __HVX_ARCH__ > 79 +#define WORKER_POOL_POLL_COUNT 2000 +#else +#define WORKER_POOL_POLL_COUNT 1 +#endif + // Initialize worker pool. WORKERPOOL_API AEEResult worker_pool_init(worker_pool_context_t * context, uint32_t n_threads); From e7ba42c381d51606a31b878637ede13d4b7597ba Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Tue, 7 Jul 2026 22:02:12 -0700 Subject: [PATCH 32/32] hex-mm: fix trailing spaces --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 8 ++++---- ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 2d13ff6dc4a4..3d9fc9b56dd0 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -2627,7 +2627,7 @@ static void ggml_hexagon_precompute_fused_qkv_params( struct htp_mm_hvx_vtcm_layout L; bool try_tiled = (opt_mm_select >= 2); - + // Test tiled first htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, @@ -2646,7 +2646,7 @@ static void ggml_hexagon_precompute_fused_qkv_params( } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - + htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false @@ -2697,7 +2697,7 @@ static void ggml_hexagon_precompute_fused_ffn_params( struct htp_mm_hvx_vtcm_layout L; bool try_tiled = (opt_mm_select >= 2); - + // Test tiled first htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, @@ -2715,7 +2715,7 @@ static void ggml_hexagon_precompute_fused_ffn_params( } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - + htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true 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 96371f2b89bf..740a8f87d61f 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -628,7 +628,7 @@ static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *res } } -static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, +static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { if (n_dot_tiles <= 32) {