已合并
[sync] PR-39974: bugfix for profiler FLOPs #40031
ascend-robot创建于 7月3日
[sync] PR-39974: bugfix for profiler FLOPs #40031
已合并
共 3 个文件变更+26-3
| @@ -33,8 +33,9 @@ def npu_fusion_attention_flops( | |||
| 33 | return _calculate_tnd_layout_flops( | 33 | return _calculate_tnd_layout_flops( |
| 34 | q_shape, k_shape, v_shape, actual_seq_qlen, actual_seq_kvlen | 34 | q_shape, k_shape, v_shape, actual_seq_qlen, actual_seq_kvlen |
| 35 | ) | 35 | ) |
| 36 | + kv_heads = _infer_kv_heads(q_shape, k_shape, input_layout, head_num) | ||
| 36 | return _calculate_common_layout_flops( | 37 | return _calculate_common_layout_flops( |
| 37 | - q_shape, k_shape, v_shape, input_layout, sparse_mode, head_num, head_num | 38 | + q_shape, k_shape, v_shape, input_layout, sparse_mode, head_num, kv_heads |
| 38 | ) | 39 | ) |
| 39 | 40 | ||
| 40 | 41 | ||
| @@ -424,6 +425,24 @@ def _calculate_tnd_layout_flops( | |||
| 424 | return int(2 * (q_heads or shape_q_heads) * (q_d + v_d) * attention_scores) | 425 | return int(2 * (q_heads or shape_q_heads) * (q_d + v_d) * attention_scores) |
| 425 | 426 | ||
| 426 | 427 | ||
| 428 | +def _infer_kv_heads(q_shape, k_shape, input_layout, q_heads): | ||
| 429 | + if input_layout == "BNSD": | ||
| 430 | + return k_shape[1] | ||
| 431 | + if input_layout == "BSND": | ||
| 432 | + return k_shape[2] | ||
| 433 | + if input_layout == "BSH": | ||
| 434 | + _, _, q_hidden = q_shape | ||
| 435 | + _, _, k_hidden = k_shape | ||
| 436 | + elif input_layout == "SBH": | ||
| 437 | + _, _, q_hidden = q_shape | ||
| 438 | + _, _, k_hidden = k_shape | ||
| 439 | + else: | ||
| 440 | + return q_heads | ||
| 441 | + | ||
| 442 | + q_head_dim = _head_dim(q_hidden, q_heads) | ||
| 443 | + return _head_dim(k_hidden, q_head_dim) | ||
| 444 | + | ||
| 445 | + | ||
| 427 | def _calculate_attention_scores(q_s, k_s, sparse_mode): | 446 | def _calculate_attention_scores(q_s, k_s, sparse_mode): |
| 428 | if sparse_mode == 0: | 447 | if sparse_mode == 0: |
| 429 | return q_s * k_s | 448 | return q_s * k_s |
| @@ -148,6 +148,10 @@ class _ExperimentalConfig: | |||
| 148 | def gc_detect_threshold(self): | 148 | def gc_detect_threshold(self): |
| 149 | return self._gc_detect_threshold | 149 | return self._gc_detect_threshold |
| 150 | 150 | ||
| 151 | + | ||
| 152 | + def tx_enabled(self): | ||
| 153 | + return self._msprof_tx or self._mstx | ||
| 154 | + | ||
| 151 | def _conver_export_type_to_list(self, export_type: Union[str, list]) -> list: | 155 | def _conver_export_type_to_list(self, export_type: Union[str, list]) -> list: |
| 152 | if not export_type: | 156 | if not export_type: |
| 153 | print_warn_msg( | 157 | print_warn_msg( |
| @@ -156,7 +156,7 @@ class _ProfInterface: | |||
| 156 | self.start_monotonic = _get_monotonic() | 156 | self.start_monotonic = _get_monotonic() |
| 157 | _enable_event_record() | 157 | _enable_event_record() |
| 158 | _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: | 159 | + if self.with_flops and self.experimental_config.tx_enabled: |
| 160 | FlopsHookManager.install() | 160 | FlopsHookManager.install() |
| 161 | self.start_gc_detect() | 161 | self.start_gc_detect() |
| 162 | 162 | ||
| @@ -165,7 +165,7 @@ class _ProfInterface: | |||
| 165 | return | 165 | return |
| 166 | if ProfilerActivity.NPU in self.activities: | 166 | if ProfilerActivity.NPU in self.activities: |
| 167 | torch.npu.synchronize() | 167 | torch.npu.synchronize() |
| 168 | - if self.with_flops and self.experimental_config._msprof_tx: | 168 | + if self.with_flops and self.experimental_config.tx_enabled: |
| 169 | FlopsHookManager.uninstall() | 169 | FlopsHookManager.uninstall() |
| 170 | _stop_profiler() | 170 | _stop_profiler() |
| 171 | self.stop_gc_detect() | 171 | self.stop_gc_detect() |