Commit 78a8ae14b for llama.cpp
commit 78a8ae14be559b45f5f13ddee82c188b683d5d8d
Author: Ehsan Bateni <ebateni@qti.qualcomm.com>
Date: Sat Oct 24 15:19:32 2026 -0400
hexagon: KV-head team split for multicore HMX flash attention (#31055)
* hexagon: KV-head team split for HMX flash attention in multicore
* hexagon: merge FA_HEAD_SPLIT and FA_SPLIT into single GGML_HEXAGON_FA_SPLIT flag
0 = off (token-split only); 1 = head-split, plus team-split when
n_kv_heads < N_cores so stranded cores still participate (e.g. gemma-4
with 2 KV heads on 4 cores). Where heads divide evenly across cores,
plain head-split is used.
diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index 0821bce16..c7e8ea85d 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -133,7 +133,7 @@ size_t opt_mbuf_total = 0; // total buffer space limit (
int opt_mm_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
-int opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
+int opt_fa_split = 1; // 1 = head-split + team-split when n_kv_heads < N_cores (default), 0 = off
int opt_gdn_select = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
int opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -3393,7 +3393,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
const char * str_nhmx = getenv("GGML_HEXAGON_NHMX");
const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
- const char * str_fa_head_split = getenv("GGML_HEXAGON_FA_HEAD_SPLIT");
+ const char * str_fa_split = getenv("GGML_HEXAGON_FA_SPLIT");
const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -3458,7 +3458,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
opt_nhmx = str_nhmx ? atoi(str_nhmx) : opt_nhmx;
opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select;
opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select;
- opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split) : opt_fa_head_split;
+ opt_fa_split = str_fa_split ? atoi(str_fa_split) : opt_fa_split;
opt_gdn_select = str_gdn_select ? atoi(str_gdn_select) : opt_gdn_select;
opt_ar_select = str_ar_select ? atoi(str_ar_select) : opt_ar_select;
opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter) : opt_ar_scatter;
diff --git a/ggml/src/ggml-hexagon/htp-common.h b/ggml/src/ggml-hexagon/htp-common.h
index 71a690a16..e51340d8a 100644
--- a/ggml/src/ggml-hexagon/htp-common.h
+++ b/ggml/src/ggml-hexagon/htp-common.h
@@ -28,7 +28,7 @@ extern size_t opt_mbuf_total;
extern int opt_mm_select;
extern int opt_fa_select;
-extern int opt_fa_head_split;
+extern int opt_fa_split;
extern int opt_gdn_select;
extern int opt_ar_select;
extern int opt_ar_scatter;
diff --git a/ggml/src/ggml-hexagon/htp-kparams.cpp b/ggml/src/ggml-hexagon/htp-kparams.cpp
index 589a8accf..a8f01a4c4 100644
--- a/ggml/src/ggml-hexagon/htp-kparams.cpp
+++ b/ggml/src/ggml-hexagon/htp-kparams.cpp
@@ -397,9 +397,20 @@ static bool htp_kparams_solve_flash_attn(
kparams->scale = scale;
kparams->max_bias = max_bias;
kparams->logit_softcap = logit_softcap;
- kparams->head_split = (opt_fa_head_split != 0) ? 1 : 0;
+ kparams->head_split = (opt_fa_split != 0) ? 1 : 0;
+ kparams->flags = 0;
kparams->G = G;
+ // team-split only when a team has >1 member (n_kv_heads < N_cores)
+ uint32_t fa_teams = 1, fa_team_size = 1;
+ const uint32_t mdev_count = htp_session_mdev_count(sess);
+ if (opt_fa_split && mdev_count > 1) {
+ htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &fa_teams, &fa_team_size);
+ if (fa_team_size > 1) {
+ kparams->flags |= HTP_FA_FLAG_TEAM_SPLIT;
+ }
+ }
+
const uint32_t n_head = q->ne[2];
kparams->n_head_log2 = 1u << (uint32_t) std::floor(std::log2(n_head));
kparams->m0 = std::pow(2.0f, -(max_bias) / kparams->n_head_log2);
@@ -409,7 +420,7 @@ static bool htp_kparams_solve_flash_attn(
const uint32_t DK_pad = hex_round_up(DK, 64);
const uint32_t DV_pad = hex_round_up(DV, 64);
size_t Br = 0, Bc = 0;
- int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
+ int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head, fa_team_size);
if (ret == 0) {
kparams->kernel_type = HTP_FA_KERNEL_HMX;
if (logit_softcap == 0.0f) {
diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c
index fc5664185..ee258172b 100644
--- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c
+++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c
@@ -2094,7 +2094,29 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
const bool can_split = htp_tensor_can_row_partition(dst, dst_e_size);
- if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
+ if ((kparams->flags & HTP_FA_FLAG_TEAM_SPLIT) && can_split) {
+ // KV heads split across teams, Q-blocks split within a team
+ uint32_t n_teams, team_size;
+ htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &n_teams, &team_size);
+
+ const uint32_t team = mdev_idx / team_size;
+ const uint32_t member = mdev_idx - team * team_size;
+ const uint32_t kv_per_team = n_kv_heads / n_teams;
+ kv_head_min = team * kv_per_team;
+ kv_head_max = kv_head_min + kv_per_team;
+
+ const uint32_t base = n_q_blocks / team_size;
+ const uint32_t rem = n_q_blocks - base * team_size;
+ const uint32_t b0 = member * base + MIN(member, rem);
+ const uint32_t nblk = base + (member < rem ? 1 : 0);
+
+ if (nblk == 0) {
+ return HTP_STATUS_OK;
+ }
+
+ q_start_min = b0 * Br;
+ q_start_max = MIN((b0 + nblk) * Br, neq1);
+ } else if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
const uint32_t kv_per_core = n_kv_heads / mdev_count;
kv_head_min = mdev_idx * kv_per_core;
kv_head_max = kv_head_min + kv_per_core;
diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h
index a3d380501..f2654f0a8 100644
--- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h
+++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h
@@ -32,9 +32,20 @@ enum htp_fa_kernel_type {
HTP_FA_KERNEL_HMX
};
+// Multicore: gcd(n_kv_heads, N) teams own the KV heads, team members share the Q-blocks
+#define HTP_FA_FLAG_TEAM_SPLIT 0x01
+
+static inline void htp_fa_mdev_teams(uint32_t n_kv_heads, uint32_t mdev_count, bool head_split,
+ uint32_t * n_teams, uint32_t * team_size) {
+ const uint32_t t = (head_split && n_kv_heads > 0 && mdev_count > 1) ? hex_gcd_u32(n_kv_heads, mdev_count) : 1;
+ *n_teams = t;
+ *team_size = mdev_count > 1 ? mdev_count / t : 1;
+}
+
struct htp_fa_kernel_params {
uint8_t kernel_type; // enum htp_fa_kernel_type
uint8_t head_split; // 1 = partition by KV heads in multicore, 0 = token partition
+ uint8_t flags; // HTP_FA_FLAG_*
uint8_t n_threads; // Number of threads to run
// Common parameters
@@ -321,7 +332,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
size_t n_threads,
bool is_q_fp32,
bool has_sinks,
- size_t n_heads) {
+ size_t n_heads,
+ size_t q_split) {
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
@@ -361,7 +373,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
const size_t vecs_per_t = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
const size_t c_iter_actual = c_iter_base + c_softmax * vecs_per_t;
- const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
+ const size_t q_blocks_nsp = (q_blocks + q_split - 1) / (q_split ? q_split : 1);
+ const size_t cost = q_blocks_nsp * (c_q_fixed + kv_blocks * c_iter_actual);
const size_t mn = Br * Bc;
if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 5b583aa4c..92808cd8c 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12051,6 +12051,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
+ // KV-head counts that do not divide the device count
+ for (int hs : { 128, 256, 512, }) {
+ for (int nh : { 1, 2, 6, }) {
+ for (int nr2 : { 4, 8, }) {
+ for (int nb : { 75, 1024, }) {
+ test_cases.emplace_back(new test_flash_attn_ext(hs, hs, nh, {nr2, 1}, 1024, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+ }
+ }
+ }
+ }
+
for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) {
for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) {
if (hsk != 96 && hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;