hexagon: support for multi-NPU devices (IQ9, IQ10) and fully asynchronous backend (#26501)
* hexagon: use non-host bufs by default and make the backend fully async * hex-hb: remove optional hostbuf support and fix async copy * hex-unary: relax supported unary check * hex-bufs: use same get_alignment for host bufs * snapdragon: bump android_platform to 34 * hex-rows: super hacky get/set rows for q8_0 * hex-get-rows: fix q8_0 * hex-get-rows: supprot for f16 and cleanup for q8_0 * hex-get-rows: generic macros and specialized thread funcs * hex-get-rows: add DMA pipeline, vtcm_layout and kernel params * hex-set-rows: fix q8_0 support, add dma and tracing * hex-tests: override nmse threshold for HTP of Q8_0 quants * hex-fa: add support for Q8_0 with inplace dequantizers * hex-get-rows: simplify type dispatch * hex-rows: simplify GET/SET_ROWS DMA pipeline * hex-async: add events, set/get-tensor-async and rest of the async api support * hex-repack: use slice instead of expert in repack functions * hex-cpy: update event/async-cpy logging * hex-set-rows: optimize smaller tensors * hex-geglu: fix perf regression with larger tensors * hex-get-rows: add missing header * hex-set-rows: add missing header * hex-bufs: ressurect GGML_HEXAGON_HOSTBUF but disable it by default * hexagon: do not reject ops with non-heaxon buffers * hex-get-rows: apply >=32 restriction only for q8_0 * hex-res: bump vtcm acquire timeout to 10 seconds * hex-bufs: add support for cloning buffers between sessions to speed up tensor copies * hex-async: rework event recording and batch flushing and integrate with meta backend * hex-bufs: improved handling of repacked tensors * hex-repack: handle get_tensor_2d offsets * hex-dev: add support for devices with multiple NPUs * hex-sync: add support for sync tokens to synchronize npu devices for async splits * hex-mmap: cleanup mmap calls and add a retry for robustness * hex-sync: add failsafe if sync wait gets stuck * hex-sync: use sync_seq to check for completed events * hex-sync: rotate tokens for extra robustness * hex-devs: add supprot for legacy device names for now * hex-bufs: add support for auto-cloning buffers from diff sessions * hex-fusion: simplify and optimize htp-opnode fusion handling * hex-sync: override opnode name so that it shows up in the profiles * hex-trace: update scripts to handle multiple devices * hex-sync: bump the size of the opbatch queue and number of sync tokens * hex-cpy-sync: do not explicitly flush opbatches in cpy_tensor_async and add support for cpy-dma * hex-sync: add graph-flush threshold to avoid single op batches * hex-sync: add sync_peer so that we can flush peers we depend on during cross-device ops * hex-bufs: introduce tensor->extra and shadow_bufs for repacking * hex-l2: flush tiny tensors inline * hex-sync: use explicit l2flush for sync tokens * hex-extra: track weight flags via tensor extra * hex-fence: rename sync to fence * hex-repack: proper handling of set-tensor-2d in the shadow_buf * hex-trace: remove obsolete opstage mask that we used for profiling * hex-env: remove obsolete use_hmx variable * hexagon: new unified run.py and build.py and updated docs * snapdragon: update run script to auto-escapt test-backend-op -p argument * hex-scripts: fix trailing spaces * hex-scripts: fix flake8 warnings * snapdragon: cleanup dst lib/bin dirs before copying new build * hex-ops: add support for allreduce * hex-ar: improved allreduce with dma pipeline * hex-ar: align macros * hex-ar: consistent use of fence_seq * hex-ar: add AR_SELECT env var to select ALLREDUCE kernel or fallback * hex-ar: add proper synchronize handling for ALLREDUCE * hex-opbatch: looks like we now just rely on backend.synchronise to flush the batches, no need to flush them by threshold * hex-ar: bump block size to improve dma efficiency * hex-ar: fused ALLREDUCE+ADD * hex-ar: cleaner fence buffer management * hex-ar: futher allreduce tweaking to remove race conditions * hex-ar: add simple solver and remove non-dma kernels * hex-ar: add row-broadcast to fuse with bias ADD * hex-fence: pass seq numbers via op_params * hex-ar: allow for both entry/exit seq for completing entry wait * hex-ar: align macros * hex-ar: do not refetch broadcast row * hex-fusion: move all fusion into opbatch::add_op for consistency with ALLREDUCE and things * hex-fusion: fix incorrect MUL_MAT reordering * hex-mm: make fused 2x and 3x matmuls more generic * hex-fusion: move tensor fusion tagging to graph_compute * hexagon: make sure to copy tensor->extra by value * hex-get-rows: fix offset calc with row-chunking * hex-repack: get_tensor_2d fixes for non-zero offsets * snapdragon: make profile/trace scripts more robust and donot mix stdout/stderr by default * hex-devices: use legacy device nameing by default to ease the transition * hex-devices: hardcode CDSP domain IDs for current devices for now * hex-optrace: improve multi-NPU timestamp alignment and overall handling of cycle values * hex-optrace: more robust handling of the fence events
This commit is contained in:
Max Krasnyansky
@@ -30,6 +30,8 @@
|
||||
#include "ggml-common.h"
|
||||
#include "htp-ctx.h"
|
||||
#include "htp-ops.h"
|
||||
#include "htp-tensor.h"
|
||||
#include "hvx-quant.h"
|
||||
|
||||
#include "flash-attn-ops.h"
|
||||
#include "hvx-fa-kernels.h"
|
||||
@@ -85,12 +87,17 @@ struct htp_fa_context {
|
||||
uint8_t * spad_m;
|
||||
uint8_t * spad_a;
|
||||
|
||||
const struct htp_tensor * k;
|
||||
const struct htp_tensor * v;
|
||||
|
||||
uint64_t t_start;
|
||||
};
|
||||
|
||||
struct hmx_fa_context {
|
||||
const struct htp_ops_context * octx;
|
||||
const struct htp_tensor * sinks; // attention sinks (src[4]), NULL if absent
|
||||
const struct htp_tensor * k;
|
||||
const struct htp_tensor * v;
|
||||
bool pipeline; // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2
|
||||
uint32_t n_threads;
|
||||
|
||||
@@ -214,8 +221,8 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
|
||||
const uint32_t DV = nev0;
|
||||
|
||||
const size_t size_q_row = DK * ((q->type == HTP_TYPE_F32) ? 4 : 2);
|
||||
const size_t size_k_row = DK * sizeof(__fp16);
|
||||
const size_t size_v_row = DV * sizeof(__fp16);
|
||||
const size_t size_k_row = htp_tensor_get_row_size(k->type, DK);
|
||||
const size_t size_v_row = htp_tensor_get_row_size(v->type, DV);
|
||||
|
||||
// Scratchpad buffers for Q, K, V, Mask, and VKQ32 accumulator
|
||||
uint8_t * spad_q = factx->spad_q + factx->size_q_block * ith;
|
||||
@@ -364,6 +371,23 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
|
||||
uint8_t * v_base = dma_queue_pop(dma).dst; // V
|
||||
__fp16 * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M
|
||||
|
||||
if (factx->k->type == HTP_TYPE_Q8_0) {
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
|
||||
for (uint32_t r = 0; r < current_block_size; ++r) {
|
||||
__fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded);
|
||||
hvx_dequantize_row_q8_0_f16(row_k, row_k, DK);
|
||||
}
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
|
||||
}
|
||||
if (factx->v->type == HTP_TYPE_Q8_0) {
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
|
||||
for (uint32_t r = 0; r < current_block_size; ++r) {
|
||||
__fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded);
|
||||
hvx_dequantize_row_q8_0_f16(row_v, row_v, DV);
|
||||
}
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
|
||||
}
|
||||
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);
|
||||
|
||||
// Inner loop processing the block from VTCM
|
||||
@@ -625,6 +649,12 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
|
||||
|
||||
struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
|
||||
if (factx->k->type == HTP_TYPE_Q8_0) {
|
||||
for (uint32_t r = start; r < end; ++r) {
|
||||
__fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16));
|
||||
hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK);
|
||||
}
|
||||
}
|
||||
hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
|
||||
args->src_stride, start, end);
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
|
||||
@@ -673,6 +703,12 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
|
||||
|
||||
struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
|
||||
if (factx->v->type == HTP_TYPE_Q8_0) {
|
||||
for (uint32_t r = start; r < end; ++r) {
|
||||
__fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16));
|
||||
hvx_dequantize_row_q8_0_f16(row_v, row_v, 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));
|
||||
@@ -1809,6 +1845,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
memset(&factx, 0, sizeof(factx));
|
||||
factx.octx = octx;
|
||||
factx.sinks = octx->src[4]; // NULL if this op has no attention sinks
|
||||
factx.k = k;
|
||||
factx.v = v;
|
||||
factx.n_threads = kparams->n_threads;
|
||||
factx.DK = DK;
|
||||
factx.DV = DV;
|
||||
@@ -1853,10 +1891,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
// ======== 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 size_k_row = htp_tensor_get_row_size(k->type, DK);
|
||||
const size_t size_v_row = htp_tensor_get_row_size(v->type, DV);
|
||||
const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
|
||||
const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
|
||||
|
||||
// Build the VTCM layout once (shared with the host estimator) and place every
|
||||
// scratch buffer at its computed offset.
|
||||
@@ -2348,7 +2386,9 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
const struct htp_tensor * dst = octx->dst;
|
||||
|
||||
// Check support
|
||||
if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || k->type != HTP_TYPE_F16 || v->type != HTP_TYPE_F16) {
|
||||
if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) ||
|
||||
(k->type != HTP_TYPE_F16 && k->type != HTP_TYPE_Q8_0) ||
|
||||
(v->type != HTP_TYPE_F16 && v->type != HTP_TYPE_Q8_0)) {
|
||||
return HTP_STATUS_NO_SUPPORT;
|
||||
}
|
||||
|
||||
@@ -2364,6 +2404,8 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
|
||||
struct htp_fa_context factx;
|
||||
factx.octx = octx;
|
||||
factx.k = k;
|
||||
factx.v = v;
|
||||
|
||||
factx.t_start = HAP_perf_get_qtimer_count();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user