已合并
[feature]FA MXFP8 特性提交，新注册标识字符串“MXFP8_DYNAMIC”，并同步提交量化代码类 #375
乐加维林创建于 6月22日
[feature]FA MXFP8 特性提交，新注册标识字符串“MXFP8_DYNAMIC”，并同步提交量化代码类 #375
已合并
共 4 个文件变更+87-0
| @@ -494,6 +494,76 @@ class FP8RotateQuantFA(nn.Module): | |||
| 494 | return x | 494 | return x |
| 495 | 495 | ||
| 496 | 496 | ||
| 497 | +class MXFP8RotateQuantFA(nn.Module): | ||
| 498 | + def __init__(self, prefix=None, weights=None): | ||
| 499 | + super().__init__() | ||
| 500 | + | ||
| 501 | + q_rot = get_quant_weight(weights, f'{prefix}.q_rot') | ||
| 502 | + self.register_buffer("q_rot", q_rot, persistent=False) | ||
| 503 | + k_rot = get_quant_weight(weights, f'{prefix}.k_rot') | ||
| 504 | + self.register_buffer("k_rot", k_rot, persistent=False) | ||
| 505 | + | ||
| 506 | + def forward(self, query, key, value, **kwargs): | ||
| 507 | + query = torch.matmul(query, self.q_rot) | ||
| 508 | + key = torch.matmul(key, self.k_rot) | ||
| 509 | + | ||
| 510 | + layout = kwargs.get("layout", "BNSD") | ||
| 511 | + if layout == "BNSD": | ||
| 512 | + b, n, s, d = query.shape | ||
| 513 | + query = query.permute(0, 2, 1, 3).reshape(b * s, n, d) | ||
| 514 | + key = key.permute(0, 2, 1, 3).reshape(b * s, n, d) | ||
| 515 | + value = value.permute(0, 2, 1, 3).reshape(b * s, n, d) | ||
| 516 | + elif layout == "BSND": | ||
| 517 | + b, s, n, d = query.shape | ||
| 518 | + query = query.reshape(b * s, n, d) | ||
| 519 | + key = key.reshape(b * s, n, d) | ||
| 520 | + value = value.reshape(b * s, n, d) | ||
| 521 | + else: | ||
| 522 | + raise ValueError(f"Unsupported layout: {layout}, expected 'BNSD' or 'BSND'.") | ||
| 523 | + | ||
| 524 | + actual_seq_qlen = torch.arange(s, s * (b + 1), s, dtype=torch.int64, device=query.device) | ||
| 525 | + actual_seq_kvlen = torch.arange(s, s * (b + 1), s, dtype=torch.int64, device=key.device) | ||
| 526 | + | ||
| 527 | + q, q_scale = torch_npu.npu_dynamic_mx_quant(query, dst_type=torch.float8_e4m3fn, axis=-1) | ||
| 528 | + k, k_scale = torch_npu.npu_dynamic_mx_quant(key, dst_type=torch.float8_e4m3fn, axis=-1) | ||
| 529 | + v, v_scale = torch_npu.npu_dynamic_mx_quant(value, dst_type=torch.float8_e4m3fn, axis=0) | ||
| 530 | + | ||
| 531 | + x = torch_npu.npu_fused_infer_attention_score_v2( | ||
| 532 | + q, | ||
| 533 | + k, | ||
| 534 | + v, | ||
| 535 | + input_layout="TND", | ||
| 536 | + num_query_heads=n, | ||
| 537 | + num_key_value_heads=n, | ||
| 538 | + softmax_scale=1.0 / math.sqrt(d), | ||
| 539 | + dequant_scale_query=q_scale, | ||
| 540 | + dequant_scale_key=k_scale, | ||
| 541 | + dequant_scale_value=v_scale, | ||
| 542 | + actual_seq_qlen=actual_seq_qlen, | ||
| 543 | + actual_seq_kvlen=actual_seq_kvlen, | ||
| 544 | + sparse_mode=0, # could be 0/3, atten_mask is needed if set 3 | ||
| 545 | + query_quant_mode=6, | ||
| 546 | + key_quant_mode=6, | ||
| 547 | + value_quant_mode=8, | ||
| 548 | + query_dtype=torch.float8_e4m3fn, | ||
| 549 | + key_dtype=torch.float8_e4m3fn, | ||
| 550 | + value_dtype=torch.float8_e4m3fn, | ||
| 551 | + dequant_scale_query_dtype=torch_npu.float8_e8m0fnu, | ||
| 552 | + dequant_scale_key_dtype=torch_npu.float8_e8m0fnu, | ||
| 553 | + dequant_scale_value_dtype=torch_npu.float8_e8m0fnu, | ||
| 554 | + out_dtype=query.dtype, | ||
| 555 | + )[0] | ||
| 556 | + | ||
| 557 | + if layout == "BNSD": | ||
| 558 | + # [B*S, N, D] -> [B, S, N, D] -> [B, N, S, D] | ||
| 559 | + x = x.reshape(b, s, n, d).permute(0, 2, 1, 3) | ||
| 560 | + elif layout == "BSND": | ||
| 561 | + # [B*S, N, D] -> [B, S, N, D] | ||
| 562 | + x = x.reshape(b, s, n, d) | ||
| 563 | + | ||
| 564 | + return x | ||
| 565 | + | ||
| 566 | + | ||
| 497 | class MXFP4QuantFA(nn.Module): | 567 | class MXFP4QuantFA(nn.Module): |
| 498 | def __init__(self, prefix=None, weights=None, **kwargs): | 568 | def __init__(self, prefix=None, weights=None, **kwargs): |
| 499 | super().__init__() | 569 | super().__init__() |
| @@ -43,6 +43,7 @@ class QuantAlgorithm(StrEnum): | |||
| 43 | INT8 = "INT8" | 43 | INT8 = "INT8" |
| 44 | MIXED_PERCISION = "MIXED_PERCISION" | 44 | MIXED_PERCISION = "MIXED_PERCISION" |
| 45 | FP8_DYNAMIC = "FP8_DYNAMIC" | 45 | FP8_DYNAMIC = "FP8_DYNAMIC" |
| 46 | + MXFP8_DYNAMIC = "MXFP8_DYNAMIC" | ||
| 46 | MXFP4_DYNAMIC = "MXFP4_DYNAMIC" | 47 | MXFP4_DYNAMIC = "MXFP4_DYNAMIC" |
| 47 | NO_QUANT = "NO_QUANT" | 48 | NO_QUANT = "NO_QUANT" |
| 48 | W4A4_MXFP4_SVD = "W4A4_MXFP4_SVD" | 49 | W4A4_MXFP4_SVD = "W4A4_MXFP4_SVD" |
| @@ -185,6 +186,7 @@ class QuantMode: | |||
| 185 | QuantAlgorithm.W8A8_PER_CHANNEL_PER_TOKEN: QuantMode.use_smooth_quant(per_token=True, per_channel=True), | 186 | QuantAlgorithm.W8A8_PER_CHANNEL_PER_TOKEN: QuantMode.use_smooth_quant(per_token=True, per_channel=True), |
| 186 | QuantAlgorithm.W8A8_PER_TENSOR_PER_TOKEN: QuantMode.use_smooth_quant(per_token=True, per_channel=False), | 187 | QuantAlgorithm.W8A8_PER_TENSOR_PER_TOKEN: QuantMode.use_smooth_quant(per_token=True, per_channel=False), |
| 187 | QuantAlgorithm.FP8_DYNAMIC: QuantMode.from_descriptor(QuantModeDescriptor(use_fa_quant=True)), | 188 | QuantAlgorithm.FP8_DYNAMIC: QuantMode.from_descriptor(QuantModeDescriptor(use_fa_quant=True)), |
| 189 | + QuantAlgorithm.MXFP8_DYNAMIC: QuantMode.from_descriptor(QuantModeDescriptor(use_fa_quant=True)), | ||
| 188 | QuantAlgorithm.MXFP4_DYNAMIC: QuantMode.from_descriptor(QuantModeDescriptor(use_fa_quant=True)), | 190 | QuantAlgorithm.MXFP4_DYNAMIC: QuantMode.from_descriptor(QuantModeDescriptor(use_fa_quant=True)), |
| 189 | QuantAlgorithm.W8A8_MXFP8: QuantMode.use_smooth_quant(per_token=False, per_channel=False), | 191 | QuantAlgorithm.W8A8_MXFP8: QuantMode.use_smooth_quant(per_token=False, per_channel=False), |
| 190 | QuantAlgorithm.W4A4_MXFP4_SVD: QuantMode.use_smooth_quant(per_token=False, per_channel=False), | 192 | QuantAlgorithm.W4A4_MXFP4_SVD: QuantMode.use_smooth_quant(per_token=False, per_channel=False), |
| @@ -29,6 +29,7 @@ from .layer import ( | |||
| 29 | W8A8TimeStepQuantLinear, | 29 | W8A8TimeStepQuantLinear, |
| 30 | WeightQuantLinear, | 30 | WeightQuantLinear, |
| 31 | FP8RotateQuantFA, | 31 | FP8RotateQuantFA, |
| 32 | + MXFP8RotateQuantFA, | ||
| 32 | MXFP4QuantFA, | 33 | MXFP4QuantFA, |
| 33 | W8A8MXFP8QuantLinear, | 34 | W8A8MXFP8QuantLinear, |
| 34 | W4A4MXFP4QuantLinear, | 35 | W4A4MXFP4QuantLinear, |
| @@ -160,6 +161,8 @@ def add_fa_quant(layer, cfg, prefix, quant_weights, **kwargs): | |||
| 160 | layer.fa_quant = MXFP4QuantFA(prefix, quant_weights, **kwargs) | 161 | layer.fa_quant = MXFP4QuantFA(prefix, quant_weights, **kwargs) |
| 161 | elif cfg.quant_algo in [QuantAlgorithm.FP8_DYNAMIC]: | 162 | elif cfg.quant_algo in [QuantAlgorithm.FP8_DYNAMIC]: |
| 162 | layer.fa_quant = FP8RotateQuantFA(prefix, quant_weights) | 163 | layer.fa_quant = FP8RotateQuantFA(prefix, quant_weights) |
| 164 | + elif cfg.quant_algo in [QuantAlgorithm.MXFP8_DYNAMIC]: | ||
| 165 | + layer.fa_quant = MXFP8RotateQuantFA(prefix, quant_weights) | ||
| 163 | 166 | ||
| 164 | 167 | ||
| 165 | def normalize_quant_config(kwargs): | 168 | def normalize_quant_config(kwargs): |
| @@ -22,6 +22,7 @@ from mindiesd.quantization.layer import ( | |||
| 22 | W8A8QuantBaseLinear, | 22 | W8A8QuantBaseLinear, |
| 23 | WeightQuantLinear, | 23 | WeightQuantLinear, |
| 24 | FP8RotateQuantFA, | 24 | FP8RotateQuantFA, |
| 25 | + MXFP8RotateQuantFA, | ||
| 25 | W8A8MXFP8QuantLinear, | 26 | W8A8MXFP8QuantLinear, |
| 26 | W4A4QuantLinear, | 27 | W4A4QuantLinear, |
| 27 | W4A4MXFP4QuantLinear, | 28 | W4A4MXFP4QuantLinear, |
| @@ -518,6 +519,17 @@ class TestAddFAQuant(unittest.TestCase): | |||
| 518 | add_fa_quant(layer, cfg, "test_layer", self.weights) | 519 | add_fa_quant(layer, cfg, "test_layer", self.weights) |
| 519 | self.assertFalse(hasattr(layer, 'fa_quant')) | 520 | self.assertFalse(hasattr(layer, 'fa_quant')) |
| 520 | 521 | ||
| 522 | + def test_add_fa_quant_with_mxfp8_layer(self): | ||
| 523 | + # 创建一个具有必要属性的模拟层 | ||
| 524 | + class MockLayer(nn.Module): | ||
| 525 | + pass | ||
| 526 | + | ||
| 527 | + layer = MockLayer() | ||
| 528 | + cfg = QuantConfig(quant_algo=QuantAlgorithm.MXFP8_DYNAMIC) | ||
| 529 | + add_fa_quant(layer, cfg, "test_layer", create_mock_handler(self.weights)) | ||
| 530 | + self.assertTrue(hasattr(layer, 'fa_quant')) | ||
| 531 | + self.assertIsInstance(layer.fa_quant, MXFP8RotateQuantFA) | ||
| 532 | + | ||
| 521 | 533 | ||
| 522 | 534 | ||
| 523 | os.environ.get("MINDIE_TEST_MODE", "ALL") == "CPU", "Skip NPU-dependent tests when MINDIE_TEST_MODE is CPU." | 535 | os.environ.get("MINDIE_TEST_MODE", "ALL") == "CPU", "Skip NPU-dependent tests when MINDIE_TEST_MODE is CPU." |