Star181
852
代码
代码
Issues161
Pull Requests323
流水线
项目讨论
Wiki
分析
Star181
852
已合并
prolog support hifloat8 #4495
Le_666创建于 3月18日
prolog support hifloat8 #4495
已合并
Le_666创建于 3月18日
从Le_666/op-pluginmaster合入到Ascend/op-pluginmaster
共 3 个文件变更+330-138
Mop_plugin/config/op_plugin_functions.yaml+2-2
@@ -6255,11 +6255,11 @@ custom:
6255 exposed: [v2.1, newest]6255 exposed: [v2.1, newest]
6256 internal_format_opapi: [v2.1, newest]6256 internal_format_opapi: [v2.1, newest]
6257 6257 
6258- - func: npu_mla_prolog_v3(Tensor token_x, Tensor weight_dq, Tensor weight_uq_qr, Tensor weight_uk, Tensor weight_dkv_kr, Tensor rmsnorm_gamma_cq, Tensor rmsnorm_gamma_ckv, Tensor rope_sin, Tensor rope_cos, Tensor(a!) kv_cache, Tensor(b!) kr_cache, *, Tensor? cache_index=None, Tensor? dequant_scale_x=None, Tensor? dequant_scale_w_dq=None, Tensor? dequant_scale_w_uq_qr=None, Tensor? dequant_scale_w_dkv_kr=None, Tensor? quant_scale_ckv=None, Tensor? quant_scale_ckr=None, Tensor? smooth_scales_cq=None, Tensor? actual_seq_len=None, Tensor? k_nope_clip_alpha=None, float rmsnorm_epsilon_cq=1e-05, float rmsnorm_epsilon_ckv=1e-05, str cache_mode="PA_BSND", bool query_norm_flag=False, int weight_quant_mode=0, int kv_cache_quant_mode=0, int query_quant_mode=0, int ckvkr_repo_mode=0, int quant_scale_repo_mode=0, int tile_size=128, float qc_qr_scale=1.0, float kc_scale=1.0) -> (Tensor, Tensor, Tensor, Tensor, Tensor)6258+ - func: npu_mla_prolog_v3(Tensor token_x, Tensor weight_dq, Tensor weight_uq_qr, Tensor weight_uk, Tensor weight_dkv_kr, Tensor rmsnorm_gamma_cq, Tensor rmsnorm_gamma_ckv, Tensor rope_sin, Tensor rope_cos, Tensor(a!) kv_cache, Tensor(b!) kr_cache, *, Tensor? cache_index=None, Tensor? dequant_scale_x=None, Tensor? dequant_scale_w_dq=None, Tensor? dequant_scale_w_uq_qr=None, Tensor? dequant_scale_w_dkv_kr=None, Tensor? quant_scale_ckv=None, Tensor? quant_scale_ckr=None, Tensor? smooth_scales_cq=None, Tensor? actual_seq_len=None, Tensor? k_nope_clip_alpha=None, float rmsnorm_epsilon_cq=1e-05, float rmsnorm_epsilon_ckv=1e-05, str cache_mode="PA_BSND", bool query_norm_flag=False, int weight_quant_mode=0, int kv_cache_quant_mode=0, int query_quant_mode=0, int ckvkr_repo_mode=0, int quant_scale_repo_mode=0, int tile_size=128, float qc_qr_scale=1.0, float kc_scale=1.0, int? token_x_dtype=None, int? weight_dq_dtype=None, int? weight_uq_qr_dtype=None, int? weight_dkv_kr_dtype=None, int? kv_cache_dtype=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor)
6259 op_api: [v2.1, newest]6259 op_api: [v2.1, newest]
6260 internal_format_opapi: [v2.1, newest]6260 internal_format_opapi: [v2.1, newest]
6261 6261 
6262- - func: npu_mla_prolog_v3_functional(Tensor token_x, Tensor weight_dq, Tensor weight_uq_qr, Tensor weight_uk, Tensor weight_dkv_kr, Tensor rmsnorm_gamma_cq, Tensor rmsnorm_gamma_ckv, Tensor rope_sin, Tensor rope_cos, Tensor kv_cache, Tensor kr_cache, *, Tensor? cache_index=None, Tensor? dequant_scale_x=None, Tensor? dequant_scale_w_dq=None, Tensor? dequant_scale_w_uq_qr=None, Tensor? dequant_scale_w_dkv_kr=None, Tensor? quant_scale_ckv=None, Tensor? quant_scale_ckr=None, Tensor? smooth_scales_cq=None, Tensor? actual_seq_len=None, Tensor? k_nope_clip_alpha=None, float rmsnorm_epsilon_cq=1e-05, float rmsnorm_epsilon_ckv=1e-05, str cache_mode="PA_BSND", bool query_norm_flag=False, int weight_quant_mode=0, int kv_cache_quant_mode=0, int query_quant_mode=0, int ckvkr_repo_mode=0, int quant_scale_repo_mode=0, int tile_size=128, float qc_qr_scale=1.0, float kc_scale=1.0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)6262+ - func: npu_mla_prolog_v3_functional(Tensor token_x, Tensor weight_dq, Tensor weight_uq_qr, Tensor weight_uk, Tensor weight_dkv_kr, Tensor rmsnorm_gamma_cq, Tensor rmsnorm_gamma_ckv, Tensor rope_sin, Tensor rope_cos, Tensor kv_cache, Tensor kr_cache, *, Tensor? cache_index=None, Tensor? dequant_scale_x=None, Tensor? dequant_scale_w_dq=None, Tensor? dequant_scale_w_uq_qr=None, Tensor? dequant_scale_w_dkv_kr=None, Tensor? quant_scale_ckv=None, Tensor? quant_scale_ckr=None, Tensor? smooth_scales_cq=None, Tensor? actual_seq_len=None, Tensor? k_nope_clip_alpha=None, float rmsnorm_epsilon_cq=1e-05, float rmsnorm_epsilon_ckv=1e-05, str cache_mode="PA_BSND", bool query_norm_flag=False, int weight_quant_mode=0, int kv_cache_quant_mode=0, int query_quant_mode=0, int ckvkr_repo_mode=0, int quant_scale_repo_mode=0, int tile_size=128, float qc_qr_scale=1.0, float kc_scale=1.0, int? token_x_dtype=None, int? weight_dq_dtype=None, int? weight_uq_qr_dtype=None, int? weight_dkv_kr_dtype=None, int? kv_cache_dtype=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
6263 op_api: [v2.1, newest]6263 op_api: [v2.1, newest]
6264 internal_format_opapi: [v2.1, newest]6264 internal_format_opapi: [v2.1, newest]
6265 6265 
Mop_plugin/ops/opapi/MlaPrologV3KernelNpuOpApi.cpp+223-110
@@ -25,9 +25,12 @@ const int DIM_3 = 3;
25const int MODE_1 = 1;25const int MODE_1 = 1;
26const int MODE_2 = 2;26const int MODE_2 = 2;
27const int MODE_3 = 3;27const int MODE_3 = 3;
28+const int MODE_4 = 4;
29+const int MODE_5 = 5;
28const int FP8_E4M3_BLOCK_SIZE = 32;30const int FP8_E4M3_BLOCK_SIZE = 32;
29const char* const REQUIRED_CANN_VERSION = "8.5.0.alpha003";31const char* const REQUIRED_CANN_VERSION = "8.5.0.alpha003";
30const char* const CANN_PRODUCT = "CANN";32const char* const CANN_PRODUCT = "CANN";
33+bool is_hifloat8_dtype = false;
31bool is_cann_version_gte_required = IsGteCANNVersion(REQUIRED_CANN_VERSION, CANN_PRODUCT); // whether cann version >= 8.5.0.alpha00334bool is_cann_version_gte_required = IsGteCANNVersion(REQUIRED_CANN_VERSION, CANN_PRODUCT); // whether cann version >= 8.5.0.alpha003
32 35 
33std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_prolog_v3(36std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_prolog_v3(
@@ -40,11 +43,47 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_p
40 const c10::optional<at::Tensor>& quant_scale_ckr, const c10::optional<at::Tensor>& smooth_scales_cq,43 const c10::optional<at::Tensor>& quant_scale_ckr, const c10::optional<at::Tensor>& smooth_scales_cq,
41 const c10::optional<at::Tensor>& actual_seq_len, const c10::optional<at::Tensor>& k_nope_clip_alpha,44 const c10::optional<at::Tensor>& actual_seq_len, const c10::optional<at::Tensor>& k_nope_clip_alpha,
42 double rmsnorm_epsilon_cq, double rmsnorm_epsilon_ckv, c10::string_view cache_mode, bool query_norm_flag, int64_t weight_quant_mode, int64_t kv_cache_quant_mode,45 double rmsnorm_epsilon_cq, double rmsnorm_epsilon_ckv, c10::string_view cache_mode, bool query_norm_flag, int64_t weight_quant_mode, int64_t kv_cache_quant_mode,
43- int64_t query_quant_mode, int64_t ckvkr_repo_mode, int64_t quant_scale_repo_mode, int64_t tile_size, double qc_qr_scale, double kc_scale)46+ int64_t query_quant_mode, int64_t ckvkr_repo_mode, int64_t quant_scale_repo_mode, int64_t tile_size, double qc_qr_scale, double kc_scale,
47+ c10::optional<int64_t> token_x_dtype, c10::optional<int64_t> weight_dq_dtype, c10::optional<int64_t> weight_uq_qr_dtype,
48+ c10::optional<int64_t> weight_dkv_kr_dtype, c10::optional<int64_t> kv_cache_dtype)
44{49{
45 // construct the output tensor50 // construct the output tensor
46 if (weight_quant_mode == MODE_3) {51 if (weight_quant_mode == MODE_3) {
47 TORCH_CHECK(c10_npu::IsAclnnOnly(), "When weight_quant_mode is 3, not support on this soc version.", OPS_ERROR(ErrCode::NOT_SUPPORT));52 TORCH_CHECK(c10_npu::IsAclnnOnly(), "When weight_quant_mode is 3, not support on this soc version.", OPS_ERROR(ErrCode::NOT_SUPPORT));
53+ 
54+ auto dequant_scale_x_dtype = dequant_scale_x.value().dtype();
55+ auto dequant_scale_w_dq_dtype = dequant_scale_w_dq.value().dtype();
56+ auto dequant_scale_w_uq_qr_dtype = dequant_scale_w_uq_qr.value().dtype();
57+ auto dequant_scale_w_dkv_kr_dtype = dequant_scale_w_dkv_kr.value().dtype();
58+ 
59+#if VERSION_BETWEEN(V2R7, VERSION_NEWEST)
60+ TORCH_CHECK(dequant_scale_x_dtype == at::kFloat8_e8m0fnu && dequant_scale_w_dq_dtype == at::kFloat8_e8m0fnu &&
61+ dequant_scale_w_uq_qr_dtype == at::kFloat8_e8m0fnu && dequant_scale_w_dkv_kr_dtype == at::kFloat8_e8m0fnu,
62+ "torch_npu supports the float8_e8m0 only in version later than v2.7., dequant_scale_x, dequant_scale_w_dq, dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr Dtype in weight_quant_mode=3 must be torch.float8_e8m0fnu",
63+ OPS_ERROR(ErrCode::PARAM));
64+#endif
65+#if VERSION_BETWEEN(V2R1, V2R6)
66+ TORCH_CHECK(false, "torch_npu supports the float8_e8m0 only in version later than v2.7.", OPS_ERROR(ErrCode::PARAM));
67+#endif
68+ }
69+ 
70+ if (weight_quant_mode == MODE_5 && (token_x.dtype() == at::kByte || weight_dq.dtype() == at::kByte ||
71+ weight_uq_qr.dtype() == at::kByte || weight_dkv_kr.dtype() == at::kByte)) {
72+ TORCH_CHECK(token_x_dtype.has_value() && weight_dq_dtype.has_value() && weight_uq_qr_dtype.has_value() && weight_dkv_kr_dtype.has_value(),
73+ "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype,weight_dq_dtype,weight_uq_qr_dtype,weight_dkv_kr_dtype cannot be null.", OPS_ERROR(ErrCode::PARAM));
74+ 
75+ TORCH_CHECK(c10_npu::GetAclDataType(token_x_dtype.value()) == aclDataType::ACL_HIFLOAT8 && c10_npu::GetAclDataType(weight_dq_dtype.value()) == aclDataType::ACL_HIFLOAT8 &&
76+ c10_npu::GetAclDataType(weight_uq_qr_dtype.value()) == aclDataType::ACL_HIFLOAT8 && c10_npu::GetAclDataType(weight_dkv_kr_dtype.value()) == aclDataType::ACL_HIFLOAT8,
77+ "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype, weight_dq_dtype, weight_uq_qr_dtype, weight_dkv_kr_dtype value must be torch_npu.hifloat8", OPS_ERROR(ErrCode::PARAM));
78+ 
79+ if (kv_cache_quant_mode == 1 || kv_cache_quant_mode == 3) {
80+ TORCH_CHECK(kv_cache_dtype.has_value(),
81+ "when weight_quant_mode is 5 and kv_cache_quant_mode is 1 or 3 and input dtype is hifloat8, kv_cache_dtype cannot be null.", OPS_ERROR(ErrCode::PARAM));
82+ 
83+ TORCH_CHECK(c10_npu::GetAclDataType(kv_cache_dtype.value()) == aclDataType::ACL_HIFLOAT8,
84+ "when weight_quant_mode is 5 and input dtype is hifloat8, kv_cache_dtype value must be torch_npu.hifloat8", OPS_ERROR(ErrCode::PARAM));
85+ }
86+ is_hifloat8_dtype = true;
48 }87 }
49 88 
50 auto token_x_dim = token_x.dim();89 auto token_x_dim = token_x.dim();
@@ -60,94 +99,113 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_p
60 at::Tensor dequant_scale_q_nope {nullptr};99 at::Tensor dequant_scale_q_nope {nullptr};
61 at::Tensor query_norm {nullptr};100 at::Tensor query_norm {nullptr};
62 at::Tensor dequant_scale_q_norm {nullptr};101 at::Tensor dequant_scale_q_norm {nullptr};
63- const at::Tensor& dequant_scale_x_tensor = c10::value_or_else(dequant_scale_x, [] {return at::Tensor();});
64 102 
65 if (token_x_dim == DIM_3) {103 if (token_x_dim == DIM_3) {
66 TORCH_CHECK(rope_sin_dim == DIM_3, "when token_x dim num is 3, rope_sin dim num should be 3, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));104 TORCH_CHECK(rope_sin_dim == DIM_3, "when token_x dim num is 3, rope_sin dim num should be 3, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));
67- if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3) && kv_cache_quant_mode == MODE_1) {105+ if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) && kv_cache_quant_mode == MODE_1) {
68- // weight_quant_mode=2且kv_cache_quant_mode=1时为全量化kv量化场景106+ // weight_quant_mode=2,4,5且kv_cache_quant_mode=1时为全量化kv量化场景(int8,fp8,hif8)
69 // weight_quant_mode=3且kv_cache_quant_mode=1时为mxfp8全量化kv量化场景107 // weight_quant_mode=3且kv_cache_quant_mode=1时为mxfp8全量化kv量化场景
70- // kvcache量化108+ if (is_hifloat8_dtype) {
71- query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));109+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kByte));
110+ } else {
111+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));
112+ }
72 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_uk.size(0), 1}, at::kFloat);113 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_uk.size(0), 1}, at::kFloat);
73 } else {114 } else {
74- query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(rope_sin.dtype()));115+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kBFloat16));
75 }116 }
76 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), rope_sin.size(2)}, at::kBFloat16);117 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), rope_sin.size(2)}, at::kBFloat16);
77 if (query_norm_flag) {118 if (query_norm_flag) {
78- query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));119+ if (is_hifloat8_dtype) {
79- if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2) { // weight_quant_mode=1 半量化场景,weight_quant_mode=2 全量化场景120+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_dq.size(1)}, token_x.options().dtype(at::kByte));
121+ } else {
122+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));
123+ }
124+ if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) {
125+ // weight_quant_mode=1 半量化场景, weight_quant_mode=2,4,5 全量化场景(int8,fp8,hif8)
80 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), 1}, at::kFloat);126 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), 1}, at::kFloat);
81 } else if (weight_quant_mode == MODE_3) {127 } else if (weight_quant_mode == MODE_3) {
82- dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x_tensor.options().dtype(dequant_scale_x_tensor.dtype()));128+ dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x.value().options().dtype(dequant_scale_x.value().dtype()));
83 }129 }
84 }130 }
85 } else {131 } else {
86 TORCH_CHECK(rope_sin_dim == DIM_2, "when token_x dim num is 2, rope_sin dim num should be 2, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));132 TORCH_CHECK(rope_sin_dim == DIM_2, "when token_x dim num is 2, rope_sin dim num should be 2, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));
87- if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3) && kv_cache_quant_mode == MODE_1) {133+ if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) && kv_cache_quant_mode == MODE_1) {
88- // weight_quant_mode=2且kv_cache_quant_mode=1时为全量化kv量化场景134+ if (is_hifloat8_dtype) {
89- // weight_quant_mode=3且kv_cache_quant_mode=1时为mxfp8全量化kv量化场景135+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kByte));
90- // kvcache量化136+ } else {
91- query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));137+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));
138+ }
92 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), 1}, at::kFloat);139 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), 1}, at::kFloat);
93 } else {140 } else {
94- query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(rope_sin.dtype()));141+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kBFloat16));
95 }142 }
96 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), rope_sin.size(1)}, at::kBFloat16);143 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), rope_sin.size(1)}, at::kBFloat16);
97 if (query_norm_flag) {144 if (query_norm_flag) {
98- query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));145+ if (is_hifloat8_dtype) {
99- if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2) { // weight_quant_mode=1 半量化场景,weight_quant_mode=2 全量化场景146+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1)}, token_x.options().dtype(at::kByte));
147+ } else {
148+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));
149+ }
150+ if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) {
100 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), 1}, at::kFloat);151 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), 1}, at::kFloat);
101 } else if (weight_quant_mode == MODE_3) {152 } else if (weight_quant_mode == MODE_3) {
102- dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x_tensor.options().dtype(dequant_scale_x_tensor.dtype()));153+ dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x.value().options().dtype(dequant_scale_x.value().dtype()));
103 }154 }
104 }155 }
105 }156 }
106 157 
107 char *cache_mode_ptr = const_cast<char *>(cache_mode.data());158 char *cache_mode_ptr = const_cast<char *>(cache_mode.data());
108 159 
109- if (weight_quant_mode == MODE_3) {160+ if (is_cann_version_gte_required) {
110- const at::Tensor& dequant_scale_w_dq_tensor = c10::value_or_else(dequant_scale_w_dq, [] {return at::Tensor();});161+ if (is_hifloat8_dtype) {
111- const at::Tensor& dequant_scale_w_uq_qr_tensor = c10::value_or_else(dequant_scale_w_uq_qr, [] {return at::Tensor();});162+ TensorWrapper token_x_wrapper = make_wrapper(token_x, token_x_dtype);
112- const at::Tensor& dequant_scale_w_dkv_kr_tensor = c10::value_or_else(dequant_scale_w_dkv_kr, [] {return at::Tensor();});163+ TensorWrapper weight_dq_wrapper = make_wrapper(weight_dq, weight_dq_dtype);
113- 164+ TensorWrapper weight_uq_qr_wrapper = make_wrapper(weight_uq_qr, weight_uq_qr_dtype);
114- TensorWrapper dequant_scale_x_wrapper = {dequant_scale_x_tensor, aclDataType::ACL_FLOAT8_E8M0};165+ TensorWrapper weight_dkv_kr_wrapper = make_wrapper(weight_dkv_kr, weight_dkv_kr_dtype);
115- TensorWrapper dequant_scale_w_dq_wrapper = {dequant_scale_w_dq_tensor, aclDataType::ACL_FLOAT8_E8M0};166+ if (kv_cache_quant_mode == MODE_1) {
116- TensorWrapper dequant_scale_w_uq_qr_wrapper = {dequant_scale_w_uq_qr_tensor, aclDataType::ACL_FLOAT8_E8M0};167+ TensorWrapper kv_cache_wrapper = make_wrapper(kv_cache, kv_cache_dtype);
117- TensorWrapper dequant_scale_w_dkv_kr_wrapper = {dequant_scale_w_dkv_kr_tensor, aclDataType::ACL_FLOAT8_E8M0};168+ TensorWrapper query_wrapper = make_wrapper(query, token_x_dtype);
118- TensorWrapper dequant_scale_q_norm_wrapper = {dequant_scale_q_norm, aclDataType::ACL_FLOAT8_E8M0};169+ TensorWrapper query_norm_wrapper = make_wrapper(query_norm, token_x_dtype);
119- 170+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x_wrapper, weight_dq_wrapper, weight_uq_qr_wrapper, weight_uk, weight_dkv_kr_wrapper, rmsnorm_gamma_cq,
120- if (is_cann_version_gte_required) {171+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_wrapper, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,
121- EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,172+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
122- rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x_wrapper, dequant_scale_w_dq_wrapper,173+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
123- dequant_scale_w_uq_qr_wrapper, dequant_scale_w_dkv_kr_wrapper, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,174+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query_wrapper, query_rope,
124- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,175+ dequant_scale_q_nope, query_norm_wrapper, dequant_scale_q_norm);
125- ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,176+ } else if (kv_cache_quant_mode == MODE_3) {
126- dequant_scale_q_nope, query_norm, dequant_scale_q_norm_wrapper);177+ TensorWrapper kv_cache_wrapper = make_wrapper(kv_cache, kv_cache_dtype);
127- } else {178+ TensorWrapper query_norm_wrapper = make_wrapper(query_norm, token_x_dtype);
128- EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,179+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x_wrapper, weight_dq_wrapper, weight_uq_qr_wrapper, weight_uk, weight_dkv_kr_wrapper, rmsnorm_gamma_cq,
129- rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x_wrapper, dequant_scale_w_dq_wrapper,180+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_wrapper, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,
130- dequant_scale_w_uq_qr_wrapper, dequant_scale_w_dkv_kr_wrapper, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,181+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
131- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, query_norm_flag, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,182+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
132- ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,183+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
133- dequant_scale_q_nope, query_norm, dequant_scale_q_norm_wrapper);184+ dequant_scale_q_nope, query_norm_wrapper, dequant_scale_q_norm);
134- }185+ } else {
135- } else {186+ TensorWrapper query_norm_wrapper = make_wrapper(query_norm, token_x_dtype);
136- if (is_cann_version_gte_required) {187+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x_wrapper, weight_dq_wrapper, weight_uq_qr_wrapper, weight_uk, weight_dkv_kr_wrapper, rmsnorm_gamma_cq,
137- EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,188+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,
138- rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,189+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
139- dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,190+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
140- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,191+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
141- ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,192+ dequant_scale_q_nope, query_norm_wrapper, dequant_scale_q_norm);
142- dequant_scale_q_nope, query_norm, dequant_scale_q_norm);193+ }
143 } else {194 } else {
144 EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,195 EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,
145 rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,196 rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,
146 dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,197 dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
147- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, query_norm_flag, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,198+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
148 ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,199 ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
149 dequant_scale_q_nope, query_norm, dequant_scale_q_norm);200 dequant_scale_q_nope, query_norm, dequant_scale_q_norm);
150 }201 }
202+ } else {
203+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,
204+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache, kr_cache, cache_index, dequant_scale_x, dequant_scale_w_dq,
205+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
206+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, query_norm_flag, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
207+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
208+ dequant_scale_q_nope, query_norm, dequant_scale_q_norm);
151 }209 }
152 210 
153 if (!query_norm.defined()) {211 if (!query_norm.defined()) {
@@ -158,7 +216,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_p
158 }216 }
159 if (!dequant_scale_q_norm.defined()) {217 if (!dequant_scale_q_norm.defined()) {
160 if (weight_quant_mode == MODE_3) {218 if (weight_quant_mode == MODE_3) {
161- dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, dequant_scale_x_tensor.options().dtype(dequant_scale_x_tensor.dtype()));219+ dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, dequant_scale_x.value().options().dtype(dequant_scale_x.value().dtype()));
162 } else {220 } else {
163 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, at::kFloat);221 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, at::kFloat);
164 }222 }
@@ -177,11 +235,47 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
177 const c10::optional<at::Tensor>& quant_scale_ckr, const c10::optional<at::Tensor>& smooth_scales_cq,235 const c10::optional<at::Tensor>& quant_scale_ckr, const c10::optional<at::Tensor>& smooth_scales_cq,
178 const c10::optional<at::Tensor>& actual_seq_len, const c10::optional<at::Tensor>& k_nope_clip_alpha,236 const c10::optional<at::Tensor>& actual_seq_len, const c10::optional<at::Tensor>& k_nope_clip_alpha,
179 double rmsnorm_epsilon_cq, double rmsnorm_epsilon_ckv, c10::string_view cache_mode, bool query_norm_flag, int64_t weight_quant_mode, int64_t kv_cache_quant_mode,237 double rmsnorm_epsilon_cq, double rmsnorm_epsilon_ckv, c10::string_view cache_mode, bool query_norm_flag, int64_t weight_quant_mode, int64_t kv_cache_quant_mode,
180- int64_t query_quant_mode, int64_t ckvkr_repo_mode, int64_t quant_scale_repo_mode, int64_t tile_size, double qc_qr_scale, double kc_scale)238+ int64_t query_quant_mode, int64_t ckvkr_repo_mode, int64_t quant_scale_repo_mode, int64_t tile_size, double qc_qr_scale, double kc_scale,
239+ c10::optional<int64_t> token_x_dtype, c10::optional<int64_t> weight_dq_dtype, c10::optional<int64_t> weight_uq_qr_dtype,
240+ c10::optional<int64_t> weight_dkv_kr_dtype, c10::optional<int64_t> kv_cache_dtype)
181{241{
182 // construct the output tensor242 // construct the output tensor
183 if (weight_quant_mode == MODE_3) {243 if (weight_quant_mode == MODE_3) {
184 TORCH_CHECK(c10_npu::IsAclnnOnly(), "When weight_quant_mode is 3, not support on this soc version.", OPS_ERROR(ErrCode::NOT_SUPPORT));244 TORCH_CHECK(c10_npu::IsAclnnOnly(), "When weight_quant_mode is 3, not support on this soc version.", OPS_ERROR(ErrCode::NOT_SUPPORT));
245+ 
246+ auto dequant_scale_x_dtype = dequant_scale_x.value().dtype();
247+ auto dequant_scale_w_dq_dtype = dequant_scale_w_dq.value().dtype();
248+ auto dequant_scale_w_uq_qr_dtype = dequant_scale_w_uq_qr.value().dtype();
249+ auto dequant_scale_w_dkv_kr_dtype = dequant_scale_w_dkv_kr.value().dtype();
250+ 
251+#if VERSION_BETWEEN(V2R7, VERSION_NEWEST)
252+ TORCH_CHECK(dequant_scale_x_dtype == at::kFloat8_e8m0fnu && dequant_scale_w_dq_dtype == at::kFloat8_e8m0fnu &&
253+ dequant_scale_w_uq_qr_dtype == at::kFloat8_e8m0fnu && dequant_scale_w_dkv_kr_dtype == at::kFloat8_e8m0fnu,
254+ "torch_npu supports the float8_e8m0 only in version later than v2.7., dequant_scale_x, dequant_scale_w_dq, dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr Dtype in weight_quant_mode=3 must be torch.float8_e8m0fnu",
255+ OPS_ERROR(ErrCode::PARAM));
256+#endif
257+#if VERSION_BETWEEN(V2R1, V2R6)
258+ TORCH_CHECK(false, "torch_npu supports the float8_e8m0 only in version later than v2.7.", OPS_ERROR(ErrCode::PARAM));
259+#endif
260+ }
261+ 
262+ if (weight_quant_mode == MODE_5 && (token_x.dtype() == at::kByte || weight_dq.dtype() == at::kByte ||
263+ weight_uq_qr.dtype() == at::kByte || weight_dkv_kr.dtype() == at::kByte)) {
264+ TORCH_CHECK(token_x_dtype.has_value() && weight_dq_dtype.has_value() && weight_uq_qr_dtype.has_value() && weight_dkv_kr_dtype.has_value(),
265+ "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype,weight_dq_dtype,weight_uq_qr_dtype,weight_dkv_kr_dtype cannot be null.", OPS_ERROR(ErrCode::PARAM));
266+ 
267+ TORCH_CHECK(c10_npu::GetAclDataType(token_x_dtype.value()) == aclDataType::ACL_HIFLOAT8 && c10_npu::GetAclDataType(weight_dq_dtype.value()) == aclDataType::ACL_HIFLOAT8 &&
268+ c10_npu::GetAclDataType(weight_uq_qr_dtype.value()) == aclDataType::ACL_HIFLOAT8 && c10_npu::GetAclDataType(weight_dkv_kr_dtype.value()) == aclDataType::ACL_HIFLOAT8,
269+ "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype, weight_dq_dtype, weight_uq_qr_dtype, weight_dkv_kr_dtype value must be torch_npu.hifloat8", OPS_ERROR(ErrCode::PARAM));
270+ 
271+ if (kv_cache_quant_mode == 1 or kv_cache_quant_mode == 3) {
272+ TORCH_CHECK(kv_cache_dtype.has_value(),
273+ "when weight_quant_mode is 5 and kv_cache_quant_mode is 1 or 3 and input dtype is hifloat8, kv_cache_dtype cannot be null.", OPS_ERROR(ErrCode::PARAM));
274+ 
275+ TORCH_CHECK(c10_npu::GetAclDataType(kv_cache_dtype.value()) == aclDataType::ACL_HIFLOAT8,
276+ "when weight_quant_mode is 5 and input dtype is hifloat8, kv_cache_dtype value must be torch_npu.hifloat8", OPS_ERROR(ErrCode::PARAM));
277+ }
278+ is_hifloat8_dtype = true;
185 }279 }
186 280 
187 auto token_x_dim = token_x.dim();281 auto token_x_dim = token_x.dim();
@@ -197,46 +291,58 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
197 at::Tensor dequant_scale_q_nope {nullptr};291 at::Tensor dequant_scale_q_nope {nullptr};
198 at::Tensor query_norm {nullptr};292 at::Tensor query_norm {nullptr};
199 at::Tensor dequant_scale_q_norm {nullptr};293 at::Tensor dequant_scale_q_norm {nullptr};
200- const at::Tensor& dequant_scale_x_tensor = c10::value_or_else(dequant_scale_x, [] {return at::Tensor();});
201 294 
202 if (token_x_dim == DIM_3) {295 if (token_x_dim == DIM_3) {
203 TORCH_CHECK(rope_sin_dim == DIM_3, "when token_x dim num is 3, rope_sin dim num should be 3, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));296 TORCH_CHECK(rope_sin_dim == DIM_3, "when token_x dim num is 3, rope_sin dim num should be 3, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));
204- if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3) && kv_cache_quant_mode == MODE_1) {297+ if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) && kv_cache_quant_mode == MODE_1) {
205- // weight_quant_mode=2且kv_cache_quant_mode=1时为全量化kv量化场景298+ // weight_quant_mode=2,4,5且kv_cache_quant_mode=1时为全量化kv量化场景(int8,fp8,hif8)
206 // weight_quant_mode=3且kv_cache_quant_mode=1时为mxfp8全量化kv量化场景299 // weight_quant_mode=3且kv_cache_quant_mode=1时为mxfp8全量化kv量化场景
207- // kvcache量化300+ if (is_hifloat8_dtype) {
208- query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));301+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kByte));
302+ } else {
303+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));
304+ }
209 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_uk.size(0), 1}, at::kFloat);305 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_uk.size(0), 1}, at::kFloat);
210 } else {306 } else {
211- query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(rope_sin.dtype()));307+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kBFloat16));
212 }308 }
213 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), rope_sin.size(2)}, at::kBFloat16);309 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_uk.size(0), rope_sin.size(2)}, at::kBFloat16);
214 if (query_norm_flag) {310 if (query_norm_flag) {
215- query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));311+ if (is_hifloat8_dtype) {
216- if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2) { // weight_quant_mode=1 半量化场景,weight_quant_mode=2 全量化场景312+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_dq.size(1)}, token_x.options().dtype(at::kByte));
313+ } else {
314+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), token_x.size(1), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));
315+ }
316+ if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) {
317+ // weight_quant_mode=1 半量化场景,weight_quant_mode=2,4,5 全量化场景(int8,fp8,hif8)
217 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), 1}, at::kFloat);318 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), 1}, at::kFloat);
218 } else if (weight_quant_mode == MODE_3) {319 } else if (weight_quant_mode == MODE_3) {
219- dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x_tensor.options().dtype(dequant_scale_x_tensor.dtype()));320+ dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0) * token_x.size(1), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x.value().options().dtype(dequant_scale_x.value().dtype()));
220 }321 }
221 }322 }
222 } else {323 } else {
223 TORCH_CHECK(rope_sin_dim == DIM_2, "when token_x dim num is 2, rope_sin dim num should be 2, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));324 TORCH_CHECK(rope_sin_dim == DIM_2, "when token_x dim num is 2, rope_sin dim num should be 2, but the actual value is ", rope_sin_dim, OPS_ERROR(ErrCode::PARAM));
224- if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3) && kv_cache_quant_mode == MODE_1) {325+ if ((weight_quant_mode == MODE_2 || weight_quant_mode == MODE_3 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) && kv_cache_quant_mode == MODE_1) {
225- // weight_quant_mode=2且kv_cache_quant_mode=1时为全量化kv量化场景326+ if (is_hifloat8_dtype) {
226- // weight_quant_mode=3且kv_cache_quant_mode=1时为mxfp8全量化kv量化场景327+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kByte));
227- // kvcache量化328+ } else {
228- query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));329+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(token_x.dtype()));
330+ }
229 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), 1}, at::kFloat);331 dequant_scale_q_nope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), 1}, at::kFloat);
230 } else {332 } else {
231- query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(rope_sin.dtype()));333+ query = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), weight_uk.size(2)}, token_x.options().dtype(at::kBFloat16));
232 }334 }
233 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), rope_sin.size(1)}, at::kBFloat16);335 query_rope = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_uk.size(0), rope_sin.size(1)}, at::kBFloat16);
234 if (query_norm_flag) {336 if (query_norm_flag) {
235- query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));337+ if (is_hifloat8_dtype) {
236- if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2) { // weight_quant_mode=1 半量化场景,weight_quant_mode=2 全量化场景338+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1)}, token_x.options().dtype(at::kByte));
339+ } else {
340+ query_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1)}, token_x.options().dtype(weight_uq_qr.dtype()));
341+ }
342+ if (weight_quant_mode == MODE_1 || weight_quant_mode == MODE_2 || weight_quant_mode == MODE_4 || weight_quant_mode == MODE_5) {
237 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), 1}, at::kFloat);343 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), 1}, at::kFloat);
238 } else if (weight_quant_mode == MODE_3) {344 } else if (weight_quant_mode == MODE_3) {
239- dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x_tensor.options().dtype(dequant_scale_x_tensor.dtype()));345+ dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({token_x.size(0), weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE}, dequant_scale_x.value().options().dtype(dequant_scale_x.value().dtype()));
240 }346 }
241 }347 }
242 }348 }
@@ -245,48 +351,55 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
245 at::Tensor kv_cache_inplace = kv_cache.clone();351 at::Tensor kv_cache_inplace = kv_cache.clone();
246 at::Tensor kr_cache_inplace = kr_cache.clone();352 at::Tensor kr_cache_inplace = kr_cache.clone();
247 353 
248- if (weight_quant_mode == MODE_3) {354+ if (is_cann_version_gte_required) {
249- const at::Tensor& dequant_scale_w_dq_tensor = c10::value_or_else(dequant_scale_w_dq, [] {return at::Tensor();});355+ if (is_hifloat8_dtype) {
250- const at::Tensor& dequant_scale_w_uq_qr_tensor = c10::value_or_else(dequant_scale_w_uq_qr, [] {return at::Tensor();});356+ TensorWrapper token_x_wrapper = make_wrapper(token_x, token_x_dtype);
251- const at::Tensor& dequant_scale_w_dkv_kr_tensor = c10::value_or_else(dequant_scale_w_dkv_kr, [] {return at::Tensor();});357+ TensorWrapper weight_dq_wrapper = make_wrapper(weight_dq, weight_dq_dtype);
252- 358+ TensorWrapper weight_uq_qr_wrapper = make_wrapper(weight_uq_qr, weight_uq_qr_dtype);
253- TensorWrapper dequant_scale_x_wrapper = {dequant_scale_x_tensor, aclDataType::ACL_FLOAT8_E8M0};359+ TensorWrapper weight_dkv_kr_wrapper = make_wrapper(weight_dkv_kr, weight_dkv_kr_dtype);
254- TensorWrapper dequant_scale_w_dq_wrapper = {dequant_scale_w_dq_tensor, aclDataType::ACL_FLOAT8_E8M0};360+ if (kv_cache_quant_mode == MODE_1) {
255- TensorWrapper dequant_scale_w_uq_qr_wrapper = {dequant_scale_w_uq_qr_tensor, aclDataType::ACL_FLOAT8_E8M0};361+ TensorWrapper kv_cache_wrapper = make_wrapper(kv_cache_inplace, kv_cache_dtype);
256- TensorWrapper dequant_scale_w_dkv_kr_wrapper = {dequant_scale_w_dkv_kr_tensor, aclDataType::ACL_FLOAT8_E8M0};362+ TensorWrapper query_wrapper = make_wrapper(query, token_x_dtype);
257- TensorWrapper dequant_scale_q_norm_wrapper = {dequant_scale_q_norm, aclDataType::ACL_FLOAT8_E8M0};363+ TensorWrapper query_norm_wrapper = make_wrapper(query_norm, token_x_dtype);
258- 364+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x_wrapper, weight_dq_wrapper, weight_uq_qr_wrapper, weight_uk, weight_dkv_kr_wrapper, rmsnorm_gamma_cq,
259- if (is_cann_version_gte_required) {365+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_wrapper, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,
260- EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,366+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
261- rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x_wrapper, dequant_scale_w_dq_wrapper,367+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
262- dequant_scale_w_uq_qr_wrapper, dequant_scale_w_dkv_kr_wrapper, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,368+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query_wrapper, query_rope,
263- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,369+ dequant_scale_q_nope, query_norm_wrapper, dequant_scale_q_norm);
264- ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,370+ } else if (kv_cache_quant_mode == MODE_3) {
265- dequant_scale_q_nope, query_norm, dequant_scale_q_norm_wrapper);371+ TensorWrapper kv_cache_wrapper = make_wrapper(kv_cache_inplace, kv_cache_dtype);
266- } else {372+ TensorWrapper query_norm_wrapper = make_wrapper(query_norm, token_x_dtype);
267- EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,373+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x_wrapper, weight_dq_wrapper, weight_uq_qr_wrapper, weight_uk, weight_dkv_kr_wrapper, rmsnorm_gamma_cq,
268- rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x_wrapper, dequant_scale_w_dq_wrapper,374+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_wrapper, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,
269- dequant_scale_w_uq_qr_wrapper, dequant_scale_w_dkv_kr_wrapper, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,375+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
270- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, query_norm_flag, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,376+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
271- ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,377+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
272- dequant_scale_q_nope, query_norm, dequant_scale_q_norm_wrapper);378+ dequant_scale_q_nope, query_norm_wrapper, dequant_scale_q_norm);
273- }379+ } else {
274- } else {380+ TensorWrapper query_norm_wrapper = make_wrapper(query_norm, token_x_dtype);
275- if (is_cann_version_gte_required) {381+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x_wrapper, weight_dq_wrapper, weight_uq_qr_wrapper, weight_uk, weight_dkv_kr_wrapper, rmsnorm_gamma_cq,
276- EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,382+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,
277- rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,383+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
278- dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,384+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
279- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,385+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
280- ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,386+ dequant_scale_q_nope, query_norm_wrapper, dequant_scale_q_norm);
281- dequant_scale_q_nope, query_norm, dequant_scale_q_norm);387+ }
282 } else {388 } else {
283 EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,389 EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,
284 rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,390 rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,
285 dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,391 dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
286- k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, query_norm_flag, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,392+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
287 ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,393 ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
288 dequant_scale_q_nope, query_norm, dequant_scale_q_norm);394 dequant_scale_q_nope, query_norm, dequant_scale_q_norm);
289 }395 }
396+ } else {
397+ EXEC_NPU_CMD(aclnnMlaPrologV3WeightNz, token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq,
398+ rmsnorm_gamma_ckv, rope_sin, rope_cos, kv_cache_inplace, kr_cache_inplace, cache_index, dequant_scale_x, dequant_scale_w_dq,
399+ dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr, quant_scale_ckv, quant_scale_ckr, smooth_scales_cq, actual_seq_len,
400+ k_nope_clip_alpha, rmsnorm_epsilon_cq, rmsnorm_epsilon_ckv, cache_mode_ptr, query_norm_flag, weight_quant_mode, kv_cache_quant_mode, query_quant_mode,
401+ ckvkr_repo_mode, quant_scale_repo_mode, tile_size, qc_qr_scale, kc_scale, query, query_rope,
402+ dequant_scale_q_nope, query_norm, dequant_scale_q_norm);
290 }403 }
291 404 
292 if (!query_norm.defined()) {405 if (!query_norm.defined()) {
@@ -297,7 +410,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
297 }410 }
298 if (!dequant_scale_q_norm.defined()) {411 if (!dequant_scale_q_norm.defined()) {
299 if (weight_quant_mode == MODE_3) {412 if (weight_quant_mode == MODE_3) {
300- dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, dequant_scale_x_tensor.options().dtype(dequant_scale_x_tensor.dtype()));413+ dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, dequant_scale_x.value().options().dtype(dequant_scale_x.value().dtype()));
301 } else {414 } else {
302 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, at::kFloat);415 dequant_scale_q_norm = npu_preparation::apply_tensor_without_format({0}, at::kFloat);
303 }416 }
Mop_plugin/python/meta/_meta_registrations.py+105-26
@@ -423,8 +423,8 @@ def npu_mla_prolog_v2_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weigh
423def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq, rmsnorm_gamma_ckv,423def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq, rmsnorm_gamma_ckv,
424 rope_sin, rope_cos, kv_cache, kr_cache, *, cache_index=None, dequant_scale_x=None, dequant_scale_w_dq=None, dequant_scale_w_uq_qr=None, dequant_scale_w_dkv_kr=None,424 rope_sin, rope_cos, kv_cache, kr_cache, *, cache_index=None, dequant_scale_x=None, dequant_scale_w_dq=None, dequant_scale_w_uq_qr=None, dequant_scale_w_dkv_kr=None,
425 quant_scale_ckv=None, quant_scale_ckr=None, smooth_scales_cq=None, actual_seq_len=None, k_nope_clip_alpha=None, rmsnorm_epsilon_cq=1e-5, rmsnorm_epsilon_ckv=1e-5,425 quant_scale_ckv=None, quant_scale_ckr=None, smooth_scales_cq=None, actual_seq_len=None, k_nope_clip_alpha=None, rmsnorm_epsilon_cq=1e-5, rmsnorm_epsilon_ckv=1e-5,
426- cache_mode="PA_BSND", query_norm_flag=False, weight_quant_mode=0, kv_cache_quant_mode=0, query_quant_mode=0, ckvkr_repo_mode=0, quant_scale_repo_mode=0, tile_size=128, qc_qr_scale=1.0, kc_scale=1.0):426+ cache_mode="PA_BSND", query_norm_flag=False, weight_quant_mode=0, kv_cache_quant_mode=0, query_quant_mode=0, ckvkr_repo_mode=0, quant_scale_repo_mode=0, tile_size=128,
427- 427+ qc_qr_scale=1.0, kc_scale=1.0, token_x_dtype=None, weight_dq_dtype=None, weight_uq_qr_dtype=None, weight_dkv_kr_dtype=None, kv_cache_dtype=None):
428 require_param = {"token_x": token_x, "weight_dq": weight_dq, "weight_uq_qr": weight_uq_qr, "weight_uk": weight_uk, "weight_dkv_kr": weight_dkv_kr, "rmsnorm_gamma_cq": rmsnorm_gamma_cq, "rmsnorm_gamma_ckv": rmsnorm_gamma_ckv, "rope_sin": rope_sin, "rope_cos": rope_cos, "kv_cache": kv_cache, "kr_cache": kr_cache}428 require_param = {"token_x": token_x, "weight_dq": weight_dq, "weight_uq_qr": weight_uq_qr, "weight_uk": weight_uk, "weight_dkv_kr": weight_dkv_kr, "rmsnorm_gamma_cq": rmsnorm_gamma_cq, "rmsnorm_gamma_ckv": rmsnorm_gamma_ckv, "rope_sin": rope_sin, "rope_cos": rope_cos, "kv_cache": kv_cache, "kr_cache": kr_cache}
429 429 
430 if weight_quant_mode == 3:430 if weight_quant_mode == 3:
@@ -433,6 +433,34 @@ def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weigh
433 lambda: "When weight_quant_mode is 3, not support on this soc version." + ops_error(ErrCode.NOT_SUPPORT),433 lambda: "When weight_quant_mode is 3, not support on this soc version." + ops_error(ErrCode.NOT_SUPPORT),
434 )434 )
435 435 
436+ torch._check(
437+ dequant_scale_x.dtype == torch.float8_e8m0fnu and dequant_scale_w_dq.dtype == torch.float8_e8m0fnu and
438+ dequant_scale_w_uq_qr.dtype == torch.float8_e8m0fnu and dequant_scale_w_dkv_kr.dtype == torch.float8_e8m0fnu,
439+ lambda: "torch_npu supports the float8_e8m0 only in version later than v2.7., dequant_scale_x, dequant_scale_w_dq, dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr Dtype in weight_quant_mode=3 must be torch.float8_e8m0fnu",
440+ )
441+ 
442+ is_hifloat8_dtype = False
443+ if weight_quant_mode == 5 and (token_x.dtype == torch.uint8 or weight_dq.dtype == torch.uint8 or weight_uq_qr.dtype == torch.uint8 or weight_dkv_kr.dtype == torch.uint8): # hifloat8伪装为uint8
444+ torch._check(
445+ token_x_dtype is not None and weight_dq_dtype is not None and weight_uq_qr_dtype is not None and weight_dkv_kr_dtype is not None,
446+ lambda: "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype,weight_dq_dtype,weight_uq_qr_dtype,weight_dkv_kr_dtype cannot be null."
447+ )
448+ torch._check(
449+ token_x_dtype == torch_npu.hifloat8 and weight_dq_dtype == torch_npu.hifloat8 and
450+ weight_uq_qr_dtype == torch_npu.hifloat8 and weight_dkv_kr_dtype == torch_npu.hifloat8,
451+ lambda: "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype, weight_dq_dtype, weight_uq_qr_dtype, weight_dkv_kr_dtype value must be torch_npu.hifloat8",
AtlasAccount
AtlasAccountAtlasAccount3月18日
代码质量: 错误消息字符串中使用了中文顿号（、）作为分隔符，这与第443行的问题相同。应该保持整个代码库中错误消息格式的一致性。
问题类型: 代码质量
文件路径: op_plugin/python/meta/_meta_registrations.py
行号: 451
问题代码:
"when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype, weight_dq_dtype, weight_uq_qr_dtype, weight_dkv_kr_dtype value must be torch_npu.hifloat8"
修改建议:
将中文顿号替换为英文逗号："when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype, weight_dq_dtype, weight_uq_qr_dtype, weight_dkv_kr_dtype value must be torch_npu.hifloat8"
---
此评论由代码审查工具自动生成
likedislike
452+ )
453+ if kv_cache_quant_mode == 1 or kv_cache_quant_mode == 3:
454+ torch._check(
455+ kv_cache_dtype is not None,
456+ lambda: "when weight_quant_mode is 5 and kv_cache_quant_mode is 1 or 3 and input dtype is hifloat8, kv_cache_dtype cannot be null."
457+ )
458+ torch._check(
459+ kv_cache_dtype == torch_npu.hifloat8,
460+ lambda: "when weight_quant_mode is 5 and input dtype is hifloat8, kv_cache_dtype value must be torch_npu.hifloat8",
461+ )
462+ is_hifloat8_dtype = True
463+ 
436 for item_name, item in require_param.items():464 for item_name, item in require_param.items():
437 torch._check(465 torch._check(
438 item is not None,466 item is not None,
@@ -486,7 +514,6 @@ def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weigh
486 dequant_scale_q_norm_shape.append(int(weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE))514 dequant_scale_q_norm_shape.append(int(weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE))
487 else:515 else:
488 dequant_scale_q_norm_shape.append(1)516 dequant_scale_q_norm_shape.append(1)
489- 
490 else:517 else:
491 torch._check(518 torch._check(
492 rope_sin_dim == 2,519 rope_sin_dim == 2,
@@ -518,27 +545,35 @@ def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weigh
518 dequant_scale_q_norm_shape.append(int(weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE))545 dequant_scale_q_norm_shape.append(int(weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE))
519 else:546 else:
520 dequant_scale_q_norm_shape.append(1)547 dequant_scale_q_norm_shape.append(1)
521-
522 548 
523 is_cann_version_gte_required = torch_npu.npu.utils._is_gte_cann_version("8.5.0.alpha003", "CANN") # whether cann version >= 8.5.0.alpha003549 is_cann_version_gte_required = torch_npu.npu.utils._is_gte_cann_version("8.5.0.alpha003", "CANN") # whether cann version >= 8.5.0.alpha003
524 # kvcache量化550 # kvcache量化
525 if weight_quant_mode == 3 and kv_cache_quant_mode == 1:551 if weight_quant_mode == 3 and kv_cache_quant_mode == 1:
526 query = torch.empty(query_shape, dtype=torch.float8_e4m3fn, device='meta')552 query = torch.empty(query_shape, dtype=torch.float8_e4m3fn, device='meta')
527 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')553 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')
528- elif weight_quant_mode == 2 and kv_cache_quant_mode == 1:554+ elif (weight_quant_mode == 2 or weight_quant_mode == 4 or weight_quant_mode == 5) and kv_cache_quant_mode == 1:
529- query = torch.empty(query_shape, dtype=torch.int8, device='meta')555+ if is_hifloat8_dtype:
556+ query_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[token_x_dtype]
557+ query = torch.empty(query_shape, dtype=query_type, device='meta')
558+ else:
559+ query = torch.empty(query_shape, dtype=token_x.dtype, device='meta')
530 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')560 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')
531 else:561 else:
532- query = torch.empty(query_shape, dtype=rope_sin.dtype, device='meta')562+ query = torch.empty(query_shape, dtype=torch.bfloat16, device='meta')
533 if is_cann_version_gte_required:563 if is_cann_version_gte_required:
534 dequant_scale_q_nope = torch.empty([0], dtype=torch.float32, device='meta')564 dequant_scale_q_nope = torch.empty([0], dtype=torch.float32, device='meta')
535 else:565 else:
536 dequant_scale_q_nope = torch.empty([1], dtype=torch.float32, device='meta')566 dequant_scale_q_nope = torch.empty([1], dtype=torch.float32, device='meta')
567+ query_rope = torch.empty(query_rope_shape, dtype=torch.bfloat16, device='meta') # default dtype bfloat16
537 568 
538 # 输出query_norm569 # 输出query_norm
539 if query_norm_flag:570 if query_norm_flag:
540- query_norm = torch.empty(query_norm_shape, dtype=weight_uq_qr.dtype, device='meta')571+ if is_hifloat8_dtype:
541- if weight_quant_mode == 1 or weight_quant_mode == 2:572+ query_norm_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[token_x_dtype]
573+ query_norm = torch.empty(query_norm_shape, dtype=query_norm_type, device='meta')
574+ else:
575+ query_norm = torch.empty(query_norm_shape, dtype=weight_uq_qr.dtype, device='meta')
576+ if weight_quant_mode == 1 or weight_quant_mode == 2 or weight_quant_mode == 4 or weight_quant_mode == 5:
542 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float32, device='meta')577 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float32, device='meta')
543 elif weight_quant_mode == 3:578 elif weight_quant_mode == 3:
544 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float8_e8m0fnu, device='meta')579 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float8_e8m0fnu, device='meta')
@@ -549,13 +584,15 @@ def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weigh
549 dequant_scale_q_norm = torch.empty([1], dtype=torch.float32, device='meta')584 dequant_scale_q_norm = torch.empty([1], dtype=torch.float32, device='meta')
550 else:585 else:
551 if is_cann_version_gte_required:586 if is_cann_version_gte_required:
552- query_norm = torch.empty([0], dtype=weight_uq_qr.dtype, device='meta')587+ if is_hifloat8_dtype:
588+ query_norm_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[token_x_dtype]
589+ query_norm = torch.empty([0], dtype=query_norm_type, device='meta')
590+ else:
591+ query_norm = torch.empty([0], dtype=weight_uq_qr.dtype, device='meta')
553 dequant_scale_q_norm = torch.empty([0], dtype=torch.float8_e8m0fnu if weight_quant_mode == 3 else torch.float32, device='meta')592 dequant_scale_q_norm = torch.empty([0], dtype=torch.float8_e8m0fnu if weight_quant_mode == 3 else torch.float32, device='meta')
554 else:593 else:
555 query_norm = torch.empty([1], dtype=weight_uq_qr.dtype, device='meta')594 query_norm = torch.empty([1], dtype=weight_uq_qr.dtype, device='meta')
556- dequant_scale_q_norm = torch.empty([1], dtype=torch.float8_e8m0fnu if weight_quant_mode == 3 else torch.float32, device='meta')595+ dequant_scale_q_norm = torch.empty([1], dtype=torch.float32, device='meta')
557- 
558- query_rope = torch.empty(query_rope_shape, dtype=torch.bfloat16, device='meta') # default dtype bfloat16
559 596 
560 return (query, query_rope, dequant_scale_q_nope, query_norm, dequant_scale_q_norm)597 return (query, query_rope, dequant_scale_q_nope, query_norm, dequant_scale_q_norm)
561 598 
@@ -564,9 +601,8 @@ def npu_mla_prolog_v3_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weigh
564def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq, rmsnorm_gamma_ckv,601def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weight_uk, weight_dkv_kr, rmsnorm_gamma_cq, rmsnorm_gamma_ckv,
565 rope_sin, rope_cos, kv_cache, kr_cache, *, cache_index=None, dequant_scale_x=None, dequant_scale_w_dq=None, dequant_scale_w_uq_qr=None, dequant_scale_w_dkv_kr=None,602 rope_sin, rope_cos, kv_cache, kr_cache, *, cache_index=None, dequant_scale_x=None, dequant_scale_w_dq=None, dequant_scale_w_uq_qr=None, dequant_scale_w_dkv_kr=None,
566 quant_scale_ckv=None, quant_scale_ckr=None, smooth_scales_cq=None, actual_seq_len=None, k_nope_clip_alpha=None, rmsnorm_epsilon_cq=1e-5, rmsnorm_epsilon_ckv=1e-5,603 quant_scale_ckv=None, quant_scale_ckr=None, smooth_scales_cq=None, actual_seq_len=None, k_nope_clip_alpha=None, rmsnorm_epsilon_cq=1e-5, rmsnorm_epsilon_ckv=1e-5,
567- cache_mode="PA_BSND", query_norm_flag=False, weight_quant_mode=0, kv_cache_quant_mode=0, query_quant_mode=0, ckvkr_repo_mode=0, quant_scale_repo_mode=0, tile_size=128, qc_qr_scale=1.0, kc_scale=1.0):604+ cache_mode="PA_BSND", query_norm_flag=False, weight_quant_mode=0, kv_cache_quant_mode=0, query_quant_mode=0, ckvkr_repo_mode=0, quant_scale_repo_mode=0, tile_size=128,
568- 605+ qc_qr_scale=1.0, kc_scale=1.0, token_x_dtype=None, weight_dq_dtype=None, weight_uq_qr_dtype=None, weight_dkv_kr_dtype=None, kv_cache_dtype=None):
569- 
570 require_param = {"token_x": token_x, "weight_dq": weight_dq, "weight_uq_qr": weight_uq_qr, "weight_uk": weight_uk, "weight_dkv_kr": weight_dkv_kr, "rmsnorm_gamma_cq": rmsnorm_gamma_cq, "rmsnorm_gamma_ckv": rmsnorm_gamma_ckv, "rope_sin": rope_sin, "rope_cos": rope_cos, "kv_cache": kv_cache, "kr_cache": kr_cache}606 require_param = {"token_x": token_x, "weight_dq": weight_dq, "weight_uq_qr": weight_uq_qr, "weight_uk": weight_uk, "weight_dkv_kr": weight_dkv_kr, "rmsnorm_gamma_cq": rmsnorm_gamma_cq, "rmsnorm_gamma_ckv": rmsnorm_gamma_ckv, "rope_sin": rope_sin, "rope_cos": rope_cos, "kv_cache": kv_cache, "kr_cache": kr_cache}
571 607 
572 if weight_quant_mode == 3:608 if weight_quant_mode == 3:
@@ -575,6 +611,34 @@ def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weigh
575 lambda: "When weight_quant_mode is 3, not support on this soc version." + ops_error(ErrCode.NOT_SUPPORT),611 lambda: "When weight_quant_mode is 3, not support on this soc version." + ops_error(ErrCode.NOT_SUPPORT),
576 )612 )
577 613 
614+ torch._check(
615+ dequant_scale_x.dtype == torch.float8_e8m0fnu and dequant_scale_w_dq.dtype == torch.float8_e8m0fnu and
616+ dequant_scale_w_uq_qr.dtype == torch.float8_e8m0fnu and dequant_scale_w_dkv_kr.dtype == torch.float8_e8m0fnu,
617+ lambda: "torch_npu supports the float8_e8m0 only in version later than v2.7., dequant_scale_x, dequant_scale_w_dq, dequant_scale_w_uq_qr, dequant_scale_w_dkv_kr Dtype in weight_quant_mode=3 must be torch.float8_e8m0fnu",
618+ )
619+ 
620+ is_hifloat8_dtype = False
621+ if weight_quant_mode == 5 and (token_x.dtype == torch.uint8 or weight_dq.dtype == torch.uint8 or weight_uq_qr.dtype == torch.uint8 or weight_dkv_kr.dtype == torch.uint8): # hifloat8伪装为uint8
622+ torch._check(
623+ token_x_dtype is not None and weight_dq_dtype is not None and weight_uq_qr_dtype is not None and weight_dkv_kr_dtype is not None,
624+ lambda: "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype,weight_dq_dtype,weight_uq_qr_dtype,weight_dkv_kr_dtype cannot be null."
625+ )
626+ torch._check(
627+ token_x_dtype == torch_npu.hifloat8 and weight_dq_dtype == torch_npu.hifloat8 and
628+ weight_uq_qr_dtype == torch_npu.hifloat8 and weight_dkv_kr_dtype == torch_npu.hifloat8,
629+ lambda: "when weight_quant_mode is 5 and input dtype is hifloat8, token_x_dtype, weight_dq_dtype, weight_uq_qr_dtype, weight_dkv_kr_dtype value must be torch_npu.hifloat8",
630+ )
631+ if kv_cache_quant_mode == 1 or kv_cache_quant_mode == 3:
632+ torch._check(
633+ kv_cache_dtype is not None,
634+ lambda: "when weight_quant_mode is 5 and kv_cache_quant_mode is 1 or 3 and input dtype is hifloat8, kv_cache_dtype cannot be null."
635+ )
636+ torch._check(
637+ kv_cache_dtype == torch_npu.hifloat8,
638+ lambda: "when weight_quant_mode is 5 and input dtype is hifloat8, kv_cache_dtype value must be torch_npu.hifloat8",
639+ )
640+ is_hifloat8_dtype = True
641+ 
578 for item_name, item in require_param.items():642 for item_name, item in require_param.items():
579 torch._check(643 torch._check(
580 item is not None,644 item is not None,
@@ -628,7 +692,6 @@ def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weigh
628 dequant_scale_q_norm_shape.append(int(weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE))692 dequant_scale_q_norm_shape.append(int(weight_dq.size(1) / FP8_E4M3_BLOCK_SIZE))
629 else:693 else:
630 dequant_scale_q_norm_shape.append(1)694 dequant_scale_q_norm_shape.append(1)
631- 
632 else:695 else:
633 torch._check(696 torch._check(
634 rope_sin_dim == 2,697 rope_sin_dim == 2,
@@ -666,11 +729,15 @@ def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weigh
666 if weight_quant_mode == 3 and kv_cache_quant_mode == 1:729 if weight_quant_mode == 3 and kv_cache_quant_mode == 1:
667 query = torch.empty(query_shape, dtype=torch.float8_e4m3fn, device='meta')730 query = torch.empty(query_shape, dtype=torch.float8_e4m3fn, device='meta')
668 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')731 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')
669- elif weight_quant_mode == 2 and kv_cache_quant_mode == 1:732+ elif (weight_quant_mode == 2 or weight_quant_mode == 4 or weight_quant_mode == 5) and kv_cache_quant_mode == 1:
670- query = torch.empty(query_shape, dtype=torch.int8, device='meta')733+ if is_hifloat8_dtype:
734+ query_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[token_x_dtype]
735+ query = torch.empty(query_shape, dtype=query_type, device='meta')
736+ else:
737+ query = torch.empty(query_shape, dtype=token_x.dtype, device='meta')
671 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')738 dequant_scale_q_nope = torch.empty(dequant_scale_q_nope_shape, dtype=torch.float32, device='meta')
672 else:739 else:
673- query = torch.empty(query_shape, dtype=rope_sin.dtype, device='meta')740+ query = torch.empty(query_shape, dtype=torch.bfloat16, device='meta')
674 if is_cann_version_gte_required:741 if is_cann_version_gte_required:
675 dequant_scale_q_nope = torch.empty([0], dtype=torch.float32, device='meta')742 dequant_scale_q_nope = torch.empty([0], dtype=torch.float32, device='meta')
676 else:743 else:
@@ -680,9 +747,13 @@ def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weigh
680 747 
681 # 输出query_norm748 # 输出query_norm
682 if query_norm_flag:749 if query_norm_flag:
683- query_norm = torch.empty(query_norm_shape, dtype=weight_uq_qr.dtype, device='meta')750+ if is_hifloat8_dtype:
751+ query_norm_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[token_x_dtype]
752+ query_norm = torch.empty(query_norm_shape, dtype=query_norm_type, device='meta')
753+ else:
754+ query_norm = torch.empty(query_norm_shape, dtype=weight_uq_qr.dtype, device='meta')
684 # 动态量化755 # 动态量化
685- if weight_quant_mode == 1 or weight_quant_mode == 2:756+ if weight_quant_mode == 1 or weight_quant_mode == 2 or weight_quant_mode == 4 or weight_quant_mode == 5:
686 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float32, device='meta')757 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float32, device='meta')
687 elif weight_quant_mode == 3:758 elif weight_quant_mode == 3:
688 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float8_e8m0fnu, device='meta')759 dequant_scale_q_norm = torch.empty(dequant_scale_q_norm_shape, dtype=torch.float8_e8m0fnu, device='meta')
@@ -693,13 +764,21 @@ def npu_mla_prolog_v3_functional_forward(token_x, weight_dq, weight_uq_qr, weigh
693 dequant_scale_q_norm = torch.empty([1], dtype=torch.float32, device='meta')764 dequant_scale_q_norm = torch.empty([1], dtype=torch.float32, device='meta')
694 else:765 else:
695 if is_cann_version_gte_required:766 if is_cann_version_gte_required:
696- query_norm = torch.empty([0], dtype=weight_uq_qr.dtype, device='meta')767+ if is_hifloat8_dtype:
768+ query_norm_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[token_x_dtype]
769+ query_norm = torch.empty([0], dtype=query_norm_type, device='meta')
770+ else:
771+ query_norm = torch.empty([0], dtype=weight_uq_qr.dtype, device='meta')
697 dequant_scale_q_norm = torch.empty([0], dtype=torch.float8_e8m0fnu if weight_quant_mode == 3 else torch.float32, device='meta')772 dequant_scale_q_norm = torch.empty([0], dtype=torch.float8_e8m0fnu if weight_quant_mode == 3 else torch.float32, device='meta')
698 else:773 else:
699 query_norm = torch.empty([1], dtype=weight_uq_qr.dtype, device='meta')774 query_norm = torch.empty([1], dtype=weight_uq_qr.dtype, device='meta')
700- dequant_scale_q_norm = torch.empty([1], dtype=torch.float8_e8m0fnu if weight_quant_mode == 3 else torch.float32, device='meta')775+ dequant_scale_q_norm = torch.empty([1], dtype=torch.float32, device='meta')
701 776 
702- kv_cache_out = torch.empty_like(kv_cache, dtype=kv_cache.dtype, device='meta')777+ if is_hifloat8_dtype and (kv_cache_quant_mode == 1 or kv_cache_quant_mode == 3):
778+ kv_cache_type = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP[kv_cache_dtype]
779+ kv_cache_out = torch.empty_like(kv_cache, dtype=kv_cache_type, device='meta')
780+ else:
781+ kv_cache_out = torch.empty_like(kv_cache, dtype=kv_cache.dtype, device='meta')
703 kr_cache_out = torch.empty_like(kr_cache, dtype=kr_cache.dtype, device='meta')782 kr_cache_out = torch.empty_like(kr_cache, dtype=kr_cache.dtype, device='meta')
704 783 
705 return (query, query_rope, dequant_scale_q_nope, query_norm, dequant_scale_q_norm, kv_cache_out, kr_cache_out)784 return (query, query_rope, dequant_scale_q_nope, query_norm, dequant_scale_q_norm, kv_cache_out, kr_cache_out)
@@ -1175,7 +1254,7 @@ def npu_moe_init_routing_v2_meta(x, expert_idx, *, scale=None, offset=None, acti
1175 scale_dim == 1,1254 scale_dim == 1,
1176 lambda: "the scale shape support only 1D (bs,) in no quant mode" + ops_error(ErrCode.VALUE),1255 lambda: "the scale shape support only 1D (bs,) in no quant mode" + ops_error(ErrCode.VALUE),
1177 )1256 )
1178- # else: quant_mode为2、3、8时不使用scale也不需要校验1257+ # else: quant_mode为2,3,8时不使用scale也不需要校验
1179 1258 
1180 bs = x.size(0)1259 bs = x.size(0)
1181 h = x.size(1)1260 h = x.size(1)