已合并
feat(profiler): add fused op flops formulas #38360
hhz0创建于 6月12日
feat(profiler): add fused op flops formulas #38360
已合并
共 7 个文件变更+1069-0
| @@ -0,0 +1,234 @@ | |||
| 1 | +from types import SimpleNamespace | ||
| 2 | +from unittest import mock | ||
| 3 | + | ||
| 4 | +from torch_npu.profiler._flops_formulas import ( | ||
| 5 | + _calculate_common_layout_flops, | ||
| 6 | + matmul_flops, | ||
| 7 | + npu_all_gather_base_mm_flops, | ||
| 8 | + npu_alltoallv_gmm_flops, | ||
| 9 | + npu_block_sparse_attention_flops, | ||
| 10 | + npu_fusion_attention_flops, | ||
| 11 | + npu_gmm_alltoallv_flops, | ||
| 12 | + npu_grouped_matmul_flops, | ||
| 13 | + npu_grouped_matmul_swiglu_quant_v2_flops, | ||
| 14 | + npu_quant_matmul_gelu_flops, | ||
| 15 | + npu_transpose_batchmatmul_flops, | ||
| 16 | +) | ||
| 17 | +from torch_npu.profiler._flops_hook import FlopsHookManager | ||
| 18 | +from torch_npu.testing.testcase import run_tests, TestCase | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +class _Tensor: | ||
| 22 | + def __init__(self, shape): | ||
| 23 | + self.shape = shape | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +class _TensorWithList: | ||
| 27 | + def __init__(self, value): | ||
| 28 | + self._value = value | ||
| 29 | + | ||
| 30 | + def tolist(self): | ||
| 31 | + return self._value | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +class TestFlopsHook(TestCase): | ||
| 35 | + def tearDown(self): | ||
| 36 | + FlopsHookManager.uninstall() | ||
| 37 | + | ||
| 38 | + def test_fusion_attention_formula_accepts_real_positional_arguments(self): | ||
| 39 | + query = _Tensor((2, 4, 8, 16)) | ||
| 40 | + key = _Tensor((2, 4, 8, 16)) | ||
| 41 | + value = _Tensor((2, 4, 8, 16)) | ||
| 42 | + self.assertEqual( | ||
| 43 | + 2 * 2 * 4 * 8 * 8 * (16 + 16), | ||
| 44 | + npu_fusion_attention_flops(query, key, value, 4, "BNSD"), | ||
| 45 | + ) | ||
| 46 | + | ||
| 47 | + def test_fusion_attention_formula_uses_value_dim(self): | ||
| 48 | + query = _Tensor((1, 2, 3, 4)) | ||
| 49 | + key = _Tensor((1, 2, 5, 4)) | ||
| 50 | + value = _Tensor((1, 2, 5, 7)) | ||
| 51 | + self.assertEqual( | ||
| 52 | + 2 * 1 * 2 * 3 * 5 * (4 + 7), | ||
| 53 | + npu_fusion_attention_flops(query, key, value, 2, "BNSD"), | ||
| 54 | + ) | ||
| 55 | + | ||
| 56 | + def test_sparse_formula_uses_sequence_lengths_for_non_square_attention(self): | ||
| 57 | + self.assertEqual( | ||
| 58 | + 2 * 1 * 2 * (8 * 4 - 4 * 4 / 2) * (16 + 16), | ||
| 59 | + _calculate_common_layout_flops( | ||
| 60 | + (1, 2, 8, 16), | ||
| 61 | + (1, 2, 4, 16), | ||
| 62 | + (1, 2, 4, 16), | ||
| 63 | + "BNSD", | ||
| 64 | + 2, | ||
| 65 | + 2, | ||
| 66 | + 2, | ||
| 67 | + ), | ||
| 68 | + ) | ||
| 69 | + | ||
| 70 | + def test_all_gather_base_mm_formula_counts_gathered_gemm_only(self): | ||
| 71 | + self.assertEqual( | ||
| 72 | + 2 * (3 * 2) * 4 * 5, | ||
| 73 | + npu_all_gather_base_mm_flops( | ||
| 74 | + _Tensor((3, 4)), _Tensor((4, 5)), "hcom", 2 | ||
| 75 | + ), | ||
| 76 | + ) | ||
| 77 | + | ||
| 78 | + def test_transpose_batchmatmul_formula_counts_batch_gemm_only(self): | ||
| 79 | + self.assertEqual( | ||
| 80 | + 2 * 2 * 3 * 4 * 5, | ||
| 81 | + npu_transpose_batchmatmul_flops( | ||
| 82 | + _Tensor((3, 2, 4)), | ||
| 83 | + _Tensor((2, 4, 5)), | ||
| 84 | + perm_x1=(1, 0, 2), | ||
| 85 | + ), | ||
| 86 | + ) | ||
| 87 | + | ||
| 88 | + def test_grouped_matmul_formula_sums_group_gemms(self): | ||
| 89 | + self.assertEqual( | ||
| 90 | + 2 * 2 * 3 * 7 + 2 * 4 * 5 * 11, | ||
| 91 | + npu_grouped_matmul_flops( | ||
| 92 | + [_Tensor((2, 3)), _Tensor((4, 5))], | ||
| 93 | + [_Tensor((3, 7)), _Tensor((5, 11))], | ||
| 94 | + ), | ||
| 95 | + ) | ||
| 96 | + | ||
| 97 | + def test_grouped_matmul_formula_uses_group_list_for_split_groups(self): | ||
| 98 | + self.assertEqual( | ||
| 99 | + 2 * 2 * 3 * 7 + 2 * 4 * 3 * 11, | ||
| 100 | + npu_grouped_matmul_flops( | ||
| 101 | + _Tensor((6, 3)), | ||
| 102 | + [_Tensor((3, 7)), _Tensor((3, 11))], | ||
| 103 | + group_list=[2, 6], | ||
| 104 | + ), | ||
| 105 | + ) | ||
| 106 | + | ||
| 107 | + def test_grouped_matmul_formula_rejects_scalar_group_list(self): | ||
| 108 | + with self.assertRaisesRegex(ValueError, "returning a sequence"): | ||
| 109 | + npu_grouped_matmul_flops( | ||
| 110 | + _Tensor((6, 3)), | ||
| 111 | + [_Tensor((3, 7)), _Tensor((3, 11))], | ||
| 112 | + group_list=_TensorWithList(6), | ||
| 113 | + ) | ||
| 114 | + | ||
| 115 | + def test_quant_matmul_gelu_formula_counts_matmul_only(self): | ||
| 116 | + self.assertEqual( | ||
| 117 | + 2 * 6 * 4 * 5, | ||
| 118 | + npu_quant_matmul_gelu_flops( | ||
| 119 | + _Tensor((2, 3, 4)), _Tensor((4, 5)), _Tensor((6,)), _Tensor((5,)) | ||
| 120 | + ), | ||
| 121 | + ) | ||
| 122 | + | ||
| 123 | + def test_grouped_matmul_swiglu_quant_formula_counts_grouped_gemm_only(self): | ||
| 124 | + self.assertEqual( | ||
| 125 | + 2 * 8 * 4 * 16, | ||
| 126 | + npu_grouped_matmul_swiglu_quant_v2_flops( | ||
| 127 | + _Tensor((8, 4)), | ||
| 128 | + _Tensor((2, 4, 16)), | ||
| 129 | + _Tensor((2, 16)), | ||
| 130 | + _Tensor((8,)), | ||
| 131 | + _Tensor((2,)), | ||
| 132 | + ), | ||
| 133 | + ) | ||
| 134 | + | ||
| 135 | + def test_alltoallv_gmm_formula_adds_optional_shared_mm(self): | ||
| 136 | + self.assertEqual( | ||
| 137 | + 2 * 7 * 4 * 6 + 2 * 5 * 2 * 9, | ||
| 138 | + npu_alltoallv_gmm_flops( | ||
| 139 | + _Tensor((7, 4)), | ||
| 140 | + _Tensor((3, 4, 6)), | ||
| 141 | + "hcom", | ||
| 142 | + 2, | ||
| 143 | + [3, 4], | ||
| 144 | + [4, 3], | ||
| 145 | + mm_x=_Tensor((5, 2)), | ||
| 146 | + mm_weight=_Tensor((2, 9)), | ||
| 147 | + ), | ||
| 148 | + ) | ||
| 149 | + | ||
| 150 | + def test_gmm_alltoallv_formula_counts_route_gmm_only_when_no_shared_mm(self): | ||
| 151 | + self.assertEqual( | ||
| 152 | + 2 * 7 * 4 * 6, | ||
| 153 | + npu_gmm_alltoallv_flops( | ||
| 154 | + _Tensor((7, 4)), | ||
| 155 | + _Tensor((3, 4, 6)), | ||
| 156 | + "hcom", | ||
| 157 | + 2, | ||
| 158 | + [3, 4], | ||
| 159 | + [4, 3], | ||
| 160 | + ), | ||
| 161 | + ) | ||
| 162 | + | ||
| 163 | + def test_block_sparse_attention_formula_counts_valid_blocks(self): | ||
| 164 | + mask = [ | ||
| 165 | + [ | ||
| 166 | + [[1, 0], [0, 1], [1, 1]], | ||
| 167 | + [[0, 1], [1, 0], [0, 0]], | ||
| 168 | + ] | ||
| 169 | + ] | ||
| 170 | + self.assertEqual( | ||
| 171 | + 2 * 30 * (4 + 8), | ||
| 172 | + npu_block_sparse_attention_flops( | ||
| 173 | + _Tensor((1, 2, 5, 4)), | ||
| 174 | + _Tensor((1, 2, 6, 4)), | ||
| 175 | + _Tensor((1, 2, 6, 8)), | ||
| 176 | + mask, | ||
| 177 | + (2, 3), | ||
| 178 | + q_input_layout="BNSD", | ||
| 179 | + kv_input_layout="BNSD", | ||
| 180 | + ), | ||
| 181 | + ) | ||
| 182 | + | ||
| 183 | + def test_block_sparse_attention_formula_keeps_bnsd_lengths_per_batch(self): | ||
| 184 | + mask = [ | ||
| 185 | + [[[1, 0], [0, 1], [1, 1]]], | ||
| 186 | + [[[1, 1], [0, 0], [1, 0]]], | ||
| 187 | + ] | ||
| 188 | + self.assertEqual( | ||
| 189 | + 2 * 21 * (4 + 4), | ||
| 190 | + npu_block_sparse_attention_flops( | ||
| 191 | + _Tensor((2, 1, 5, 4)), | ||
| 192 | + _Tensor((2, 1, 5, 4)), | ||
| 193 | + _Tensor((2, 1, 5, 4)), | ||
| 194 | + mask, | ||
| 195 | + (2, 3), | ||
| 196 | + q_input_layout="BNSD", | ||
| 197 | + kv_input_layout="BNSD", | ||
| 198 | + actual_seq_lengths=[3, 5], | ||
| 199 | + actual_seq_lengths_kv=[5, 5], | ||
| 200 | + ), | ||
| 201 | + ) | ||
| 202 | + | ||
| 203 | + def test_matmul_formula_supports_broadcast_batch_dimensions(self): | ||
| 204 | + self.assertEqual( | ||
| 205 | + 2 * 5 * 2 * 3 * 4 * 6, | ||
| 206 | + matmul_flops(_Tensor((5, 1, 3, 4)), _Tensor((2, 4, 6))), | ||
| 207 | + ) | ||
| 208 | + | ||
| 209 | + def test_hook_records_flops_and_op_name_in_range_label(self): | ||
| 210 | + target = SimpleNamespace() | ||
| 211 | + | ||
| 212 | + def original(value): | ||
| 213 | + return value + 1 | ||
| 214 | + | ||
| 215 | + target.op = original | ||
| 216 | + with ( | ||
| 217 | + mock.patch( | ||
| 218 | + "torch_npu.profiler._flops_hook.get_flop_func", | ||
| 219 | + return_value=lambda value: value * 2, | ||
| 220 | + ), | ||
| 221 | + mock.patch( | ||
| 222 | + "torch_npu.profiler._flops_hook.mstx.range_start", return_value=7 | ||
| 223 | + ) as mock_start, | ||
| 224 | + mock.patch("torch_npu.profiler._flops_hook.mstx.range_end") as mock_end, | ||
| 225 | + ): | ||
| 226 | + FlopsHookManager.install({"demo_op": (target, "op")}) | ||
| 227 | + self.assertEqual(4, target.op(3)) | ||
| 228 | + | ||
| 229 | + mock_start.assert_called_once_with("6-demo_op", domain="mfu_flops") | ||
| 230 | + mock_end.assert_called_once_with(7, domain="mfu_flops") | ||
| 231 | + | ||
| 232 | + | ||
| 233 | +if __name__ == "__main__": | ||
| 234 | + run_tests() | ||
| @@ -0,0 +1,69 @@ | |||
| 1 | +from unittest import mock | ||
| 2 | + | ||
| 3 | +from torch_npu.profiler._flops_registry import ( | ||
| 4 | + _default_npu_flop_registry, | ||
| 5 | + _npu_flop_registry, | ||
| 6 | + get_flop_func, | ||
| 7 | + get_npu_flop_targets, | ||
| 8 | + register_npu_flop, | ||
| 9 | +) | ||
| 10 | +from torch_npu.testing.testcase import run_tests, TestCase | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class TestFlopsRegistry(TestCase): | ||
| 14 | + def tearDown(self): | ||
| 15 | + for op_name in ( | ||
| 16 | + "test_default_override", | ||
| 17 | + "test_external_conflict", | ||
| 18 | + "test_external_formula_only", | ||
| 19 | + ): | ||
| 20 | + _default_npu_flop_registry.pop(op_name, None) | ||
| 21 | + _npu_flop_registry.pop(op_name, None) | ||
| 22 | + | ||
| 23 | + def test_external_registration_overrides_default_registration(self): | ||
| 24 | + | ||
| 25 | + def external_flops(): | ||
| 26 | + return 2 | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + target="torch:bmm", op_name="test_default_override", is_default=True | ||
| 30 | + ) | ||
| 31 | + def default_flops(): | ||
| 32 | + return 1 | ||
| 33 | + | ||
| 34 | + self.assertIs(external_flops, get_flop_func("test_default_override")) | ||
| 35 | + self.assertEqual("torch:mm", get_npu_flop_targets()["test_default_override"]) | ||
| 36 | + | ||
| 37 | + def test_duplicate_external_registration_logs_error_and_uses_later_one(self): | ||
| 38 | + | ||
| 39 | + def first_flops(): | ||
| 40 | + return 1 | ||
| 41 | + | ||
| 42 | + with mock.patch( | ||
| 43 | + "torch_npu.profiler._flops_registry.logger.error" | ||
| 44 | + ) as mock_error: | ||
| 45 | + | ||
| 46 | + | ||
| 47 | + def second_flops(): | ||
| 48 | + return 2 | ||
| 49 | + | ||
| 50 | + mock_error.assert_called_once() | ||
| 51 | + self.assertIs(second_flops, get_flop_func("test_external_conflict")) | ||
| 52 | + | ||
| 53 | + def test_external_formula_uses_default_target_when_omitted(self): | ||
| 54 | + | ||
| 55 | + target="torch:mm", op_name="test_external_formula_only", is_default=True | ||
| 56 | + ) | ||
| 57 | + def default_flops(): | ||
| 58 | + return 1 | ||
| 59 | + | ||
| 60 | + | ||
| 61 | + def external_flops(): | ||
| 62 | + return 2 | ||
| 63 | + | ||
| 64 | + self.assertIs(external_flops, get_flop_func("test_external_formula_only")) | ||
| 65 | + self.assertEqual("torch:mm", get_npu_flop_targets()["test_external_formula_only"]) | ||
| 66 | + | ||
| 67 | + | ||
| 68 | +if __name__ == "__main__": | ||
| 69 | + run_tests() | ||
| @@ -68,6 +68,33 @@ class TestActionController(TestCase): | |||
| 68 | self.prof_if.stop_trace() | 68 | self.prof_if.stop_trace() |
| 69 | mock_stop.assert_called_once() | 69 | mock_stop.assert_called_once() |
| 70 | 70 | ||
| 71 | + def test_start_and_stop_trace_with_flops_and_msprof_tx(self): | ||
| 72 | + self.prof_if.with_flops = True | ||
| 73 | + self.prof_if.experimental_config._msprof_tx = True | ||
| 74 | + with ( | ||
| 75 | + mock.patch(self.namespace + "._start_profiler"), | ||
| 76 | + mock.patch(self.namespace + "._stop_profiler"), | ||
| 77 | + mock.patch(self.namespace + ".FlopsHookManager.install") as mock_install, | ||
| 78 | + mock.patch(self.namespace + ".FlopsHookManager.uninstall") as mock_uninstall, | ||
| 79 | + ): | ||
| 80 | + self.prof_if.start_trace() | ||
| 81 | + self.prof_if.stop_trace() | ||
| 82 | + mock_install.assert_called_once() | ||
| 83 | + mock_uninstall.assert_called_once() | ||
| 84 | + | ||
| 85 | + def test_start_and_stop_trace_with_flops_without_msprof_tx_skips_flops_hook(self): | ||
| 86 | + self.prof_if.with_flops = True | ||
| 87 | + with ( | ||
| 88 | + mock.patch(self.namespace + "._start_profiler"), | ||
| 89 | + mock.patch(self.namespace + "._stop_profiler"), | ||
| 90 | + mock.patch(self.namespace + ".FlopsHookManager.install") as mock_install, | ||
| 91 | + mock.patch(self.namespace + ".FlopsHookManager.uninstall") as mock_uninstall, | ||
| 92 | + ): | ||
| 93 | + self.prof_if.start_trace() | ||
| 94 | + self.prof_if.stop_trace() | ||
| 95 | + mock_install.assert_not_called() | ||
| 96 | + mock_uninstall.assert_not_called() | ||
| 97 | + | ||
| 71 | def test_finalize_trace(self): | 98 | def test_finalize_trace(self): |
| 72 | with ( | 99 | with ( |
| 73 | mock.patch(self.namespace + "._init_profiler"), | 100 | mock.patch(self.namespace + "._init_profiler"), |
| @@ -0,0 +1,526 @@ | |||
| 1 | +from itertools import pairwise | ||
| 2 | +from math import prod | ||
| 3 | + | ||
| 4 | +from ._flops_registry import register_npu_flop | ||
| 5 | + | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def npu_fusion_attention_flops( | ||
| 9 | + query, | ||
| 10 | + key, | ||
| 11 | + value, | ||
| 12 | + head_num, | ||
| 13 | + input_layout, | ||
| 14 | + pse=None, | ||
| 15 | + padding_mask=None, | ||
| 16 | + atten_mask=None, | ||
| 17 | + scale=1.0, | ||
| 18 | + keep_prob=1.0, | ||
| 19 | + pre_tockens=2147483647, | ||
| 20 | + next_tockens=2147483647, | ||
| 21 | + inner_precise=0, | ||
| 22 | + prefix=None, | ||
| 23 | + actual_seq_qlen=None, | ||
| 24 | + actual_seq_kvlen=None, | ||
| 25 | + sparse_mode=0, | ||
| 26 | + *args, | ||
| 27 | + **kwargs, | ||
| 28 | +): | ||
| 29 | + q_shape = query.shape | ||
| 30 | + k_shape = key.shape | ||
| 31 | + v_shape = value.shape | ||
| 32 | + if input_layout == "TND": | ||
| 33 | + return _calculate_tnd_layout_flops( | ||
| 34 | + q_shape, k_shape, v_shape, actual_seq_qlen, actual_seq_kvlen | ||
| 35 | + ) | ||
| 36 | + return _calculate_common_layout_flops( | ||
| 37 | + q_shape, k_shape, v_shape, input_layout, sparse_mode, head_num, head_num | ||
| 38 | + ) | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +def npu_fused_infer_attention_score_flops( | ||
| 43 | + query, | ||
| 44 | + key, | ||
| 45 | + value, | ||
| 46 | + *, | ||
| 47 | + input_layout, | ||
| 48 | + num_heads, | ||
| 49 | + num_key_value_heads=0, | ||
| 50 | + actual_seq_lengths=None, | ||
| 51 | + actual_seq_lengths_kv=None, | ||
| 52 | + sparse_mode=0, | ||
| 53 | + **kwargs, | ||
| 54 | +): | ||
| 55 | + num_key_value_heads = num_key_value_heads or num_heads | ||
| 56 | + q_shape = query.shape | ||
| 57 | + k_shape = key.shape | ||
| 58 | + v_shape = value.shape | ||
| 59 | + if input_layout == "TND": | ||
| 60 | + return _calculate_tnd_layout_flops( | ||
| 61 | + q_shape, | ||
| 62 | + k_shape, | ||
| 63 | + v_shape, | ||
| 64 | + actual_seq_lengths, | ||
| 65 | + actual_seq_lengths_kv, | ||
| 66 | + num_heads, | ||
| 67 | + ) | ||
| 68 | + return _calculate_common_layout_flops( | ||
| 69 | + q_shape, | ||
| 70 | + k_shape, | ||
| 71 | + v_shape, | ||
| 72 | + input_layout, | ||
| 73 | + sparse_mode, | ||
| 74 | + num_heads, | ||
| 75 | + num_key_value_heads, | ||
| 76 | + ) | ||
| 77 | + | ||
| 78 | + | ||
| 79 | + | ||
| 80 | +def npu_all_gather_base_mm_flops( | ||
| 81 | + x1, | ||
| 82 | + x2, | ||
| 83 | + hcom, | ||
| 84 | + world_size, | ||
| 85 | + bias=None, | ||
| 86 | + x1_scale=None, | ||
| 87 | + x2_scale=None, | ||
| 88 | + gather_index=0, | ||
| 89 | + gather_output=True, | ||
| 90 | + comm_turn=0, | ||
| 91 | + output_dtype=None, | ||
| 92 | + comm_mode=None, | ||
| 93 | + **kwargs, | ||
| 94 | +): | ||
| 95 | + x1_shape = _shape(x1) | ||
| 96 | + x2_shape = _shape(x2) | ||
| 97 | + m_local, k = x1_shape[-2:] | ||
| 98 | + n = x2_shape[-1] | ||
| 99 | + return 2 * m_local * int(world_size) * k * n | ||
| 100 | + | ||
| 101 | + | ||
| 102 | + | ||
| 103 | +def npu_transpose_batchmatmul_flops( | ||
| 104 | + input, | ||
| 105 | + weight, | ||
| 106 | + *, | ||
| 107 | + bias=None, | ||
| 108 | + scale=None, | ||
| 109 | + perm_x1=(0, 1, 2), | ||
| 110 | + perm_x2=(0, 1, 2), | ||
| 111 | + perm_y=(1, 0, 2), | ||
| 112 | + batch_split_factor=1, | ||
| 113 | + **kwargs, | ||
| 114 | +): | ||
| 115 | + input_shape = _permute_shape(_shape(input), perm_x1) | ||
| 116 | + weight_shape = _permute_shape(_shape(weight), perm_x2) | ||
| 117 | + return _matmul_shape_flops(input_shape, weight_shape) | ||
| 118 | + | ||
| 119 | + | ||
| 120 | + | ||
| 121 | +def npu_grouped_matmul_flops( | ||
| 122 | + x, | ||
| 123 | + weight, | ||
| 124 | + *, | ||
| 125 | + bias=None, | ||
| 126 | + scale=None, | ||
| 127 | + offset=None, | ||
| 128 | + antiquant_scale=None, | ||
| 129 | + antiquant_offset=None, | ||
| 130 | + per_token_scale=None, | ||
| 131 | + group_list=None, | ||
| 132 | + activation_input=None, | ||
| 133 | + activation_quant_scale=None, | ||
| 134 | + activation_quant_offset=None, | ||
| 135 | + split_item=0, | ||
| 136 | + group_type=None, | ||
| 137 | + group_list_type=0, | ||
| 138 | + act_type=0, | ||
| 139 | + output_dtype=None, | ||
| 140 | + tuning_config=None, | ||
| 141 | + **kwargs, | ||
| 142 | +): | ||
| 143 | + return _grouped_matmul_flops(x, weight, group_list) | ||
| 144 | + | ||
| 145 | + | ||
| 146 | + | ||
| 147 | +def npu_quant_matmul_gelu_flops( | ||
| 148 | + x1, | ||
| 149 | + x2, | ||
| 150 | + x1_scale, | ||
| 151 | + x2_scale, | ||
| 152 | + *, | ||
| 153 | + bias=None, | ||
| 154 | + approximate="gelu_erf", | ||
| 155 | + **kwargs, | ||
| 156 | +): | ||
| 157 | + return _matrix_tensor_flops(x1, x2) | ||
| 158 | + | ||
| 159 | + | ||
| 160 | + | ||
| 161 | + target="torch_npu:npu_grouped_matmul_swiglu_quant_v2", is_default=True | ||
| 162 | +) | ||
| 163 | +def npu_grouped_matmul_swiglu_quant_v2_flops( | ||
| 164 | + x, | ||
| 165 | + weight, | ||
| 166 | + weight_scale, | ||
| 167 | + x_scale, | ||
| 168 | + group_list, | ||
| 169 | + *, | ||
| 170 | + smooth_scale=None, | ||
| 171 | + weight_assist_matrix=None, | ||
| 172 | + bias=None, | ||
| 173 | + dequant_mode=0, | ||
| 174 | + dequant_dtype=0, | ||
| 175 | + quant_mode=0, | ||
| 176 | + quant_dtype=0, | ||
| 177 | + group_list_type=0, | ||
| 178 | + tuning_config=None, | ||
| 179 | + **kwargs, | ||
| 180 | +): | ||
| 181 | + return _matrix_tensor_flops(x, weight) | ||
| 182 | + | ||
| 183 | + | ||
| 184 | + | ||
| 185 | +def npu_alltoallv_gmm_flops( | ||
| 186 | + gmm_x, | ||
| 187 | + gmm_weight, | ||
| 188 | + hcom, | ||
| 189 | + ep_world_size, | ||
| 190 | + send_counts, | ||
| 191 | + recv_counts, | ||
| 192 | + *, | ||
| 193 | + send_counts_tensor=None, | ||
| 194 | + recv_counts_tensor=None, | ||
| 195 | + mm_x=None, | ||
| 196 | + mm_weight=None, | ||
| 197 | + trans_gmm_weight=False, | ||
| 198 | + trans_mm_weight=False, | ||
| 199 | + permute_out_flag=False, | ||
| 200 | + **kwargs, | ||
| 201 | +): | ||
| 202 | + return _gmm_with_optional_mm_flops( | ||
| 203 | + gmm_x, gmm_weight, mm_x, mm_weight, trans_gmm_weight, trans_mm_weight | ||
| 204 | + ) | ||
| 205 | + | ||
| 206 | + | ||
| 207 | + | ||
| 208 | +def npu_gmm_alltoallv_flops( | ||
| 209 | + gmm_x, | ||
| 210 | + gmm_weight, | ||
| 211 | + hcom, | ||
| 212 | + ep_world_size, | ||
| 213 | + send_counts, | ||
| 214 | + recv_counts, | ||
| 215 | + *, | ||
| 216 | + send_counts_tensor=None, | ||
| 217 | + recv_counts_tensor=None, | ||
| 218 | + mm_x=None, | ||
| 219 | + mm_weight=None, | ||
| 220 | + trans_gmm_weight=False, | ||
| 221 | + trans_mm_weight=False, | ||
| 222 | + **kwargs, | ||
| 223 | +): | ||
| 224 | + return _gmm_with_optional_mm_flops( | ||
| 225 | + gmm_x, gmm_weight, mm_x, mm_weight, trans_gmm_weight, trans_mm_weight | ||
| 226 | + ) | ||
| 227 | + | ||
| 228 | + | ||
| 229 | + | ||
| 230 | +def npu_block_sparse_attention_flops( | ||
| 231 | + query, | ||
| 232 | + key, | ||
| 233 | + value, | ||
| 234 | + block_sparse_mask, | ||
| 235 | + block_shape, | ||
| 236 | + *, | ||
| 237 | + q_input_layout="TND", | ||
| 238 | + kv_input_layout="TND", | ||
| 239 | + num_key_value_heads=1, | ||
| 240 | + scale_value=0.0, | ||
| 241 | + inner_precise=1, | ||
| 242 | + actual_seq_lengths=None, | ||
| 243 | + actual_seq_lengths_kv=None, | ||
| 244 | + softmax_lse_flag=0, | ||
| 245 | + **kwargs, | ||
| 246 | +): | ||
| 247 | + q_shape = _shape(query) | ||
| 248 | + v_shape = _shape(value) | ||
| 249 | + mask = _to_nested_list(block_sparse_mask) | ||
| 250 | + q_heads = len(mask[0]) if mask else None | ||
| 251 | + _, _, q_s, q_d = _parse_attention_dims(q_shape, q_input_layout, q_heads) | ||
| 252 | + _, _, kv_s, v_d = _parse_attention_dims( | ||
| 253 | + v_shape, kv_input_layout, num_key_value_heads | ||
| 254 | + ) | ||
| 255 | + batch = len(mask) | ||
| 256 | + q_lens = _parse_actual_lengths( | ||
| 257 | + actual_seq_lengths, batch, q_s, q_input_layout == "TND" | ||
| 258 | + ) | ||
| 259 | + kv_lens = _parse_actual_lengths( | ||
| 260 | + actual_seq_lengths_kv, batch, kv_s, kv_input_layout == "TND" | ||
| 261 | + ) | ||
| 262 | + block_x, block_y = [int(dim) for dim in block_shape] | ||
| 263 | + score_elems = _count_block_sparse_score_elems( | ||
| 264 | + mask, q_lens, kv_lens, block_x, block_y | ||
| 265 | + ) | ||
| 266 | + return int(2 * score_elems * (q_d + v_d)) | ||
| 267 | + | ||
| 268 | + | ||
| 269 | + | ||
| 270 | +def mm_flops(input, other, **kwargs): | ||
| 271 | + m, k = input.shape | ||
| 272 | + _, n = other.shape | ||
| 273 | + return 2 * m * n * k | ||
| 274 | + | ||
| 275 | + | ||
| 276 | + | ||
| 277 | +def bmm_flops(input, other, **kwargs): | ||
| 278 | + b, m, k = input.shape | ||
| 279 | + _, _, n = other.shape | ||
| 280 | + return 2 * b * m * n * k | ||
| 281 | + | ||
| 282 | + | ||
| 283 | + | ||
| 284 | +def matmul_flops(input, other, **kwargs): | ||
| 285 | + input_shape = tuple(input.shape) | ||
| 286 | + other_shape = tuple(other.shape) | ||
| 287 | + if len(input_shape) == 1 and len(other_shape) == 1: | ||
| 288 | + return 2 * input_shape[0] | ||
| 289 | + if len(input_shape) == 1: | ||
| 290 | + batch_shape = other_shape[:-2] | ||
| 291 | + m, k, n = 1, input_shape[0], other_shape[-1] | ||
| 292 | + elif len(other_shape) == 1: | ||
| 293 | + batch_shape = input_shape[:-2] | ||
| 294 | + m, k, n = input_shape[-2], input_shape[-1], 1 | ||
| 295 | + else: | ||
| 296 | + batch_shape = _broadcast_shapes(input_shape[:-2], other_shape[:-2]) | ||
| 297 | + m, k, n = input_shape[-2], input_shape[-1], other_shape[-1] | ||
| 298 | + return 2 * prod(batch_shape) * m * n * k | ||
| 299 | + | ||
| 300 | + | ||
| 301 | + | ||
| 302 | +def linear_flops(input, weight, bias=None, **kwargs): | ||
| 303 | + n, k = weight.shape | ||
| 304 | + return 2 * prod(input.shape[:-1]) * n * k | ||
| 305 | + | ||
| 306 | + | ||
| 307 | + | ||
| 308 | +def addmm_flops(self, mat1, mat2, beta=1, alpha=1, **kwargs): | ||
| 309 | + m, k = mat1.shape | ||
| 310 | + _, n = mat2.shape | ||
| 311 | + return 2 * m * n * k | ||
| 312 | + | ||
| 313 | + | ||
| 314 | +def _shape(tensor): | ||
| 315 | + return tuple(int(dim) for dim in tensor.shape) | ||
| 316 | + | ||
| 317 | + | ||
| 318 | +def _matmul_shape_flops(left_shape, right_shape, trans_right=False): | ||
| 319 | + if len(left_shape) < 2 or len(right_shape) < 2: | ||
| 320 | + raise ValueError(f"Matmul FLOPs requires rank >= 2: {left_shape}, {right_shape}") | ||
| 321 | + m = prod(left_shape[:-1]) | ||
| 322 | + k = left_shape[-1] | ||
| 323 | + n = right_shape[-2] if trans_right else right_shape[-1] | ||
| 324 | + return int(2 * m * k * n) | ||
| 325 | + | ||
| 326 | + | ||
| 327 | +def _matrix_tensor_flops(left, right, trans_right=False): | ||
| 328 | + return _matmul_shape_flops(_shape(left), _shape(right), trans_right) | ||
| 329 | + | ||
| 330 | + | ||
| 331 | +def _as_tensor_list(tensors): | ||
| 332 | + return list(tensors) if isinstance(tensors, (list, tuple)) else [tensors] | ||
| 333 | + | ||
| 334 | + | ||
| 335 | +def _grouped_matmul_flops(x, weight, group_list=None): | ||
| 336 | + x_list = _as_tensor_list(x) | ||
| 337 | + weight_list = _as_tensor_list(weight) | ||
| 338 | + if len(x_list) == len(weight_list): | ||
| 339 | + return sum( | ||
| 340 | + _matrix_tensor_flops(left, right) | ||
| 341 | + for left, right in zip(x_list, weight_list) | ||
| 342 | + ) | ||
| 343 | + if len(x_list) == 1: | ||
| 344 | + left_shape = _shape(x_list[0]) | ||
| 345 | + group_lengths = _parse_group_lengths( | ||
| 346 | + group_list, len(weight_list), prod(left_shape[:-1]) | ||
| 347 | + ) | ||
| 348 | + return sum( | ||
| 349 | + _matmul_shape_flops((group_m, left_shape[-1]), _shape(right)) | ||
| 350 | + for group_m, right in zip(group_lengths, weight_list) | ||
| 351 | + ) | ||
| 352 | + raise ValueError( | ||
| 353 | + f"Grouped matmul FLOPs requires matching groups: {len(x_list)}, {len(weight_list)}" | ||
| 354 | + ) | ||
| 355 | + | ||
| 356 | + | ||
| 357 | +def _parse_group_lengths(group_list, group_count, total_m): | ||
| 358 | + if group_count == 1: | ||
| 359 | + return [total_m] | ||
| 360 | + if group_list is None: | ||
| 361 | + raise ValueError("Grouped matmul FLOPs requires group_list for split groups") | ||
| 362 | + groups = [int(group) for group in _to_sequence(group_list)] | ||
| 363 | + if len(groups) != group_count: | ||
| 364 | + raise ValueError(f"Expected {group_count} groups, got {len(groups)}") | ||
| 365 | + if groups[-1] == total_m and sum(groups) > total_m: | ||
| 366 | + groups = [groups[0]] + [curr - prev for prev, curr in pairwise(groups)] | ||
| 367 | + if sum(groups) != total_m: | ||
| 368 | + raise ValueError("group_list does not match grouped matmul token count") | ||
| 369 | + return groups | ||
| 370 | + | ||
| 371 | + | ||
| 372 | +def _gmm_with_optional_mm_flops( | ||
| 373 | + gmm_x, gmm_weight, mm_x, mm_weight, trans_gmm_weight, trans_mm_weight | ||
| 374 | +): | ||
| 375 | + flops = _matrix_tensor_flops(gmm_x, gmm_weight, trans_gmm_weight) | ||
| 376 | + if mm_x is not None and mm_weight is not None: | ||
| 377 | + flops += _matrix_tensor_flops(mm_x, mm_weight, trans_mm_weight) | ||
| 378 | + return flops | ||
| 379 | + | ||
| 380 | + | ||
| 381 | +def _permute_shape(tensor_shape, permutation): | ||
| 382 | + if len(tensor_shape) != len(permutation): | ||
| 383 | + raise ValueError( | ||
| 384 | + f"Permutation {permutation} does not match tensor shape {tensor_shape}" | ||
| 385 | + ) | ||
| 386 | + return tuple(tensor_shape[int(index)] for index in permutation) | ||
| 387 | + | ||
| 388 | + | ||
| 389 | +def _broadcast_shapes(left_shape, right_shape): | ||
| 390 | + result = [] | ||
| 391 | + for left, right in zip(reversed(left_shape), reversed(right_shape)): | ||
| 392 | + if left != right and left != 1 and right != 1: | ||
| 393 | + raise ValueError( | ||
| 394 | + f"Cannot broadcast matmul batch dimensions: {left_shape}, {right_shape}" | ||
| 395 | + ) | ||
| 396 | + result.append(max(left, right)) | ||
| 397 | + longer = left_shape if len(left_shape) > len(right_shape) else right_shape | ||
| 398 | + result.extend(reversed(longer[: abs(len(left_shape) - len(right_shape))])) | ||
| 399 | + return tuple(reversed(result)) | ||
| 400 | + | ||
| 401 | + | ||
| 402 | +def _calculate_common_layout_flops( | ||
| 403 | + q_shape, k_shape, v_shape, input_layout, sparse_mode, q_heads, kv_heads | ||
| 404 | +): | ||
| 405 | + q_b, q_n, q_s, q_d = _parse_dims(q_shape, input_layout, q_heads) | ||
| 406 | + _, _, k_s, k_d = _parse_dims(k_shape, input_layout, kv_heads) | ||
| 407 | + _, _, _, v_d = _parse_dims(v_shape, input_layout, kv_heads) | ||
| 408 | + attention_scores = _calculate_attention_scores(q_s, k_s, sparse_mode) | ||
| 409 | + return int(2 * q_b * q_n * attention_scores * (q_d + v_d)) | ||
| 410 | + | ||
| 411 | + | ||
| 412 | +def _calculate_tnd_layout_flops( | ||
| 413 | + q_shape, k_shape, v_shape, actual_seq_qlen, actual_seq_kvlen, q_heads=None | ||
| 414 | +): | ||
| 415 | + if actual_seq_qlen is None or actual_seq_kvlen is None: | ||
| 416 | + raise ValueError("TND layout requires actual_seq_qlen and actual_seq_kvlen") | ||
| 417 | + _, shape_q_heads, q_d = q_shape | ||
| 418 | + _, _, v_d = v_shape | ||
| 419 | + q_lens = _parse_seq_len(actual_seq_qlen) | ||
| 420 | + kv_lens = _parse_seq_len(actual_seq_kvlen) | ||
| 421 | + if len(q_lens) != len(kv_lens) or any(length <= 0 for length in q_lens + kv_lens): | ||
| 422 | + raise ValueError("actual_seq_qlen and actual_seq_kvlen must contain valid cumulative lengths") | ||
| 423 | + attention_scores = sum(q_len * kv_len for q_len, kv_len in zip(q_lens, kv_lens)) | ||
| 424 | + return int(2 * (q_heads or shape_q_heads) * (q_d + v_d) * attention_scores) | ||
| 425 | + | ||
| 426 | + | ||
| 427 | +def _calculate_attention_scores(q_s, k_s, sparse_mode): | ||
| 428 | + if sparse_mode == 0: | ||
| 429 | + return q_s * k_s | ||
| 430 | + if sparse_mode not in (2, 3): | ||
| 431 | + raise ValueError(f"Unknown FLOPs formula for sparse_mode={sparse_mode}") | ||
| 432 | + if sparse_mode == 2: | ||
| 433 | + return q_s * k_s - k_s * k_s / 2 if q_s >= k_s else q_s * q_s / 2 | ||
| 434 | + return k_s * k_s / 2 if q_s >= k_s else q_s * k_s - q_s * q_s / 2 | ||
| 435 | + | ||
| 436 | + | ||
| 437 | +def _parse_dims(tensor_shape, input_layout, heads): | ||
| 438 | + if input_layout == "BNSD": | ||
| 439 | + return tensor_shape | ||
| 440 | + if input_layout == "BSND": | ||
| 441 | + b, s, n, d = tensor_shape | ||
| 442 | + return b, n, s, d | ||
| 443 | + if input_layout == "BSH": | ||
| 444 | + b, s, h = tensor_shape | ||
| 445 | + return b, heads, s, _head_dim(h, heads) | ||
| 446 | + if input_layout == "SBH": | ||
| 447 | + s, b, h = tensor_shape | ||
| 448 | + return b, heads, s, _head_dim(h, heads) | ||
| 449 | + raise ValueError(f"Invalid layout for FlashAttention input tensor: {input_layout}") | ||
| 450 | + | ||
| 451 | + | ||
| 452 | +def _parse_attention_dims(tensor_shape, input_layout, heads): | ||
| 453 | + if input_layout == "TND": | ||
| 454 | + s, n, d = tensor_shape | ||
| 455 | + return None, n, s, d | ||
| 456 | + return _parse_dims(tensor_shape, input_layout, heads) | ||
| 457 | + | ||
| 458 | + | ||
| 459 | +def _head_dim(hidden_size, heads): | ||
| 460 | + if heads <= 0 or hidden_size % heads != 0: | ||
| 461 | + raise ValueError( | ||
| 462 | + f"Hidden size {hidden_size} must be divisible by the number of heads {heads}" | ||
| 463 | + ) | ||
| 464 | + return hidden_size // heads | ||
| 465 | + | ||
| 466 | + | ||
| 467 | +def _parse_seq_len(original_seq_lens): | ||
| 468 | + seq_lens = [int(length) for length in original_seq_lens] | ||
| 469 | + while seq_lens and seq_lens[-1] == 0: | ||
| 470 | + seq_lens.pop() | ||
| 471 | + if not seq_lens: | ||
| 472 | + return [] | ||
| 473 | + return [seq_lens[0]] + [ | ||
| 474 | + curr - prev for prev, curr in pairwise(seq_lens) | ||
| 475 | + ] | ||
| 476 | + | ||
| 477 | + | ||
| 478 | +def _parse_actual_lengths(seq_lens, batch, default_len, is_cumulative=False): | ||
| 479 | + if seq_lens is None: | ||
| 480 | + return [default_len] * batch | ||
| 481 | + lengths = [int(length) for length in seq_lens] | ||
| 482 | + while lengths and lengths[-1] == 0: | ||
| 483 | + lengths.pop() | ||
| 484 | + if len(lengths) != batch: | ||
| 485 | + raise ValueError(f"Expected {batch} sequence lengths, got {len(lengths)}") | ||
| 486 | + if is_cumulative: | ||
| 487 | + lengths = [lengths[0]] + [ | ||
| 488 | + curr - prev for prev, curr in pairwise(lengths) | ||
| 489 | + ] | ||
| 490 | + if any(length < 0 for length in lengths): | ||
| 491 | + raise ValueError("Sequence lengths must be non-negative") | ||
| 492 | + return lengths | ||
| 493 | + | ||
| 494 | + | ||
| 495 | +def _to_nested_list(value): | ||
| 496 | + return _to_sequence(value) | ||
| 497 | + | ||
| 498 | + | ||
| 499 | +def _to_sequence(value): | ||
| 500 | + if isinstance(value, (list, tuple)): | ||
| 501 | + return value | ||
| 502 | + if hasattr(value, "tolist"): | ||
| 503 | + sequence = value.tolist() | ||
| 504 | + if isinstance(sequence, (list, tuple)): | ||
| 505 | + return sequence | ||
| 506 | + raise ValueError("Value must be a sequence or expose tolist() returning a sequence") | ||
| 507 | + | ||
| 508 | + | ||
| 509 | +def _count_block_sparse_score_elems(mask, q_lens, kv_lens, block_x, block_y): | ||
| 510 | + score_elems = 0 | ||
| 511 | + for batch_idx, heads in enumerate(mask): | ||
| 512 | + q_len = q_lens[batch_idx] | ||
| 513 | + kv_len = kv_lens[batch_idx] | ||
| 514 | + for q_blocks in heads: | ||
| 515 | + for q_block_idx, kv_blocks in enumerate(q_blocks): | ||
| 516 | + q_start = q_block_idx * block_x | ||
| 517 | + q_tokens = min(block_x, max(q_len - q_start, 0)) | ||
| 518 | + if q_tokens == 0: | ||
| 519 | + continue | ||
| 520 | + for kv_block_idx, is_valid in enumerate(kv_blocks): | ||
| 521 | + if not is_valid: | ||
| 522 | + continue | ||
| 523 | + kv_start = kv_block_idx * block_y | ||
| 524 | + kv_tokens = min(block_y, max(kv_len - kv_start, 0)) | ||
| 525 | + score_elems += q_tokens * kv_tokens | ||
| 526 | + return score_elems | ||
| @@ -0,0 +1,153 @@ | |||
| 1 | +import functools | ||
| 2 | +import importlib | ||
| 3 | +import logging | ||
| 4 | +import sys | ||
| 5 | +import threading | ||
| 6 | +from collections.abc import Callable | ||
| 7 | +from typing import Any | ||
| 8 | + | ||
| 9 | +from torch_npu.npu.mstx import mstx | ||
| 10 | + | ||
| 11 | +from ._flops_registry import get_flop_func, get_npu_flop_targets | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +logger = logging.getLogger(__name__) | ||
| 15 | + | ||
| 16 | +_FLOPS_DOMAIN = "mfu_flops" | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def _resolve_target(target: str): | ||
| 20 | + if ":" not in target: | ||
| 21 | + logger.warning("Invalid FLOPs target format: %s", target) | ||
| 22 | + return None, None | ||
| 23 | + module_path, attr_path = target.split(":", 1) | ||
| 24 | + try: | ||
| 25 | + module_obj = importlib.import_module(module_path) | ||
| 26 | + except ImportError: | ||
| 27 | + logger.warning("Cannot import FLOPs target module: %s", module_path) | ||
| 28 | + return None, None | ||
| 29 | + | ||
| 30 | + obj = module_obj | ||
| 31 | + attrs = attr_path.split(".") | ||
| 32 | + for attr in attrs[:-1]: | ||
| 33 | + if not hasattr(obj, attr): | ||
| 34 | + logger.warning("Cannot resolve FLOPs target: %s", target) | ||
| 35 | + return None, None | ||
| 36 | + obj = getattr(obj, attr) | ||
| 37 | + return obj, attrs[-1] | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +def _build_target_ops() -> dict[str, tuple[Any, str]]: | ||
| 41 | + importlib.import_module("torch_npu.profiler._flops_formulas") | ||
| 42 | + | ||
| 43 | + target_ops = {} | ||
| 44 | + for op_name, target in get_npu_flop_targets().items(): | ||
| 45 | + parent, attr_name = _resolve_target(target) | ||
| 46 | + if parent is not None and attr_name is not None: | ||
| 47 | + target_ops[op_name] = (parent, attr_name) | ||
| 48 | + return target_ops | ||
| 49 | + | ||
| 50 | + | ||
| 51 | +def _find_existing_refs(original_func: Callable, attr_name: str): | ||
| 52 | + refs = [] | ||
| 53 | + for module in list(sys.modules.values()): | ||
| 54 | + if module is None: | ||
| 55 | + continue | ||
| 56 | + try: | ||
| 57 | + ref = getattr(module, attr_name, None) | ||
| 58 | + except Exception: | ||
| 59 | + continue | ||
| 60 | + if ref is original_func: | ||
| 61 | + refs.append((module, attr_name)) | ||
| 62 | + return refs | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +class FlopsHookManager: | ||
| 66 | + _local = threading.local() | ||
| 67 | + _original_funcs: dict[str, Callable] = {} | ||
| 68 | + _patched_targets: dict[str, tuple[Any, str]] = {} | ||
| 69 | + _extra_refs: dict[str, list[tuple[Any, str]]] = {} | ||
| 70 | + _installed = False | ||
| 71 | + | ||
| 72 | + | ||
| 73 | + def install(cls, target_ops: dict[str, tuple[Any, str]] | None = None): | ||
| 74 | + if cls._installed: | ||
| 75 | + return | ||
| 76 | + if target_ops is None: | ||
| 77 | + target_ops = _build_target_ops() | ||
| 78 | + | ||
| 79 | + try: | ||
| 80 | + for op_name, (module_obj, attr_name) in target_ops.items(): | ||
| 81 | + original = getattr(module_obj, attr_name, None) | ||
| 82 | + if original is None: | ||
| 83 | + logger.warning("Cannot find FLOPs target %s.%s", module_obj, attr_name) | ||
| 84 | + continue | ||
| 85 | + wrapped = cls._make_wrapper(op_name, original) | ||
| 86 | + setattr(module_obj, attr_name, wrapped) | ||
| 87 | + cls._original_funcs[op_name] = original | ||
| 88 | + cls._patched_targets[op_name] = (module_obj, attr_name) | ||
| 89 | + cls._extra_refs[op_name] = [] | ||
| 90 | + for ref_module, ref_attr in _find_existing_refs(original, attr_name): | ||
| 91 | + if ref_module is module_obj and ref_attr == attr_name: | ||
| 92 | + continue | ||
| 93 | + setattr(ref_module, ref_attr, wrapped) | ||
| 94 | + cls._extra_refs[op_name].append((ref_module, ref_attr)) | ||
| 95 | + except Exception: | ||
| 96 | + cls._restore() | ||
| 97 | + raise | ||
| 98 | + | ||
| 99 | + cls._installed = bool(cls._patched_targets) | ||
| 100 | + | ||
| 101 | + | ||
| 102 | + def uninstall(cls): | ||
| 103 | + cls._restore() | ||
| 104 | + | ||
| 105 | + | ||
| 106 | + def _restore(cls): | ||
| 107 | + for op_name, (module_obj, attr_name) in cls._patched_targets.items(): | ||
| 108 | + original = cls._original_funcs.get(op_name) | ||
| 109 | + if original is not None: | ||
| 110 | + setattr(module_obj, attr_name, original) | ||
| 111 | + for op_name, refs in cls._extra_refs.items(): | ||
| 112 | + original = cls._original_funcs.get(op_name) | ||
| 113 | + if original is not None: | ||
| 114 | + for ref_module, ref_attr in refs: | ||
| 115 | + setattr(ref_module, ref_attr, original) | ||
| 116 | + cls._original_funcs.clear() | ||
| 117 | + cls._patched_targets.clear() | ||
| 118 | + cls._extra_refs.clear() | ||
| 119 | + cls._installed = False | ||
| 120 | + | ||
| 121 | + | ||
| 122 | + def is_installed(cls) -> bool: | ||
| 123 | + return cls._installed | ||
| 124 | + | ||
| 125 | + | ||
| 126 | + def _make_wrapper(cls, op_name: str, original_func: Callable) -> Callable: | ||
| 127 | + | ||
| 128 | + def wrapper(*args, **kwargs): | ||
| 129 | + if getattr(cls._local, "in_hook", False): | ||
| 130 | + return original_func(*args, **kwargs) | ||
| 131 | + | ||
| 132 | + cls._local.in_hook = True | ||
| 133 | + range_id = None | ||
| 134 | + try: | ||
| 135 | + flop_func = get_flop_func(op_name) | ||
| 136 | + if flop_func is not None: | ||
| 137 | + try: | ||
| 138 | + flops = flop_func(*args, **kwargs) | ||
| 139 | + if flops is not None and flops >= 0: | ||
| 140 | + range_id = mstx.range_start( | ||
| 141 | + f"{flops}-{op_name}", domain=_FLOPS_DOMAIN | ||
| 142 | + ) | ||
| 143 | + except Exception: | ||
| 144 | + logger.warning( | ||
| 145 | + "Failed to calculate FLOPs for %s", op_name, exc_info=True | ||
| 146 | + ) | ||
| 147 | + return original_func(*args, **kwargs) | ||
| 148 | + finally: | ||
| 149 | + if isinstance(range_id, int) and range_id > 0: | ||
| 150 | + mstx.range_end(range_id, domain=_FLOPS_DOMAIN) | ||
| 151 | + cls._local.in_hook = False | ||
| 152 | + | ||
| 153 | + return wrapper | ||
| @@ -0,0 +1,55 @@ | |||
| 1 | +import logging | ||
| 2 | +from collections.abc import Callable | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +logger = logging.getLogger(__name__) | ||
| 6 | + | ||
| 7 | +_default_npu_flop_registry: dict[str, tuple[Callable, str | None]] = {} | ||
| 8 | +_npu_flop_registry: dict[str, tuple[Callable, str | None]] = {} | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +def register_npu_flop( | ||
| 12 | + target: str | None = None, | ||
| 13 | + op_name: str | None = None, | ||
| 14 | + *, | ||
| 15 | + is_default: bool = False, | ||
| 16 | +): | ||
| 17 | + def decorator(func: Callable) -> Callable: | ||
| 18 | + resolved_name = op_name | ||
| 19 | + if resolved_name is None: | ||
| 20 | + if target is not None and ":" in target: | ||
| 21 | + resolved_name = target.split(":", 1)[1] | ||
| 22 | + else: | ||
| 23 | + resolved_name = func.__name__ | ||
| 24 | + registry = _default_npu_flop_registry if is_default else _npu_flop_registry | ||
| 25 | + if not is_default and resolved_name in registry: | ||
| 26 | + logger.error( | ||
| 27 | + "Duplicate external FLOPs registration for %s. " | ||
| 28 | + "The later registration takes precedence.", | ||
| 29 | + resolved_name, | ||
| 30 | + ) | ||
| 31 | + registry[resolved_name] = (func, target) | ||
| 32 | + return func | ||
| 33 | + | ||
| 34 | + return decorator | ||
| 35 | + | ||
| 36 | + | ||
| 37 | +def get_flop_func(op_name: str) -> Callable | None: | ||
| 38 | + entry = _npu_flop_registry.get(op_name) or _default_npu_flop_registry.get(op_name) | ||
| 39 | + return entry[0] if entry else None | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +def get_npu_flop_targets() -> dict[str, str]: | ||
| 43 | + targets = { | ||
| 44 | + name: entry[1] | ||
| 45 | + for name, entry in _default_npu_flop_registry.items() | ||
| 46 | + if entry[1] is not None | ||
| 47 | + } | ||
| 48 | + targets.update( | ||
| 49 | + { | ||
| 50 | + name: entry[1] | ||
| 51 | + for name, entry in _npu_flop_registry.items() | ||
| 52 | + if entry[1] is not None | ||
| 53 | + } | ||
| 54 | + ) | ||
| 55 | + return targets | ||
| @@ -37,6 +37,7 @@ from .analysis.prof_common_func._utils import ( | |||
| 37 | no_exception_func, | 37 | no_exception_func, |
| 38 | ) | 38 | ) |
| 39 | from .experimental_config import _ExperimentalConfig | 39 | from .experimental_config import _ExperimentalConfig |
| 40 | +from ._flops_hook import FlopsHookManager | ||
| 40 | from .scheduler import ProfilerAction | 41 | from .scheduler import ProfilerAction |
| 41 | 42 | ||
| 42 | 43 | ||
| @@ -155,6 +156,8 @@ class _ProfInterface: | |||
| 155 | self.start_monotonic = _get_monotonic() | 156 | self.start_monotonic = _get_monotonic() |
| 156 | _enable_event_record() | 157 | _enable_event_record() |
| 157 | _start_profiler(npu_prof_config, self.activities) | 158 | _start_profiler(npu_prof_config, self.activities) |
| 159 | + if self.with_flops and self.experimental_config._msprof_tx: | ||
| 160 | + FlopsHookManager.install() | ||
| 158 | self.start_gc_detect() | 161 | self.start_gc_detect() |
| 159 | 162 | ||
| 160 | def stop_trace(self): | 163 | def stop_trace(self): |
| @@ -162,6 +165,8 @@ class _ProfInterface: | |||
| 162 | return | 165 | return |
| 163 | if ProfilerActivity.NPU in self.activities: | 166 | if ProfilerActivity.NPU in self.activities: |
| 164 | torch.npu.synchronize() | 167 | torch.npu.synchronize() |
| 168 | + if self.with_flops and self.experimental_config._msprof_tx: | ||
| 169 | + FlopsHookManager.uninstall() | ||
| 165 | _stop_profiler() | 170 | _stop_profiler() |
| 166 | self.stop_gc_detect() | 171 | self.stop_gc_detect() |
| 167 | _disable_event_record() | 172 | _disable_event_record() |