已合并
sdpa use component impl when the last dim of input is greater than 256 #5578
XianglongZeng创建于 26 天前
sdpa use component impl when the last dim of input is greater than 256 #5578
已合并
共 2 个文件变更+158-191
| @@ -33,41 +33,34 @@ const static int64_t LEFT_UP_CAUSAL = 2; | |||
| 33 | const static int64_t ATTN_MASK_DIM_TWO = 2; | 33 | const static int64_t ATTN_MASK_DIM_TWO = 2; |
| 34 | const static int64_t ATTN_MASK_DIM_FOUR = 4; | 34 | const static int64_t ATTN_MASK_DIM_FOUR = 4; |
| 35 | 35 | ||
| 36 | - | 36 | +inline void validate_sdpa_input(const at::Tensor &query, const at::Tensor &key, const at::Tensor &value, |
| 37 | -inline void validate_sdpa_input( | 37 | + const c10::optional<at::Tensor> &attn_mask) { |
| 38 | - const at::Tensor &query, | ||
| 39 | - const at::Tensor &key, | ||
| 40 | - const at::Tensor &value, | ||
| 41 | - const c10::optional<at::Tensor> &attn_mask) | ||
| 42 | -{ | ||
| 43 | TORCH_CHECK(query.dtype() == key.dtype() && query.dtype() == value.dtype(), | 38 | TORCH_CHECK(query.dtype() == key.dtype() && query.dtype() == value.dtype(), |
| 44 | - "Expected query, key, and value to have the same dtype, but got query.dtype: ", | 39 | + "Expected query, key, and value to have the same dtype, but got query.dtype: ", query.dtype(), |
| 45 | - query.dtype(), " key.dtype: ", key.dtype(), " and value.dtype: ", value.dtype(), " instead." + | 40 | + " key.dtype: ", key.dtype(), " and value.dtype: ", value.dtype(), |
| 46 | - OPS_ERROR(ErrCode::NOT_SUPPORT)); | 41 | + " instead." + OPS_ERROR(ErrCode::NOT_SUPPORT)); |
| 47 | TORCH_CHECK(query.device() == key.device() && query.device() == value.device(), | 42 | TORCH_CHECK(query.device() == key.device() && query.device() == value.device(), |
| 48 | - "Expected query, key, and value to have the same device type, but got query.device: ", | 43 | + "Expected query, key, and value to have the same device type, but got query.device: ", query.device(), |
| 49 | - query.device(), " key.device: ", key.device(), " and value.device: ", value.device(), " instead." + | 44 | + " key.device: ", key.device(), " and value.device: ", value.device(), |
| 50 | - OPS_ERROR(ErrCode::NOT_SUPPORT)); | 45 | + " instead." + OPS_ERROR(ErrCode::NOT_SUPPORT)); |
| 51 | TORCH_CHECK(query.dim() >= 2 && key.dim() >= 2 && value.dim() >= 2, | 46 | TORCH_CHECK(query.dim() >= 2 && key.dim() >= 2 && value.dim() >= 2, |
| 52 | - "Expected query, key, and value to all be at least 2 dimensional, but got query.dim: ", | 47 | + "Expected query, key, and value to all be at least 2 dimensional, but got query.dim: ", query.dim(), |
| 53 | - query.dim(), " key.dim: ", key.dim(), " and value.dim: ", value.dim(), " instead." + | 48 | + " key.dim: ", key.dim(), " and value.dim: ", value.dim(), " instead." + OPS_ERROR(ErrCode::NOT_SUPPORT)); |
| 54 | - OPS_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 55 | if (attn_mask.has_value()) { | 49 | if (attn_mask.has_value()) { |
| 56 | auto mask_dtype = attn_mask->dtype(); | 50 | auto mask_dtype = attn_mask->dtype(); |
| 57 | TORCH_CHECK(mask_dtype == at::kBool || mask_dtype == query.dtype(), | 51 | TORCH_CHECK(mask_dtype == at::kBool || mask_dtype == query.dtype(), |
| 58 | - "Expected attn_mask dtype to be bool or to match query dtype, but got attn_mask.dtype: ", | 52 | + "Expected attn_mask dtype to be bool or to match query dtype, but got attn_mask.dtype: ", mask_dtype, |
| 59 | - mask_dtype, " and query.dtype: ", query.dtype(), " instead." + OPS_ERROR(ErrCode::NOT_SUPPORT)); | 53 | + " and query.dtype: ", query.dtype(), " instead." + OPS_ERROR(ErrCode::NOT_SUPPORT)); |
| 60 | TORCH_CHECK(!query.is_nested() && !key.is_nested(), | 54 | TORCH_CHECK(!query.is_nested() && !key.is_nested(), |
| 61 | "Scaled_dot_product_attention: Nested tensors for query / key are not supported " | 55 | "Scaled_dot_product_attention: Nested tensors for query / key are not supported " |
| 62 | - "when an explicit attn_mask is set" + OPS_ERROR(ErrCode::NOT_SUPPORT)); | 56 | + "when an explicit attn_mask is set" + |
| 57 | + OPS_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 63 | } | 58 | } |
| 64 | return; | 59 | return; |
| 65 | } | 60 | } |
| 66 | 61 | ||
| 67 | c10::optional<at::Tensor> convert_boolean_attn_mask_math( | 62 | c10::optional<at::Tensor> convert_boolean_attn_mask_math( |
| 68 | - const c10::optional<at::Tensor> &attn_mask, | 63 | + const c10::optional<at::Tensor> &attn_mask, caffe2::TypeMeta dtype) { |
| 69 | - caffe2::TypeMeta dtype) | ||
| 70 | -{ | ||
| 71 | if (!attn_mask.has_value()) { | 64 | if (!attn_mask.has_value()) { |
| 72 | return c10::nullopt; | 65 | return c10::nullopt; |
| 73 | } | 66 | } |
| @@ -80,16 +73,12 @@ c10::optional<at::Tensor> convert_boolean_attn_mask_math( | |||
| 80 | } | 73 | } |
| 81 | 74 | ||
| 82 | c10::optional<at::Tensor> convert_boolean_attn_mask( | 75 | c10::optional<at::Tensor> convert_boolean_attn_mask( |
| 83 | - const at::Tensor &query, | 76 | + const at::Tensor &query, const c10::optional<at::Tensor> &attn_mask, bool is_causal) { |
| 84 | - const c10::optional<at::Tensor> &attn_mask, | ||
| 85 | - bool is_causal) | ||
| 86 | -{ | ||
| 87 | if (!attn_mask.has_value() && !is_causal) { | 77 | if (!attn_mask.has_value() && !is_causal) { |
| 88 | return c10::nullopt; | 78 | return c10::nullopt; |
| 89 | } | 79 | } |
| 90 | if (is_causal) { | 80 | if (is_causal) { |
| 91 | - TORCH_CHECK(!attn_mask.has_value(), | 81 | + TORCH_CHECK(!attn_mask.has_value(), "The attn_mask should be none when is_causal is true, but got ", |
| 92 | - "The attn_mask should be none when is_causal is true, but got ", | ||
| 93 | attn_mask.has_value(), "-value"); | 82 | attn_mask.has_value(), "-value"); |
| 94 | at::Tensor atten_mask_shape = at::ones({ATTENMASK_LIMIT, ATTENMASK_LIMIT}, query.options().dtype(at::kBool)); | 83 | at::Tensor atten_mask_shape = at::ones({ATTENMASK_LIMIT, ATTENMASK_LIMIT}, query.options().dtype(at::kBool)); |
| 95 | auto new_attn_mask = at::triu(atten_mask_shape, 1); | 84 | auto new_attn_mask = at::triu(atten_mask_shape, 1); |
| @@ -100,26 +89,17 @@ c10::optional<at::Tensor> convert_boolean_attn_mask( | |||
| 100 | return atten_mask; | 89 | return atten_mask; |
| 101 | } | 90 | } |
| 102 | 91 | ||
| 103 | -inline c10::SymFloat calculate_scale( | 92 | +inline c10::SymFloat calculate_scale(const at::Tensor &query, c10::optional<double> scale) { |
| 104 | - const at::Tensor &query, | 93 | + const auto softmax_scale = |
| 105 | - c10::optional<double> scale) | 94 | + scale.has_value() ? scale.value() : (c10::SymFloat(1.0) / (c10::SymFloat(query.sym_size(-1)).sqrt())); |
| 106 | -{ | ||
| 107 | - const auto softmax_scale = scale.has_value() | ||
| 108 | - ? scale.value() | ||
| 109 | - : (c10::SymFloat(1.0) / (c10::SymFloat(query.sym_size(-1)).sqrt())); | ||
| 110 | return c10::SymFloat(softmax_scale); | 95 | return c10::SymFloat(softmax_scale); |
| 111 | } | 96 | } |
| 112 | 97 | ||
| 113 | -static bool can_broadcast_to(int64_t mask_size, int64_t target_size) | 98 | +static bool can_broadcast_to(int64_t mask_size, int64_t target_size) { |
| 114 | -{ | ||
| 115 | return mask_size == target_size || mask_size == 1; | 99 | return mask_size == target_size || mask_size == 1; |
| 116 | } | 100 | } |
| 117 | 101 | ||
| 118 | -static bool can_expand_attn_mask_for_npu( | 102 | +static bool can_expand_attn_mask_for_npu(const at::Tensor &mask, const at::Tensor &query, const at::Tensor &key) { |
| 119 | - const at::Tensor &mask, | ||
| 120 | - const at::Tensor &query, | ||
| 121 | - const at::Tensor &key) | ||
| 122 | -{ | ||
| 123 | int64_t B = query.size(0), N = query.size(1), Sq = query.size(2), Skv = key.size(2); | 103 | int64_t B = query.size(0), N = query.size(1), Sq = query.size(2), Skv = key.size(2); |
| 124 | int64_t target_shape[ATTN_MASK_DIM_FOUR] = {B, N, Sq, Skv}; | 104 | int64_t target_shape[ATTN_MASK_DIM_FOUR] = {B, N, Sq, Skv}; |
| 125 | if (mask.dim() < ATTN_MASK_DIM_TWO || mask.dim() > ATTN_MASK_DIM_FOUR) { | 105 | if (mask.dim() < ATTN_MASK_DIM_TWO || mask.dim() > ATTN_MASK_DIM_FOUR) { |
| @@ -135,10 +115,7 @@ static bool can_expand_attn_mask_for_npu( | |||
| 135 | } | 115 | } |
| 136 | 116 | ||
| 137 | static c10::optional<at::Tensor> preprocess_attn_mask_for_npu( | 117 | static c10::optional<at::Tensor> preprocess_attn_mask_for_npu( |
| 138 | - const c10::optional<at::Tensor> &attn_mask, | 118 | + const c10::optional<at::Tensor> &attn_mask, const at::Tensor &query, const at::Tensor &key) { |
| 139 | - const at::Tensor &query, | ||
| 140 | - const at::Tensor &key) | ||
| 141 | -{ | ||
| 142 | if (!attn_mask.has_value()) { | 119 | if (!attn_mask.has_value()) { |
| 143 | return c10::nullopt; | 120 | return c10::nullopt; |
| 144 | } | 121 | } |
| @@ -157,7 +134,7 @@ static c10::optional<at::Tensor> preprocess_attn_mask_for_npu( | |||
| 157 | } else if (mask.dim() == ATTN_MASK_DIM_FOUR) { | 134 | } else if (mask.dim() == ATTN_MASK_DIM_FOUR) { |
| 158 | supported_shape = mask.size(2) == Sq && mask.size(3) == Skv && | 135 | supported_shape = mask.size(2) == Sq && mask.size(3) == Skv && |
| 159 | ((mask.size(0) == B && (mask.size(1) == N || mask.size(1) == 1)) || | 136 | ((mask.size(0) == B && (mask.size(1) == N || mask.size(1) == 1)) || |
| 160 | - (mask.size(0) == 1 && mask.size(1) == 1)); | 137 | + (mask.size(0) == 1 && mask.size(1) == 1)); |
| 161 | } | 138 | } |
| 162 | if (supported_shape) { | 139 | if (supported_shape) { |
| 163 | return mask; | 140 | return mask; |
| @@ -166,10 +143,7 @@ static c10::optional<at::Tensor> preprocess_attn_mask_for_npu( | |||
| 166 | } | 143 | } |
| 167 | 144 | ||
| 168 | static int64_t calculate_inner_precise_for_fa( | 145 | static int64_t calculate_inner_precise_for_fa( |
| 169 | - const at::Tensor &query, | 146 | + const at::Tensor &query, const c10::optional<at::Tensor> &attn_mask, bool is_causal) { |
| 170 | - const c10::optional<at::Tensor> &attn_mask, | ||
| 171 | - bool is_causal) | ||
| 172 | -{ | ||
| 173 | if (!is_causal && query.size(2) > 1 && attn_mask.has_value()) { | 147 | if (!is_causal && query.size(2) > 1 && attn_mask.has_value()) { |
| 174 | return 2; | 148 | return 2; |
| 175 | } | 149 | } |
| @@ -178,37 +152,32 @@ static int64_t calculate_inner_precise_for_fa( | |||
| 178 | 152 | ||
| 179 | 153 | ||
| 180 | 154 | ||
| 181 | -at::Tensor scaled_dot_product_attention( | 155 | +at::Tensor scaled_dot_product_attention(const at::Tensor &query, const at::Tensor &key, const at::Tensor &value, |
| 182 | - const at::Tensor &query, | 156 | + const c10::optional<at::Tensor> &attn_mask, double dropout_p, bool is_causal, c10::optional<double> scale) { |
| 183 | - const at::Tensor &key, | ||
| 184 | - const at::Tensor &value, | ||
| 185 | - const c10::optional<at::Tensor> &attn_mask, | ||
| 186 | - double dropout_p, | ||
| 187 | - bool is_causal, | ||
| 188 | - c10::optional<double> scale) | ||
| 189 | -{ | ||
| 190 | validate_sdpa_input(query, key, value, attn_mask); | 157 | validate_sdpa_input(query, key, value, attn_mask); |
| 191 | 158 | ||
| 192 | static auto compatible_impl = at_npu::native::env::CheckCompatibleImpl(); | 159 | static auto compatible_impl = at_npu::native::env::CheckCompatibleImpl(); |
| 193 | bool force_math = false; | 160 | bool force_math = false; |
| 194 | if (compatible_impl) { | 161 | if (compatible_impl) { |
| 195 | - auto& ctx = at::globalContext(); | 162 | + auto &ctx = at::globalContext(); |
| 196 | force_math = ctx.userEnabledMathSDP() && !ctx.userEnabledFlashSDP(); | 163 | force_math = ctx.userEnabledMathSDP() && !ctx.userEnabledFlashSDP(); |
| 197 | if (query.scalar_type() == at::kFloat && | 164 | if (query.scalar_type() == at::kFloat && |
| 198 | (query.size(-1) % 4 != 0 || key.size(-1) % 4 != 0 || value.size(-1) % 4 != 0)) { | 165 | (query.size(-1) % 4 != 0 || key.size(-1) % 4 != 0 || value.size(-1) % 4 != 0)) { |
| 199 | force_math = true; | 166 | force_math = true; |
| 200 | } | 167 | } |
| 168 | + if (query.size(-1) > 256 || key.size(-1) > 256 || value.size(-1) > 256) { | ||
| 169 | + force_math = true; | ||
| 170 | + } | ||
| 201 | } | 171 | } |
| 202 | 172 | ||
| 203 | if (!force_math && | 173 | if (!force_math && |
| 204 | - (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16 || query.scalar_type() == at::kFloat) && | 174 | + (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16 || |
| 175 | + query.scalar_type() == at::kFloat) && | ||
| 205 | ((attn_mask.has_value() && attn_mask->dtype() == at::kBool) || !attn_mask.has_value()) && | 176 | ((attn_mask.has_value() && attn_mask->dtype() == at::kBool) || !attn_mask.has_value()) && |
| 206 | - query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM && | 177 | + query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM && query.size(1) <= N_LIMIT && |
| 207 | - query.size(1) <= N_LIMIT && query.size(3) <= D_LIMIT && key.size(1) <= N_LIMIT && | 178 | + query.size(3) <= D_LIMIT && key.size(1) <= N_LIMIT && query.size(1) % key.size(1) == 0 && |
| 208 | - query.size(1) % key.size(1) == 0 && query.size(1) / key.size(1) > 0 && | 179 | + query.size(1) / key.size(1) > 0 && c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { |
| 209 | - c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { | 180 | + auto processed_mask = preprocess_attn_mask_for_npu(attn_mask, query, key); |
| 210 | - auto processed_mask = preprocess_attn_mask_for_npu( | ||
| 211 | - attn_mask, query, key); | ||
| 212 | 181 | ||
| 213 | /* The implementation of the NPU FlashAttention fusion operator without grad constraints: | 182 | /* The implementation of the NPU FlashAttention fusion operator without grad constraints: |
| 214 | 1. The FA operator must be registered. GetSocVersion api is provisional, IsExistOp aclnn api will apply. | 183 | 1. The FA operator must be registered. GetSocVersion api is provisional, IsExistOp aclnn api will apply. |
| @@ -227,100 +196,92 @@ at::Tensor scaled_dot_product_attention( | |||
| 227 | c10::optional<at::Tensor> nulltensor = c10::nullopt; | 196 | c10::optional<at::Tensor> nulltensor = c10::nullopt; |
| 228 | c10::OptionalIntArrayRef nulllen = c10::nullopt; | 197 | c10::OptionalIntArrayRef nulllen = c10::nullopt; |
| 229 | int64_t inner_precise = calculate_inner_precise_for_fa(query, attn_mask, is_causal); | 198 | int64_t inner_precise = calculate_inner_precise_for_fa(query, attn_mask, is_causal); |
| 230 | - auto output = | 199 | + auto output = at_npu::native::custom_ops::npu_fusion_attention(query, key, value, head_num, input_layout, |
| 231 | - at_npu::native::custom_ops::npu_fusion_attention(query, key, value, head_num, input_layout, nulltensor, | 200 | + nulltensor, nulltensor, atten_mask, input_scale.as_float_unchecked(), keep_prob, TOKEN_MAX, next_tockens, |
| 232 | - nulltensor, atten_mask, input_scale.as_float_unchecked(), | 201 | + inner_precise, nulllen, nulllen, nulllen, sparse_mode, true, false); |
| 233 | - keep_prob, TOKEN_MAX, next_tockens, inner_precise, nulllen, | ||
| 234 | - nulllen, nulllen, sparse_mode, true, false); | ||
| 235 | return std::get<0>(output); | 202 | return std::get<0>(output); |
| 236 | - } else if (!force_math && | 203 | + } else if (!force_math && (!query.requires_grad() && !key.requires_grad() && !value.requires_grad()) && |
| 237 | - (!query.requires_grad() && !key.requires_grad() && !value.requires_grad()) && | 204 | + (query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM) && |
| 238 | - (query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM) && | 205 | + (query.size(0) != 0 && query.size(1) != 0 && query.size(2) != 0 && query.size(3) != 0) && |
| 239 | - (query.size(0) != 0 && query.size(1) != 0 && query.size(2) != 0 && query.size(3) != 0) && | 206 | + (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16) && query.size(3) % 16 == 0 && |
| 240 | - (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16) && query.size(3) % 16 == 0 && | 207 | + ((attn_mask.has_value() && attn_mask->dtype() == at::kBool && attn_mask->size(-2) == query.size(2) && |
| 241 | - ((attn_mask.has_value() && attn_mask->dtype() == at::kBool && attn_mask->size(-2) == query.size(2) && | 208 | + attn_mask->size(-1) == key.size(2)) || |
| 242 | - attn_mask->size(-1) == key.size(2)) || !attn_mask.has_value()) && | 209 | + !attn_mask.has_value()) && |
| 243 | - (query.size(0) == key.size(0) && query.size(3) == key.size(3) && key.size(0) == value.size(0) && | 210 | + (query.size(0) == key.size(0) && query.size(3) == key.size(3) && key.size(0) == value.size(0) && |
| 244 | - key.size(1) == value.size(1) && key.size(2) == value.size(2) && key.size(3) == value.size(3)) && | 211 | + key.size(1) == value.size(1) && key.size(2) == value.size(2) && key.size(3) == value.size(3)) && |
| 245 | - (query.size(0) <= B_LIMIT && query.size(1) <= FIA_N_LIMIT && query.size(3) <= D_LIMIT && key.size(1) <= FIA_N_LIMIT) && | 212 | + (query.size(0) <= B_LIMIT && query.size(1) <= FIA_N_LIMIT && query.size(3) <= D_LIMIT && |
| 246 | - (key.size(1) != 0 && query.size(1) % key.size(1) == 0 && query.size(1) / key.size(1) <= 64) && | 213 | + key.size(1) <= FIA_N_LIMIT) && |
| 247 | - c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { | 214 | + (key.size(1) != 0 && query.size(1) % key.size(1) == 0 && query.size(1) / key.size(1) <= 64) && |
| 248 | - /* The implementation of the NPU Fused Infer Attention operator constraints: | 215 | + c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { |
| 249 | - 1. It only supports the type of fp16 or bf16, and dim = BNSD_DIM | 216 | + /* The implementation of the NPU Fused Infer Attention operator constraints: |
| 250 | - 2. The shape [B, N1, S1, D] of the query is suppported, where B <= B_LIMIT, N1 <= FIA_N_LIMIT, | 217 | + 1. It only supports the type of fp16 or bf16, and dim = BNSD_DIM |
| 251 | - D <= D_LIMIT and must be a positive integer multiple of 16. | 218 | + 2. The shape [B, N1, S1, D] of the query is suppported, where B <= B_LIMIT, N1 <= FIA_N_LIMIT, |
| 252 | - 3. The shape of the key and value should be the same, and dim = BNSD_DIM. For GQA, the key shape is [B, N2, S2, D], | 219 | + D <= D_LIMIT and must be a positive integer multiple of 16. |
| 253 | - where 0 < N2 <= FIA_N_LIMIT, N1 is a positive integer multiple of N2 and N1/N2 <= 64. | 220 | + 3. The shape of the key and value should be the same, and dim = BNSD_DIM. For GQA, the key shape is [B, N2, S2, D], |
| 254 | - 4. The attn_mask supports only the bool data type, and shpae[-2] = S1, shpae[-1] = S2. | 221 | + where 0 < N2 <= FIA_N_LIMIT, N1 is a positive integer multiple of N2 and N1/N2 <= 64. |
| 255 | - 5. It only supports SocVersion after Ascend910B1. */ | 222 | + 4. The attn_mask supports only the bool data type, and shpae[-2] = S1, shpae[-1] = S2. |
| 256 | - int64_t inner_precise = 0; | 223 | + 5. It only supports SocVersion after Ascend910B1. */ |
| 257 | - if (query.size(2) == 1) { | 224 | + int64_t inner_precise = 0; |
| 258 | - is_causal = false; | 225 | + if (query.size(2) == 1) { |
| 259 | - } else { | 226 | + is_causal = false; |
| 260 | - if (!is_causal) { | 227 | + } else { |
| 261 | - inner_precise = 2; | 228 | + if (!is_causal) { |
| 262 | - } | 229 | + inner_precise = 2; |
| 263 | - } | 230 | + } |
| 264 | - auto processed_mask_fia = preprocess_attn_mask_for_npu( | 231 | + } |
| 265 | - attn_mask, query, key); | 232 | + auto processed_mask_fia = preprocess_attn_mask_for_npu(attn_mask, query, key); |
| 266 | - c10::optional<at::Tensor> atten_mask = convert_boolean_attn_mask(query, processed_mask_fia, is_causal); | 233 | + c10::optional<at::Tensor> atten_mask = convert_boolean_attn_mask(query, processed_mask_fia, is_causal); |
| 267 | - int64_t head_num = query.size(1); | 234 | + int64_t head_num = query.size(1); |
| 268 | - int64_t head_num_kv = key.size(1); | 235 | + int64_t head_num_kv = key.size(1); |
| 269 | - c10::string_view input_layout = "BNSD"; | 236 | + c10::string_view input_layout = "BNSD"; |
| 270 | - auto input_scale = calculate_scale(query, scale); | 237 | + auto input_scale = calculate_scale(query, scale); |
| 271 | - int64_t next_tockens = is_causal ? 0 : TOKEN_MAX; | 238 | + int64_t next_tockens = is_causal ? 0 : TOKEN_MAX; |
| 272 | - int64_t sparse_mode = is_causal ? LEFT_UP_CAUSAL : 0; | 239 | + int64_t sparse_mode = is_causal ? LEFT_UP_CAUSAL : 0; |
| 273 | - c10::optional<at::Tensor> nulltensor = c10::nullopt; | 240 | + c10::optional<at::Tensor> nulltensor = c10::nullopt; |
| 274 | - at::OptionalSymIntArrayRef nulllen = c10::nullopt; | 241 | + at::OptionalSymIntArrayRef nulllen = c10::nullopt; |
| 275 | - auto output = | 242 | + auto output = at_npu::native::custom_ops::npu_fused_infer_attention_score(query, key, value, nulltensor, |
| 276 | - at_npu::native::custom_ops::npu_fused_infer_attention_score(query, key, value, nulltensor, atten_mask, nulllen, nulllen, nulltensor, | 243 | + atten_mask, nulllen, nulllen, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, |
| 277 | - nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, | 244 | + nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, |
| 278 | - nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, | 245 | + nulltensor, nulllen, nulltensor, nulltensor, nulltensor, head_num, input_scale.as_float_unchecked(), |
| 279 | - nulltensor, nulltensor, nulltensor, nulllen, nulltensor, nulltensor, nulltensor, head_num, input_scale.as_float_unchecked(), | 246 | + TOKEN_MAX, next_tockens, input_layout, head_num_kv, sparse_mode, inner_precise, 0, 0, 0, 0, false); |
| 280 | - TOKEN_MAX, next_tockens, input_layout, head_num_kv, sparse_mode, inner_precise, 0, 0, 0, 0, false); | ||
| 281 | 247 | ||
| 282 | - return std::get<0>(output); | 248 | + return std::get<0>(output); |
| 283 | } | 249 | } |
| 284 | c10::optional<at::Tensor> atten_mask_math = convert_boolean_attn_mask_math(attn_mask, query.dtype()); | 250 | c10::optional<at::Tensor> atten_mask_math = convert_boolean_attn_mask_math(attn_mask, query.dtype()); |
| 285 | - auto output = at::_scaled_dot_product_attention_math(query, key, value, atten_mask_math, dropout_p, is_causal, | 251 | + auto output = at::_scaled_dot_product_attention_math( |
| 286 | - c10::nullopt, scale); | 252 | + query, key, value, atten_mask_math, dropout_p, is_causal, c10::nullopt, scale); |
| 287 | return std::get<0>(output); | 253 | return std::get<0>(output); |
| 288 | } | 254 | } |
| 289 | 255 | ||
| 290 | 256 | ||
| 291 | 257 | ||
| 292 | -at::Tensor scaled_dot_product_attention( | 258 | +at::Tensor scaled_dot_product_attention(const at::Tensor &query, const at::Tensor &key, const at::Tensor &value, |
| 293 | - const at::Tensor &query, | 259 | + const c10::optional<at::Tensor> &attn_mask, double dropout_p, bool is_causal, c10::optional<double> scale, |
| 294 | - const at::Tensor &key, | 260 | + bool enable_gqa) { |
| 295 | - const at::Tensor &value, | ||
| 296 | - const c10::optional<at::Tensor> &attn_mask, | ||
| 297 | - double dropout_p, | ||
| 298 | - bool is_causal, | ||
| 299 | - c10::optional<double> scale, | ||
| 300 | - bool enable_gqa) | ||
| 301 | -{ | ||
| 302 | validate_sdpa_input(query, key, value, attn_mask); | 261 | validate_sdpa_input(query, key, value, attn_mask); |
| 303 | 262 | ||
| 304 | static auto compatible_impl = at_npu::native::env::CheckCompatibleImpl(); | 263 | static auto compatible_impl = at_npu::native::env::CheckCompatibleImpl(); |
| 305 | bool force_math = false; | 264 | bool force_math = false; |
| 306 | if (compatible_impl) { | 265 | if (compatible_impl) { |
| 307 | - auto& ctx = at::globalContext(); | 266 | + auto &ctx = at::globalContext(); |
| 308 | force_math = ctx.userEnabledMathSDP() && !ctx.userEnabledFlashSDP(); | 267 | force_math = ctx.userEnabledMathSDP() && !ctx.userEnabledFlashSDP(); |
| 309 | if (query.scalar_type() == at::kFloat && | 268 | if (query.scalar_type() == at::kFloat && |
| 310 | (query.size(-1) % 4 != 0 || key.size(-1) % 4 != 0 || value.size(-1) % 4 != 0)) { | 269 | (query.size(-1) % 4 != 0 || key.size(-1) % 4 != 0 || value.size(-1) % 4 != 0)) { |
| 311 | force_math = true; | 270 | force_math = true; |
| 312 | } | 271 | } |
| 272 | + if (query.size(-1) > 256 || key.size(-1) > 256 || value.size(-1) > 256) { | ||
| 273 | + force_math = true; | ||
| 274 | + } | ||
| 313 | } | 275 | } |
| 314 | 276 | ||
| 315 | if (!force_math && | 277 | if (!force_math && |
| 316 | - (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16 || query.scalar_type() == at::kFloat) && | 278 | + (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16 || |
| 279 | + query.scalar_type() == at::kFloat) && | ||
| 317 | ((attn_mask.has_value() && attn_mask->dtype() == at::kBool) || !attn_mask.has_value()) && | 280 | ((attn_mask.has_value() && attn_mask->dtype() == at::kBool) || !attn_mask.has_value()) && |
| 318 | - query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM && | 281 | + query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM && query.size(1) <= N_LIMIT && |
| 319 | - query.size(1) <= N_LIMIT && query.size(3) <= D_LIMIT && key.size(1) <= N_LIMIT && | 282 | + query.size(3) <= D_LIMIT && key.size(1) <= N_LIMIT && query.size(1) % key.size(1) == 0 && |
| 320 | - query.size(1) % key.size(1) == 0 && query.size(1) / key.size(1) > 0 && | 283 | + query.size(1) / key.size(1) > 0 && c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { |
| 321 | - c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { | 284 | + auto processed_mask = preprocess_attn_mask_for_npu(attn_mask, query, key); |
| 322 | - auto processed_mask = preprocess_attn_mask_for_npu( | ||
| 323 | - attn_mask, query, key); | ||
| 324 | 285 | ||
| 325 | /* The implementation of the NPU FlashAttention fusion operator without grad constraints: | 286 | /* The implementation of the NPU FlashAttention fusion operator without grad constraints: |
| 326 | 1. The FA operator must be registered. GetSocVersion api is provisional, IsExistOp aclnn api will apply. | 287 | 1. The FA operator must be registered. GetSocVersion api is provisional, IsExistOp aclnn api will apply. |
| @@ -340,33 +301,30 @@ at::Tensor scaled_dot_product_attention( | |||
| 340 | c10::OptionalArrayRef<c10::SymInt> nullsymlen = c10::nullopt; | 301 | c10::OptionalArrayRef<c10::SymInt> nullsymlen = c10::nullopt; |
| 341 | int64_t inner_precise = calculate_inner_precise_for_fa(query, attn_mask, is_causal); | 302 | int64_t inner_precise = calculate_inner_precise_for_fa(query, attn_mask, is_causal); |
| 342 | if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) { | 303 | if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) { |
| 343 | - auto outputv3 = | 304 | + auto outputv3 = at_npu::native::custom_ops::npu_fusion_attention_v3(query, key, value, head_num, |
| 344 | - at_npu::native::custom_ops::npu_fusion_attention_v3(query, key, value, head_num, input_layout, nulltensor, | 305 | + input_layout, nulltensor, nulltensor, atten_mask, input_scale.as_float_unchecked(), keep_prob, |
| 345 | - nulltensor, atten_mask, input_scale.as_float_unchecked(), | 306 | + TOKEN_MAX, next_tockens, inner_precise, nullsymlen, nulltensor, nulltensor, sparse_mode, true, false); |
| 346 | - keep_prob, TOKEN_MAX, next_tockens, inner_precise, nullsymlen, | ||
| 347 | - nulltensor, nulltensor, sparse_mode, true, false); | ||
| 348 | return std::get<0>(outputv3); | 307 | return std::get<0>(outputv3); |
| 349 | } | 308 | } |
| 350 | c10::OptionalIntArrayRef nulllen = c10::nullopt; | 309 | c10::OptionalIntArrayRef nulllen = c10::nullopt; |
| 351 | - auto output = | 310 | + auto output = at_npu::native::custom_ops::npu_fusion_attention(query, key, value, head_num, input_layout, |
| 352 | - at_npu::native::custom_ops::npu_fusion_attention(query, key, value, head_num, input_layout, nulltensor, | 311 | + nulltensor, nulltensor, atten_mask, input_scale.as_float_unchecked(), keep_prob, TOKEN_MAX, next_tockens, |
| 353 | - nulltensor, atten_mask, input_scale.as_float_unchecked(), | 312 | + inner_precise, nulllen, nulllen, nulllen, sparse_mode, true, false); |
| 354 | - keep_prob, TOKEN_MAX, next_tockens, inner_precise, nulllen, | ||
| 355 | - nulllen, nulllen, sparse_mode, true, false); | ||
| 356 | return std::get<0>(output); | 313 | return std::get<0>(output); |
| 357 | - } else if (!force_math && | 314 | + } else if (!force_math && (!query.requires_grad() && !key.requires_grad() && !value.requires_grad()) && |
| 358 | - (!query.requires_grad() && !key.requires_grad() && !value.requires_grad()) && | 315 | + (query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM) && |
| 359 | - (query.dim() == BNSD_DIM && key.dim() == BNSD_DIM && value.dim() == BNSD_DIM) && | 316 | + (query.size(0) != 0 && query.size(1) != 0 && query.size(2) != 0 && query.size(3) != 0) && |
| 360 | - (query.size(0) != 0 && query.size(1) != 0 && query.size(2) != 0 && query.size(3) != 0) && | 317 | + (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16) && query.size(3) % 16 == 0 && |
| 361 | - (query.scalar_type() == at::kHalf || query.scalar_type() == at::kBFloat16) && query.size(3) % 16 ==0 && | 318 | + ((attn_mask.has_value() && attn_mask->dtype() == at::kBool && attn_mask->size(-2) == query.size(2) && |
| 362 | - ((attn_mask.has_value() && attn_mask->dtype() == at::kBool && attn_mask->size(-2) == query.size(2) && | 319 | + attn_mask->size(-1) == key.size(2)) || |
| 363 | - attn_mask->size(-1) == key.size(2)) || !attn_mask.has_value()) && | 320 | + !attn_mask.has_value()) && |
| 364 | - (query.size(0) == key.size(0) && query.size(3) == key.size(3) && key.size(0) == value.size(0) && | 321 | + (query.size(0) == key.size(0) && query.size(3) == key.size(3) && key.size(0) == value.size(0) && |
| 365 | - key.size(1) == value.size(1) && key.size(2) == value.size(2) && key.size(3) == value.size(3)) && | 322 | + key.size(1) == value.size(1) && key.size(2) == value.size(2) && key.size(3) == value.size(3)) && |
| 366 | - (query.size(0) <= B_LIMIT && query.size(1) <= FIA_N_LIMIT && query.size(3) <= D_LIMIT && key.size(1) <= FIA_N_LIMIT) && | 323 | + (query.size(0) <= B_LIMIT && query.size(1) <= FIA_N_LIMIT && query.size(3) <= D_LIMIT && |
| 367 | - (key.size(1) != 0 && query.size(1) % key.size(1) == 0 && query.size(1) / key.size(1) <= 64) && | 324 | + key.size(1) <= FIA_N_LIMIT) && |
| 368 | - c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { | 325 | + (key.size(1) != 0 && query.size(1) % key.size(1) == 0 && query.size(1) / key.size(1) <= 64) && |
| 369 | - /* The implementation of the NPU Fused Infer Attention operator constraints: | 326 | + c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1) { |
| 327 | + /* The implementation of the NPU Fused Infer Attention operator constraints: | ||
| 370 | 1. It only supports the type of fp16 or bf16, and dim = BNSD_DIM | 328 | 1. It only supports the type of fp16 or bf16, and dim = BNSD_DIM |
| 371 | 2. The shape [B, N1, S1, D] of the query is suppported, where B <= B_LIMIT, N1 <= FIA_N_LIMIT, | 329 | 2. The shape [B, N1, S1, D] of the query is suppported, where B <= B_LIMIT, N1 <= FIA_N_LIMIT, |
| 372 | D <= D_LIMIT and must be a positive integer multiple of 16. | 330 | D <= D_LIMIT and must be a positive integer multiple of 16. |
| @@ -374,39 +332,37 @@ at::Tensor scaled_dot_product_attention( | |||
| 374 | where 0 < N2 <= FIA_N_LIMIT, N1 is a positive integer multiple of N2 and N1/N2 <= 64. | 332 | where 0 < N2 <= FIA_N_LIMIT, N1 is a positive integer multiple of N2 and N1/N2 <= 64. |
| 375 | 4. The attn_mask supports only the bool data type, and shpae[-2] = S1, shpae[-1] = S2. | 333 | 4. The attn_mask supports only the bool data type, and shpae[-2] = S1, shpae[-1] = S2. |
| 376 | 5. It only supports SocVersion after Ascend910B1. */ | 334 | 5. It only supports SocVersion after Ascend910B1. */ |
| 377 | - int64_t inner_precise = 0; | 335 | + int64_t inner_precise = 0; |
| 378 | - if (query.size(2) == 1) { | 336 | + if (query.size(2) == 1) { |
| 379 | - is_causal = false; | 337 | + is_causal = false; |
| 380 | - } else { | 338 | + } else { |
| 381 | - if (!is_causal) { | 339 | + if (!is_causal) { |
| 382 | - inner_precise = 2; | 340 | + inner_precise = 2; |
| 383 | - } | 341 | + } |
| 384 | - } | 342 | + } |
| 385 | - auto processed_mask_fia = preprocess_attn_mask_for_npu( | 343 | + auto processed_mask_fia = preprocess_attn_mask_for_npu(attn_mask, query, key); |
| 386 | - attn_mask, query, key); | 344 | + c10::optional<at::Tensor> atten_mask = convert_boolean_attn_mask(query, processed_mask_fia, is_causal); |
| 387 | - c10::optional<at::Tensor> atten_mask = convert_boolean_attn_mask(query, processed_mask_fia, is_causal); | 345 | + int64_t head_num = query.size(1); |
| 388 | - int64_t head_num = query.size(1); | 346 | + int64_t head_num_kv = key.size(1); |
| 389 | - int64_t head_num_kv = key.size(1); | 347 | + c10::string_view input_layout = "BNSD"; |
| 390 | - c10::string_view input_layout = "BNSD"; | 348 | + auto input_scale = calculate_scale(query, scale); |
| 391 | - auto input_scale = calculate_scale(query, scale); | 349 | + int64_t next_tockens = is_causal ? 0 : TOKEN_MAX; |
| 392 | - int64_t next_tockens = is_causal ? 0 : TOKEN_MAX; | 350 | + int64_t sparse_mode = is_causal ? LEFT_UP_CAUSAL : 0; |
| 393 | - int64_t sparse_mode = is_causal ? LEFT_UP_CAUSAL : 0; | 351 | + c10::optional<at::Tensor> nulltensor = c10::nullopt; |
| 394 | - c10::optional<at::Tensor> nulltensor = c10::nullopt; | 352 | + at::OptionalSymIntArrayRef nulllen = c10::nullopt; |
| 395 | - at::OptionalSymIntArrayRef nulllen = c10::nullopt; | 353 | + auto output = at_npu::native::custom_ops::npu_fused_infer_attention_score(query, key, value, nulltensor, |
| 396 | - auto output = | 354 | + atten_mask, nulllen, nulllen, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, |
| 397 | - at_npu::native::custom_ops::npu_fused_infer_attention_score(query, key, value, nulltensor, atten_mask, nulllen, nulllen, nulltensor, | 355 | + nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, |
| 398 | - nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, | 356 | + nulltensor, nulllen, nulltensor, nulltensor, nulltensor, head_num, input_scale.as_float_unchecked(), |
| 399 | - nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, nulltensor, | 357 | + TOKEN_MAX, next_tockens, input_layout, head_num_kv, sparse_mode, inner_precise, 0, 0, 0, 0, false); |
| 400 | - nulltensor, nulltensor, nulltensor, nulllen, nulltensor, nulltensor, nulltensor, head_num, input_scale.as_float_unchecked(), | ||
| 401 | - TOKEN_MAX, next_tockens, input_layout, head_num_kv, sparse_mode, inner_precise, 0, 0, 0, 0, false); | ||
| 402 | 358 | ||
| 403 | - return std::get<0>(output); | 359 | + return std::get<0>(output); |
| 404 | } | 360 | } |
| 405 | c10::optional<at::Tensor> atten_mask_math = convert_boolean_attn_mask_math(attn_mask, query.dtype()); | 361 | c10::optional<at::Tensor> atten_mask_math = convert_boolean_attn_mask_math(attn_mask, query.dtype()); |
| 406 | - auto output = at::_scaled_dot_product_attention_math(query, key, value, atten_mask_math, dropout_p, is_causal, | 362 | + auto output = at::_scaled_dot_product_attention_math( |
| 407 | - c10::nullopt, scale, enable_gqa); | 363 | + query, key, value, atten_mask_math, dropout_p, is_causal, c10::nullopt, scale, enable_gqa); |
| 408 | return std::get<0>(output); | 364 | return std::get<0>(output); |
| 409 | } | 365 | } |
| 410 | 366 | ||
| 411 | 367 | ||
| 412 | -} // namespace op_api | 368 | +} // namespace op_api |
W | |||
| @@ -175,5 +175,16 @@ class TestScaledDotProductAttention(TestCase): | |||
| 175 | npu_output = torch.nn.functional.scaled_dot_product_attention(query.npu(), key.npu(), value.npu()) | 175 | npu_output = torch.nn.functional.scaled_dot_product_attention(query.npu(), key.npu(), value.npu()) |
| 176 | self.assertRtolEqual(cpu_output, npu_output, 0.001) | 176 | self.assertRtolEqual(cpu_output, npu_output, 0.001) |
| 177 | 177 | ||
| 178 | + | ||
| 179 | + def test_sdpa_last_dim_greater_than_256(self): | ||
| 180 | + torch_npu.npu.use_compatible_impl(True) | ||
| 181 | + query = torch.rand(1, 4, 10, 257, dtype=torch.float16) | ||
| 182 | + key = torch.rand(1, 4, 10, 257, dtype=torch.float16) | ||
| 183 | + value = torch.rand(1, 4, 10, 257, dtype=torch.float16) | ||
| 184 | + cpu_output = torch.nn.functional.scaled_dot_product_attention(query.to(torch.float32), key.to(torch.float32), value.to(torch.float32)) | ||
| 185 | + npu_output = torch.nn.functional.scaled_dot_product_attention(query.npu(), key.npu(), value.npu()) | ||
| 186 | + self.assertRtolEqual(cpu_output.to(torch.float16), npu_output, 0.001) | ||
W A2/A5都支持超256维度吗   XianglongZeng 24 天前 评论： 24 天前 评论： | |||
| 187 | + | ||
| 188 | + | ||
| 178 | if __name__ == "__main__": | 189 | if __name__ == "__main__": |
| 179 | run_tests() | 190 | run_tests() |
看上去有大量code style的差异，看不出来关键修改是啥了。看下代码风格的修改是否能回退，保证最小修改