已合并
[Performance]: [v2.10.0]inductor reduce-overhead场景，有FA更新时，FA update动作和模型执行并行 #43349
dingdairong创建于 29 天前
[Performance]: [v2.10.0]inductor reduce-overhead场景，有FA更新时，FA update动作和模型执行并行 #43349
已合并
共 2 个文件变更+166-93
| @@ -1,4 +1,6 @@ | |||
| 1 | """AscendC backend: IFA v2 dynamic actual seq with aclgraph replay.""" | 1 | """AscendC backend: IFA v2 dynamic actual seq with aclgraph replay.""" |
| 2 | +from contextlib import contextmanager | ||
| 3 | +import os | ||
| 2 | import unittest | 4 | import unittest |
| 3 | 5 | ||
| 4 | import torch | 6 | import torch |
| @@ -19,6 +21,8 @@ SEQ_CONFIGS = ( | |||
| 19 | [80, 90], | 21 | [80, 90], |
| 20 | [128, 30], | 22 | [128, 30], |
| 21 | ) | 23 | ) |
| 24 | +DEFERRED_UPDATE_LOG = "NPUGRAPH-TREE ACLGraph update deferred until after replay" | ||
| 25 | +UPDATE_BEFORE_REPLAY_LOG = "NPUGRAPH-TREE ACLGraph update before replay" | ||
| 22 | 26 | ||
| 23 | 27 | ||
| 24 | def _make_inputs(): | 28 | def _make_inputs(): |
| @@ -42,23 +46,130 @@ def _ifa_v2(q, k, v, actual_seq_qlen, actual_seq_kvlen): | |||
| 42 | ) | 46 | ) |
| 43 | 47 | ||
| 44 | 48 | ||
| 49 | + | ||
| 50 | +def _temporary_env(name, value): | ||
| 51 | + old_value = os.environ.get(name) | ||
| 52 | + if value is None: | ||
| 53 | + os.environ.pop(name, None) | ||
| 54 | + else: | ||
| 55 | + os.environ[name] = value | ||
| 56 | + try: | ||
| 57 | + yield | ||
| 58 | + finally: | ||
| 59 | + if old_value is None: | ||
| 60 | + os.environ.pop(name, None) | ||
| 61 | + else: | ||
| 62 | + os.environ[name] = old_value | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +def _is_ascendc_backend_available(): | ||
| 66 | + try: | ||
| 67 | + x = torch.randn(4, 4, device="npu") | ||
| 68 | + torch.compile( | ||
| 69 | + lambda t: t + 1, | ||
| 70 | + backend="inductor", | ||
| 71 | + options={"npu_backend": "ascendc"}, | ||
| 72 | + )(x) | ||
| 73 | + torch.npu.synchronize() | ||
| 74 | + except Exception: | ||
| 75 | + return False | ||
| 76 | + return True | ||
| 77 | + | ||
| 78 | + | ||
| 79 | +def _run_dynamic_actual_seq_case(): | ||
| 80 | + import torch_npu._inductor.ascendc.config as ascendc_config | ||
| 81 | + | ||
| 82 | + old_cudagraphs = config.triton.cudagraphs | ||
| 83 | + old_cudagraph_trees = config.triton.cudagraph_trees | ||
| 84 | + old_force_disable_caches = config.force_disable_caches | ||
| 85 | + old_slow_path_asserts = config.triton.slow_path_cudagraph_asserts | ||
| 86 | + old_sync_around_fuse_kernel = ascendc_config.sync_around_fuse_kernel | ||
| 87 | + try: | ||
| 88 | + config.triton.cudagraphs = True | ||
| 89 | + config.triton.cudagraph_trees = True | ||
| 90 | + config.force_disable_caches = True | ||
| 91 | + config.triton.slow_path_cudagraph_asserts = False | ||
| 92 | + if os.getenv("ASCEND_LAUNCH_BLOCKING", None) == "1": | ||
| 93 | + ascendc_config.sync_around_fuse_kernel = False | ||
| 94 | + | ||
| 95 | + torch.manual_seed(0) | ||
| 96 | + torch_npu.npu.manual_seed(0) | ||
| 97 | + q, k, v = _make_inputs() | ||
| 98 | + compiled_ifa_v2 = torch.compile( | ||
| 99 | + _ifa_v2, | ||
| 100 | + backend="inductor", | ||
| 101 | + dynamic=True, | ||
| 102 | + fullgraph=True, | ||
| 103 | + options={"npu_backend": "ascendc", "triton.cudagraphs": True}, | ||
| 104 | + ) | ||
| 105 | + | ||
| 106 | + log_stream, ctx = logs_to_string("torch_npu.npugraph", "cudagraphs") | ||
| 107 | + with torch.no_grad(), ctx(): | ||
| 108 | + for seq in SEQ_CONFIGS: | ||
| 109 | + # Pass the same Python list object for qlen and kvlen, matching | ||
| 110 | + # decode-style actual seq inputs from the original reproducer. | ||
| 111 | + seq_arg = list(seq) | ||
| 112 | + compiled_out = compiled_ifa_v2(q, k, v, seq_arg, seq_arg) | ||
| 113 | + eager_out = _ifa_v2(q, k, v, seq_arg, seq_arg) | ||
| 114 | + diff = ( | ||
| 115 | + compiled_out[0].float() - eager_out[0].float() | ||
| 116 | + ).abs().max().item() | ||
| 117 | + if diff > 1e-2: | ||
| 118 | + raise AssertionError(f"seq={seq} max_diff={diff}") | ||
| 119 | + torch.npu.synchronize() | ||
| 120 | + | ||
| 121 | + return log_stream.getvalue() | ||
| 122 | + finally: | ||
| 123 | + config.triton.cudagraphs = old_cudagraphs | ||
| 124 | + config.triton.cudagraph_trees = old_cudagraph_trees | ||
| 125 | + config.force_disable_caches = old_force_disable_caches | ||
| 126 | + config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts | ||
| 127 | + ascendc_config.sync_around_fuse_kernel = old_sync_around_fuse_kernel | ||
| 128 | + torch._dynamo.reset() | ||
| 129 | + | ||
| 130 | + | ||
| 131 | +def _assert_dynamic_actual_seq_key_reuse(test_case, logs): | ||
| 132 | + compile_recording_lines = [ | ||
| 133 | + line for line in logs.splitlines() | ||
| 134 | + if "NPUGRAPH-TREE Compile recording" in line | ||
| 135 | + ] | ||
| 136 | + compile_recording_count = len(compile_recording_lines) | ||
| 137 | + warmup_count = logs.count("NPUGRAPH-TREE Warmup Running warmup") | ||
| 138 | + record_count = logs.count("NPUGRAPH-TREE Node Record function=") | ||
| 139 | + replay_count = logs.count("NPUGRAPH Replay graph_id=") | ||
| 140 | + update_count = logs.count("NPUGraph: updating graph") | ||
| 141 | + | ||
| 142 | + # The first concrete list may compile statically, then Dynamo may | ||
| 143 | + # generalize to a dynamic actual seq graph. Dynamic actual seq values | ||
| 144 | + # must not create one NPUGraphTree key per sequence pair. | ||
| 145 | + test_case.assertGreaterEqual(compile_recording_count, 1, logs) | ||
| 146 | + test_case.assertLessEqual(compile_recording_count, 2, logs) | ||
| 147 | + test_case.assertLessEqual(warmup_count, 2, logs) | ||
| 148 | + test_case.assertLessEqual(record_count, 2, logs) | ||
| 149 | + test_case.assertGreaterEqual(replay_count, 1, logs) | ||
| 150 | + test_case.assertEqual( | ||
| 151 | + update_count, | ||
| 152 | + len(SEQ_CONFIGS) - warmup_count, | ||
| 153 | + logs, | ||
| 154 | + ) | ||
| 155 | + test_case.assertIn("NPUGRAPH-TREE State state=EXECUTION", logs) | ||
| 156 | + for bad_key_fragment in ("100, 50", "80, 90", "128, 30"): | ||
| 157 | + test_case.assertFalse( | ||
| 158 | + any( | ||
| 159 | + bad_key_fragment in line | ||
| 160 | + for line in compile_recording_lines | ||
| 161 | + ), | ||
| 162 | + "\n".join(compile_recording_lines), | ||
| 163 | + ) | ||
| 164 | + | ||
| 165 | + | ||
| 45 | 166 | ||
| 46 | class TestAscendcIFAv2DynamicActualSeq(TestCase): | 167 | class TestAscendcIFAv2DynamicActualSeq(TestCase): |
| 47 | 168 | ||
| 48 | 169 | ||
| 49 | def setUpClass(cls): | 170 | def setUpClass(cls): |
| 50 | super().setUpClass() | 171 | super().setUpClass() |
| 51 | - cls._ascendc_ok = False | 172 | + cls._ascendc_ok = _is_ascendc_backend_available() |
| 52 | - try: | ||
| 53 | - x = torch.randn(4, 4, device="npu") | ||
| 54 | - torch.compile( | ||
| 55 | - lambda t: t + 1, | ||
| 56 | - backend="inductor", | ||
| 57 | - options={"npu_backend": "ascendc"}, | ||
| 58 | - )(x) | ||
| 59 | - cls._ascendc_ok = True | ||
| 60 | - except Exception: | ||
| 61 | - pass | ||
| 62 | 173 | ||
| 63 | def setUp(self): | 174 | def setUp(self): |
| 64 | super().setUp() | 175 | super().setUp() |
| @@ -71,83 +182,17 @@ class TestAscendcIFAv2DynamicActualSeq(TestCase): | |||
| 71 | super().tearDown() | 182 | super().tearDown() |
| 72 | 183 | ||
| 73 | def test_dynamic_actual_seq_reuses_npugraph_key(self): | 184 | def test_dynamic_actual_seq_reuses_npugraph_key(self): |
| 74 | - old_cudagraphs = config.triton.cudagraphs | 185 | + with _temporary_env("ASCEND_LAUNCH_BLOCKING", None): |
| 75 | - old_cudagraph_trees = config.triton.cudagraph_trees | 186 | + logs = _run_dynamic_actual_seq_case() |
| 76 | - old_force_disable_caches = config.force_disable_caches | 187 | + _assert_dynamic_actual_seq_key_reuse(self, logs) |
| 77 | - old_slow_path_asserts = config.triton.slow_path_cudagraph_asserts | 188 | + self.assertIn(DEFERRED_UPDATE_LOG, logs) |
| 78 | - try: | 189 | + self.assertNotIn(UPDATE_BEFORE_REPLAY_LOG, logs) |
| 79 | - config.triton.cudagraphs = True | ||
| 80 | - config.triton.cudagraph_trees = True | ||
| 81 | - config.force_disable_caches = True | ||
| 82 | - config.triton.slow_path_cudagraph_asserts = False | ||
| 83 | 190 | ||
| 84 | - torch.manual_seed(0) | 191 | + with _temporary_env("ASCEND_LAUNCH_BLOCKING", "1"): |
| 85 | - torch_npu.npu.manual_seed(0) | 192 | + logs = _run_dynamic_actual_seq_case() |
| 86 | - q, k, v = _make_inputs() | 193 | + _assert_dynamic_actual_seq_key_reuse(self, logs) |
| 87 | - compiled_ifa_v2 = torch.compile( | 194 | + self.assertIn(UPDATE_BEFORE_REPLAY_LOG, logs) |
| 88 | - _ifa_v2, | 195 | + self.assertNotIn(DEFERRED_UPDATE_LOG, logs) |
| 89 | - backend="inductor", | ||
| 90 | - dynamic=True, | ||
| 91 | - fullgraph=True, | ||
| 92 | - options={"npu_backend": "ascendc", "triton.cudagraphs": True}, | ||
| 93 | - ) | ||
| 94 | - | ||
| 95 | - log_stream, ctx = logs_to_string("torch_npu.npugraph", "cudagraphs") | ||
| 96 | - with torch.no_grad(), ctx(): | ||
| 97 | - for seq in SEQ_CONFIGS: | ||
| 98 | - # Pass the same Python list object for qlen and kvlen, matching | ||
| 99 | - # decode-style actual seq inputs from the original reproducer. | ||
| 100 | - seq_arg = list(seq) | ||
| 101 | - compiled_out = compiled_ifa_v2(q, k, v, seq_arg, seq_arg) | ||
| 102 | - eager_out = _ifa_v2(q, k, v, seq_arg, seq_arg) | ||
| 103 | - diff = ( | ||
| 104 | - compiled_out[0].float() - eager_out[0].float() | ||
| 105 | - ).abs().max().item() | ||
| 106 | - self.assertLessEqual( | ||
| 107 | - diff, | ||
| 108 | - 1e-2, | ||
| 109 | - f"seq={seq} max_diff={diff}", | ||
| 110 | - ) | ||
| 111 | - torch.npu.synchronize() | ||
| 112 | - | ||
| 113 | - logs = log_stream.getvalue() | ||
| 114 | - compile_recording_lines = [ | ||
| 115 | - line for line in logs.splitlines() | ||
| 116 | - if "NPUGRAPH-TREE Compile recording" in line | ||
| 117 | - ] | ||
| 118 | - compile_recording_count = len(compile_recording_lines) | ||
| 119 | - warmup_count = logs.count("NPUGRAPH-TREE Warmup Running warmup") | ||
| 120 | - record_count = logs.count("NPUGRAPH-TREE Node Record function=") | ||
| 121 | - replay_count = logs.count("NPUGRAPH Replay graph_id=") | ||
| 122 | - update_count = logs.count("NPUGraph: updating graph") | ||
| 123 | - | ||
| 124 | - # The first concrete list may compile statically, then Dynamo may | ||
| 125 | - # generalize to a dynamic actual seq graph. Dynamic actual seq values | ||
| 126 | - # must not create one NPUGraphTree key per sequence pair. | ||
| 127 | - self.assertGreaterEqual(compile_recording_count, 1, logs) | ||
| 128 | - self.assertLessEqual(compile_recording_count, 2, logs) | ||
| 129 | - self.assertLessEqual(warmup_count, 2, logs) | ||
| 130 | - self.assertLessEqual(record_count, 2, logs) | ||
| 131 | - self.assertGreaterEqual(replay_count, 1, logs) | ||
| 132 | - self.assertEqual( | ||
| 133 | - update_count, | ||
| 134 | - len(SEQ_CONFIGS) - warmup_count, | ||
| 135 | - logs, | ||
| 136 | - ) | ||
| 137 | - self.assertIn("NPUGRAPH-TREE State state=EXECUTION", logs) | ||
| 138 | - for bad_key_fragment in ("100, 50", "80, 90", "128, 30"): | ||
| 139 | - self.assertFalse( | ||
| 140 | - any( | ||
| 141 | - bad_key_fragment in line | ||
| 142 | - for line in compile_recording_lines | ||
| 143 | - ), | ||
| 144 | - "\n".join(compile_recording_lines), | ||
| 145 | - ) | ||
| 146 | - finally: | ||
| 147 | - config.triton.cudagraphs = old_cudagraphs | ||
| 148 | - config.triton.cudagraph_trees = old_cudagraph_trees | ||
| 149 | - config.force_disable_caches = old_force_disable_caches | ||
| 150 | - config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts | ||
| 151 | 196 | ||
| 152 | 197 | ||
| 153 | if __name__ == "__main__": | 198 | if __name__ == "__main__": |
| @@ -41,6 +41,7 @@ import functools | |||
| 41 | import gc | 41 | import gc |
| 42 | import itertools | 42 | import itertools |
| 43 | import operator | 43 | import operator |
| 44 | +import os | ||
| 44 | import sys | 45 | import sys |
| 45 | import threading | 46 | import threading |
| 46 | import traceback | 47 | import traceback |
| @@ -126,6 +127,20 @@ S = TypeVar("S", bound="StorageWeakRefWrapper") | |||
| 126 | log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs") | 127 | log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs") |
| 127 | 128 | ||
| 128 | 129 | ||
| 130 | +def _npu_launch_blocking_enabled() -> bool: | ||
| 131 | + return os.getenv("ASCEND_LAUNCH_BLOCKING", None) == "1" | ||
| 132 | + | ||
| 133 | + | ||
| 134 | +def _can_defer_aclgraph_update_after_replay(graph: Any, cpu_update_input: Any) -> bool: | ||
| 135 | + if graph is None or not getattr(graph, "auto_dispatch_capture", False): | ||
| 136 | + return False | ||
| 137 | + if not cpu_update_input: | ||
| 138 | + return False | ||
| 139 | + if _npu_launch_blocking_enabled(): | ||
| 140 | + return False | ||
| 141 | + return True | ||
| 142 | + | ||
| 143 | + | ||
| 129 | 144 | ||
| 130 | class GraphID: | 145 | class GraphID: |
| 131 | "Unique counter of a npu graph recording" | 146 | "Unique counter of a npu graph recording" |
| @@ -1088,13 +1103,26 @@ class NPUGraphNode: | |||
| 1088 | def run(self, new_inputs: List[InputType]) -> OutputType: | 1103 | def run(self, new_inputs: List[InputType]) -> OutputType: |
| 1089 | log.debug("NPUGRAPH-TREE Node Run node=%s", self.id) | 1104 | log.debug("NPUGRAPH-TREE Node Run node=%s", self.id) |
| 1090 | self.check_static_inputs_are_stable(new_inputs) | 1105 | self.check_static_inputs_are_stable(new_inputs) |
| 1091 | - aclgraph_update_submitted = update_aclgraph_records_for_graph( | 1106 | + aclgraph_cpu_update_input = resolve_aclgraph_update_plan( |
| 1092 | - resolve_aclgraph_update_plan(self.aclgraph_update_plan, new_inputs), | 1107 | + self.aclgraph_update_plan, new_inputs |
| 1093 | - self.graph, | ||
| 1094 | ) | 1108 | ) |
| 1095 | - self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs) | 1109 | + if _can_defer_aclgraph_update_after_replay(self.graph, aclgraph_cpu_update_input): |
| 1096 | - | 1110 | + log.debug("NPUGRAPH-TREE ACLGraph update deferred until after replay") |
| 1097 | - self.run_graph() | 1111 | + self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs) |
| 1112 | + self.run_graph() | ||
| 1113 | + aclgraph_update_submitted = update_aclgraph_records_for_graph( | ||
| 1114 | + aclgraph_cpu_update_input, | ||
| 1115 | + self.graph, | ||
| 1116 | + ) | ||
| 1117 | + else: | ||
| 1118 | + if aclgraph_cpu_update_input: | ||
| 1119 | + log.debug("NPUGRAPH-TREE ACLGraph update before replay") | ||
| 1120 | + aclgraph_update_submitted = update_aclgraph_records_for_graph( | ||
| 1121 | + aclgraph_cpu_update_input, | ||
| 1122 | + self.graph, | ||
| 1123 | + ) | ||
| 1124 | + self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs) | ||
| 1125 | + self.run_graph() | ||
| 1098 | if aclgraph_update_submitted: | 1126 | if aclgraph_update_submitted: |
| 1099 | # Ensure the next ACLGraph update does not record reusable external events before this replay resets them. | 1127 | # Ensure the next ACLGraph update does not record reusable external events before this replay resets them. |
| 1100 | self.graph.graph_dispatch_mode.update_stream.wait_stream( | 1128 | self.graph.graph_dispatch_mode.update_stream.wait_stream( |