Tidied the code a bit to match existing style

This commit is contained in:
RedToasty
2026-06-19 18:52:39 +01:00
parent 370c99dc53
commit b90c6c6300
+4 -5
View File
@@ -1039,11 +1039,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
for (size_t s = 0; s < src_ss[i].n_segments; s++) {
sum += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s];
}
const bool is_flash_attn_gqa_kv =
tensor->op == GGML_OP_FLASH_ATTN_EXT && (i == 1 || i == 2) &&
tensor->src[i]->ne[src_ss[i].axis] != tensor->ne[split_state.axis];
if (is_flash_attn_gqa_kv) {
// GQA/MQA splits Q/output over query heads, but K/V over fewer KV heads.
// For FLASH_ATTN_EXT, src[1] and src[2] are K and V. With GQA/MQA they are split
// over fewer KV heads than Q/output, so the generic split ratio check does not apply.
if (tensor->op == GGML_OP_FLASH_ATTN_EXT && (i == 1 || i == 2) &&
tensor->src[i]->ne[src_ss[i].axis] != tensor->ne[split_state.axis]) {
continue;
}
GGML_ASSERT(split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis]