Skip to content

Commit 09b2345

Browse files
hipuddinghipudding
committed
cann: upgrade flash attention from V2 to V5 for sinks and asymmetric head support
Upgrade aclnnFusedInferAttentionScore from V2 to V5 (requires CANN 9.0.0+): - Attention sinks: pass src[4] as learnableSinkOptional to V5. Sinks are cast from F32 to BF16 (V5 requirement on 910B4). Force innerPrecise=0 when sinks are present. Sinks + ALiBi rejected. V5 learnableSink constraint: Q/K head dim 64/128/192, V head dim 64/128. - Asymmetric head sizes: removed support check rejection for DKQ != DV. V5 natively supports unequal head dim when pse_shift is not present. Padding logic updated to pad Q/K/V independently. Output dimensions use V's head dim instead of Q's. - logit_softcap: no aclnn interface supports it (confirmed V1-V5). Support check returns false, routing softcap models to non-FA fallback. - Fix dead code in FA support check: was checking src[1] (K) instead of src[0] (Q) for allowed query types, making the check always pass since K type was already validated above. - Fix BF16 output assert inconsistency: support check allowed BF16 output but implementation would abort. Add BF16 to assert and change output buffer/cast conditions from F32-specific to non-F16, so any non-F16 output type uses F16 intermediate buffer with proper cast. Support check changes: - Removed src[4] (sinks) rejection - Removed src[1]->ne[0] != src[2]->ne[0] (asymmetric head) rejection - Added sinks + ALiBi mutual exclusion check - Added learnableSink head dim constraint check - Added mask + asymmetric head dim rejection (V5 limitation) - Added runtime guards in FA implementation for same constraints test-backend-ops results on 910B4: 912 tests pass, 0 FAIL.
1 parent 0ef6f06 commit 09b2345

3 files changed

Lines changed: 101 additions & 10071 deletions

File tree

‎ggml/src/ggml-cann/aclnn_ops.cpp‎

Lines changed: 75 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
#include <aclnnop/aclnn_eq_tensor.h>
4444
#include <aclnnop/aclnn_exp.h>
4545
#include <aclnnop/aclnn_fill_scalar.h>
46-
#include <aclnnop/aclnn_fused_infer_attention_score_v2.h>
46+
#include <aclnnop/aclnn_fused_infer_attention_score_v5.h>
4747
#include <aclnnop/aclnn_ger.h>
4848
#include <aclnnop/aclnn_group_norm.h>
4949
#include <aclnnop/aclnn_gather_v2.h>
@@ -3860,6 +3860,7 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
38603860
ggml_tensor * src1 = dst->src[1]; // k, fp16 | B, N, S, D (uncont) -> B, S, N, D (cont)
38613861
ggml_tensor * src2 = dst->src[2]; // v, fp16 | B, N, S, D (uncont) -> B, S, N, D (cont)
38623862
ggml_tensor * src3 = dst->src[3]; // mask, fp16
3863+
ggml_tensor * src4 = dst->src[4]; // sinks, f32 [n_heads]
38633864

38643865
// B, N, S, D (uncont) -> B, S, N, D (cont)
38653866
int64_t src0_bsnd_ne[GGML_MAX_DIMS];
@@ -3932,22 +3933,28 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
39323933
acl_k_tensor = ggml_cann_create_tensor(src1, src1_bsnd_ne, src1_bsnd_nb, GGML_MAX_DIMS);
39333934
acl_v_tensor = ggml_cann_create_tensor(src2, src2_bsnd_ne, src2_bsnd_nb, GGML_MAX_DIMS);
39343935

3935-
// Step 2.5: Pad Q, K, V along head dimension if D is not a multiple of 16
3936-
// (required by FusedInferAttentionScoreV2)
3937-
const int64_t D = src0->ne[0];
3938-
const int64_t D_padded = GGML_PAD(D, 16);
3939-
const bool needs_padding = (D != D_padded);
3936+
// Step 2.5: Pad Q, K, V along head dimension if D is not a multiples of 16
3937+
// (required by FusedInferAttentionScoreV5)
3938+
const int64_t DQ = src0->ne[0];
3939+
const int64_t DV = src2->ne[0];
3940+
const int64_t DQ_padded = GGML_PAD(DQ, 16);
3941+
const int64_t DK_padded = GGML_PAD(src1->ne[0], 16);
3942+
const int64_t DV_padded = GGML_PAD(DV, 16);
3943+
const bool needs_padding = (DQ != DQ_padded) || (src1->ne[0] != DK_padded) || (DV != DV_padded);
39403944

39413945
ggml_cann_pool_alloc q_pad_allocator(ctx.pool());
39423946
ggml_cann_pool_alloc k_pad_allocator(ctx.pool());
39433947
ggml_cann_pool_alloc v_pad_allocator(ctx.pool());
39443948

39453949
if (needs_padding) {
3946-
int64_t paddings[] = { 0, D_padded - D, 0, 0, 0, 0, 0, 0 };
3947-
39483950
auto pad_fa_tensor = [&](acl_tensor_ptr & tensor, const int64_t * bsnd_ne,
3951+
int64_t D_orig, int64_t D_pad,
39493952
ggml_cann_pool_alloc & allocator) {
3950-
int64_t pad_ne[GGML_MAX_DIMS] = { D_padded, bsnd_ne[1], bsnd_ne[2], bsnd_ne[3] };
3953+
if (D_orig == D_pad) {
3954+
return;
3955+
}
3956+
int64_t paddings[] = { 0, D_pad - D_orig, 0, 0, 0, 0, 0, 0 };
3957+
int64_t pad_ne[GGML_MAX_DIMS] = { D_pad, bsnd_ne[1], bsnd_ne[2], bsnd_ne[3] };
39513958
size_t pad_nb[GGML_MAX_DIMS];
39523959
pad_nb[0] = faElemSize;
39533960
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
@@ -3961,13 +3968,13 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
39613968
tensor = std::move(padded);
39623969
};
39633970

3964-
pad_fa_tensor(acl_q_tensor, src0_bsnd_ne, q_pad_allocator);
3965-
pad_fa_tensor(acl_k_tensor, src1_bsnd_ne, k_pad_allocator);
3966-
pad_fa_tensor(acl_v_tensor, src2_bsnd_ne, v_pad_allocator);
3971+
pad_fa_tensor(acl_q_tensor, src0_bsnd_ne, DQ, DQ_padded, q_pad_allocator);
3972+
pad_fa_tensor(acl_k_tensor, src1_bsnd_ne, src1->ne[0], DK_padded, k_pad_allocator);
3973+
pad_fa_tensor(acl_v_tensor, src2_bsnd_ne, DV, DV_padded, v_pad_allocator);
39673974

3968-
src0_bsnd_ne[0] = D_padded;
3969-
src1_bsnd_ne[0] = D_padded;
3970-
src2_bsnd_ne[0] = D_padded;
3975+
src0_bsnd_ne[0] = DQ_padded;
3976+
src1_bsnd_ne[0] = DK_padded;
3977+
src2_bsnd_ne[0] = DV_padded;
39713978
}
39723979

39733980
// Step 3: create the PSEShift tensor if needed
@@ -4039,30 +4046,66 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
40394046
}
40404047
}
40414048

4042-
// Step 4: set the inputs for FusedInferAttention.
4049+
// Step 5: set the inputs for FusedInferAttention.
40434050
acl_tensor_list_ptr acl_k_tensor_list = ggml_cann_create_tensor_list(acl_k_tensor);
40444051
acl_tensor_list_ptr acl_v_tensor_list = ggml_cann_create_tensor_list(acl_v_tensor);
40454052

40464053
int64_t numHeads = src0->ne[2]; // N
40474054
int64_t numKeyValueHeads = src1->ne[2];
4055+
4056+
// Step 5.5: create learnableSink tensor if src4 (sinks) is present
4057+
// V5 learnableSink: dtype must be F16, shape = (Q_N) 1D tensor
4058+
acl_tensor_ptr acl_learnable_sink_tensor;
4059+
ggml_cann_pool_alloc sink_f16_allocator(ctx.pool());
4060+
if (src4 != nullptr) {
4061+
// V5 requires learnableSink dtype to be BF16
4062+
const aclDataType sinkDataType = ACL_BF16;
4063+
const size_t sinkElemSize = sizeof(uint16_t);
4064+
void * sink_f16_buffer = sink_f16_allocator.alloc(numHeads * sinkElemSize);
4065+
4066+
// aclnn_cast requires matching shapes, so cast via a 4D tensor
4067+
// matching src4's shape, then create a 1D view for V5
4068+
int64_t sink_ne_4d[GGML_MAX_DIMS] = { numHeads, 1, 1, 1 };
4069+
size_t sink_nb_4d[GGML_MAX_DIMS];
4070+
sink_nb_4d[0] = sinkElemSize;
4071+
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
4072+
sink_nb_4d[i] = sink_nb_4d[i - 1] * sink_ne_4d[i - 1];
4073+
}
4074+
acl_tensor_ptr sink_bf16_4d =
4075+
ggml_cann_create_tensor(sink_f16_buffer, sinkDataType, sinkElemSize, sink_ne_4d, sink_nb_4d, GGML_MAX_DIMS);
4076+
4077+
acl_tensor_ptr acl_src4_f32_tensor = ggml_cann_create_tensor(src4);
4078+
aclnn_cast(ctx, acl_src4_f32_tensor.get(), sink_bf16_4d.get(), sinkDataType);
4079+
4080+
// create 1D view for V5 learnableSink (same buffer, 1D shape)
4081+
int64_t sink_ne_1d[1] = { numHeads };
4082+
size_t sink_nb_1d[1] = { sinkElemSize };
4083+
acl_learnable_sink_tensor =
4084+
ggml_cann_create_tensor(sink_f16_buffer, sinkDataType, sinkElemSize, sink_ne_1d, sink_nb_1d, 1);
4085+
}
4086+
40484087
// double scaleValue = 1 / sqrt(src0->ne[0]); // 1/sqrt(d)
40494088
int64_t preTokens = 65535;
40504089
int64_t nextTokens = 65535;
40514090
char layout[5] = { 'B', 'S', 'N', 'D', 0 };
40524091
int64_t sparseMode = 0;
4053-
int64_t innerPrecise = (src0->ne[1] == 1) ? 0 : 2;
4092+
// learnableSink requires innerPrecise = 0
4093+
int64_t innerPrecise = (src0->ne[1] == 1 || src4 != nullptr) ? 0 : 2;
40544094
int64_t blockSize = 0;
40554095
int64_t antiquantMode = 0;
40564096
bool softmaxLseFlag = false;
40574097
int64_t keyAntiquantMode = 0;
40584098
int64_t valueAntiquantMode = 0;
4099+
int64_t queryQuantMode = 0;
4100+
int64_t pseType = 0;
40594101

4060-
GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16);
4102+
GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_BF16);
40614103
acl_tensor_ptr fa_dst_tensor;
40624104
ggml_cann_pool_alloc out_f16_allocator(ctx.pool());
4063-
if (dst->type == GGML_TYPE_F32 || needs_padding) {
4064-
int64_t * out_f16_ne = src0_bsnd_ne;
4065-
size_t out_f16_nb[GGML_MAX_DIMS];
4105+
if (dst->type != GGML_TYPE_F16 || needs_padding) {
4106+
// output shape: B, S, N, DV (V's head dim, not Q's)
4107+
int64_t out_f16_ne[GGML_MAX_DIMS] = { src2_bsnd_ne[0], src0_bsnd_ne[1], src0_bsnd_ne[2], src0_bsnd_ne[3] };
4108+
size_t out_f16_nb[GGML_MAX_DIMS];
40664109
out_f16_nb[0] = faElemSize;
40674110
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
40684111
out_f16_nb[i] = out_f16_nb[i - 1] * out_f16_ne[i - 1];
@@ -4076,7 +4119,7 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
40764119
fa_dst_tensor = ggml_cann_create_tensor(dst);
40774120
}
40784121

4079-
GGML_CANN_CALL_ACLNN_OP(ctx, FusedInferAttentionScoreV2, acl_q_tensor.get(), acl_k_tensor_list.get(),
4122+
GGML_CANN_CALL_ACLNN_OP(ctx, FusedInferAttentionScoreV5, acl_q_tensor.get(), acl_k_tensor_list.get(),
40804123
acl_v_tensor_list.get(), // q, k, v
40814124
bcast_pse_tensor.get(), nullptr, // pse, mask
40824125
nullptr, nullptr, // actSeqLen, actSeqLenkv
@@ -4088,6 +4131,9 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
40884131
nullptr, nullptr, // kAntiquantScale, kAntiQuantOffset
40894132
nullptr, nullptr, // vAntiquantScale, vAntiQuantOffset
40904133
nullptr, nullptr, nullptr, // kSharedPrefix, vSharedPrefix, actSharedLen
4134+
nullptr, nullptr, nullptr, // queryRope, keyRope, keyRopeAntiquantScale
4135+
nullptr, acl_learnable_sink_tensor.get(), // dequantScaleQuery, learnableSink
4136+
nullptr, nullptr, // qStartIdx, kvStartIdx
40914137
numHeads, scaleValue, // heads, scaleValue
40924138
preTokens, nextTokens, // preTokens, nextTokens
40934139
layout, // inputLayout
@@ -4096,16 +4142,17 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
40964142
blockSize, antiquantMode, // blockSize, antiquantMode
40974143
softmaxLseFlag, // softmaxLseFlag
40984144
keyAntiquantMode, valueAntiquantMode, // keyAntiqMode, valueAntiqMode
4145+
queryQuantMode, pseType, // queryQuantMode, pseType
40994146
fa_dst_tensor.get(), // attentionOut
41004147
nullptr // softmaxLse
41014148
);
41024149

4103-
// Step 6: post-processing — slice padded output and/or cast to f32
4150+
// Step 6: post-processing — slice padded output and/or cast to dst type
41044151
if (needs_padding) {
41054152
ggml_cann_pool_alloc sliced_f16_allocator(ctx.pool());
41064153

4107-
if (dst->type == GGML_TYPE_F32) {
4108-
int64_t sliced_ne[GGML_MAX_DIMS] = { D, src0_bsnd_ne[1], src0_bsnd_ne[2], src0_bsnd_ne[3] };
4154+
if (dst->type != GGML_TYPE_F16) {
4155+
int64_t sliced_ne[GGML_MAX_DIMS] = { DV, src0_bsnd_ne[1], src0_bsnd_ne[2], src0_bsnd_ne[3] };
41094156
size_t sliced_nb[GGML_MAX_DIMS];
41104157
sliced_nb[0] = faElemSize;
41114158
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
@@ -4117,16 +4164,16 @@ void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst
41174164
sliced_ne, sliced_nb, GGML_MAX_DIMS);
41184165

41194166
GGML_CANN_CALL_ACLNN_OP(ctx, Slice, fa_dst_tensor.get(),
4120-
(int64_t) -1, (int64_t) 0, D, (int64_t) 1, sliced_f16_tensor.get());
4167+
(int64_t) -1, (int64_t) 0, DV, (int64_t) 1, sliced_f16_tensor.get());
41214168

41224169
acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst);
41234170
aclnn_cast(ctx, sliced_f16_tensor.get(), acl_dst_tensor.get(), ggml_cann_type_mapping(dst->type));
41244171
} else {
41254172
acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst);
41264173
GGML_CANN_CALL_ACLNN_OP(ctx, Slice, fa_dst_tensor.get(),
4127-
(int64_t) -1, (int64_t) 0, D, (int64_t) 1, acl_dst_tensor.get());
4174+
(int64_t) -1, (int64_t) 0, DV, (int64_t) 1, acl_dst_tensor.get());
41284175
}
4129-
} else if (dst->type == GGML_TYPE_F32) {
4176+
} else if (dst->type != GGML_TYPE_F16) {
41304177
acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst);
41314178
aclnn_cast(ctx, fa_dst_tensor.get(), acl_dst_tensor.get(), ggml_cann_type_mapping(dst->type));
41324179
}

‎ggml/src/ggml-cann/ggml-cann.cpp‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2661,23 +2661,41 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
26612661
if (op->src[1]->type != GGML_TYPE_F16 || op->src[2]->type != GGML_TYPE_F16) {
26622662
return false;
26632663
}
2664-
if (op->src[1]->type != GGML_TYPE_F16 && op->src[1]->type != GGML_TYPE_F32 &&
2665-
op->src[1]->type != GGML_TYPE_BF16) {
2664+
if (op->src[0]->type != GGML_TYPE_F16 && op->src[0]->type != GGML_TYPE_F32 &&
2665+
op->src[0]->type != GGML_TYPE_BF16) {
26662666
return false;
26672667
}
26682668
if (op->type != GGML_TYPE_F16 && op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_BF16) {
26692669
return false;
26702670
}
2671-
// TODO: support attention sinks [TAG_ATTN_SINKS]
2672-
if (op->src[4]) {
2671+
// learnableSink is not compatible with ALiBi
2672+
float scale = 0.0f;
2673+
float maxBias = 0.0f;
2674+
float logitSoftcap = 0.0f;
2675+
memcpy(&scale, (const float *) (op->op_params) + 0, sizeof(float));
2676+
memcpy(&maxBias, (const float *) (op->op_params) + 1, sizeof(float));
2677+
memcpy(&logitSoftcap,(const float *) (op->op_params) + 2, sizeof(float));
2678+
if (op->src[4] && maxBias > 0.0f) {
26732679
return false;
26742680
}
2681+
// V5 learnableSink: Q/K head dim only 64/128/192, V head dim only 64/128
2682+
if (op->src[4]) {
2683+
int64_t dq = op->src[0]->ne[0];
2684+
int64_t dv = op->src[2]->ne[0];
2685+
if ((dq != 64 && dq != 128 && dq != 192) ||
2686+
(dv != 64 && dv != 128)) {
2687+
return false;
2688+
}
2689+
}
2690+
// V5 asymmetric head dim (dq != dv) constraints:
2691+
// - With mask (pse_shift): dq must equal dv
2692+
// - Without mask: only dq=192, dv=128 is supported
26752693
if (op->src[1]->ne[0] != op->src[2]->ne[0]) {
2676-
// different head sizes of K and V are not supported yet
2677-
return false;
2694+
if (op->src[3] ||
2695+
op->src[1]->ne[0] != 192 || op->src[2]->ne[0] != 128) {
2696+
return false;
2697+
}
26782698
}
2679-
float logitSoftcap = 0.0f;
2680-
memcpy(&logitSoftcap, (const float *) (op->op_params) + 2, sizeof(float));
26812699
if (logitSoftcap != 0.0f) {
26822700
return false;
26832701
}

0 commit comments

Comments
 (0)