Star846
1320
代码
代码
Issues1492
Pull Requests2314
流水线
项目讨论
Wiki
分析
Star846
1320
已合并
[inductor][Feature] Adaptive Gear Update #35530
zhudada0120创建于 5月13日
[inductor][Feature] Adaptive Gear Update #35530
已合并
zhudada0120创建于 5月13日
从zhudada/pytorchv2.7.1-gear合入到Ascend/pytorchv2.7.1
共 7 个文件变更+2690-56
Mtest/_inductor/test_shape_handling.py+1051-29
@@ -1,5 +1,7 @@
1import unittest1import unittest
2from unittest import mock2from unittest import mock
3+import threading
4+import time
3 5 
4import torch6import torch
5from torch.nn.parallel import scatter_gather7from torch.nn.parallel import scatter_gather
@@ -62,46 +64,80 @@ shape_options = {
62 64 
63class TestShapeHandling(TestCase):65class TestShapeHandling(TestCase):
64 def test_init_no_input(self):66 def test_init_no_input(self):
67+ """Constructing NPUShapeHandling with no arguments should succeed (delay_init path)."""
65 shape_handling = torch_npu._inductor.NPUShapeHandling()68 shape_handling = torch_npu._inductor.NPUShapeHandling()
69+ # object created successfully
66 self.assertNotEqual(shape_handling, None)70 self.assertNotEqual(shape_handling, None)
67- 71+ 
68 def test_init_with_empty_conifg(self):72 def test_init_with_empty_conifg(self):
73+ """Empty config list should not crash — falls through to delay_init."""
69 configs = []74 configs = []
70 shape_handling = torch_npu._inductor.NPUShapeHandling(configs)75 shape_handling = torch_npu._inductor.NPUShapeHandling(configs)
71 self.assertNotEqual(shape_handling, None)76 self.assertNotEqual(shape_handling, None)
72- 77+ 
73 def test_init_with_gears(self):78 def test_init_with_gears(self):
79+ """Explicit gears for two dimension types should be accepted."""
74 configs = [80 configs = [
75- {81+ {"type": "BATCHSIZE", "gears": [16, 32, 64]},
76- "type": "BATCHSIZE",82+ {"type": "SEQLEN", "gears": [16, 32, 64]},
77- "gears": [16, 32, 64]
78- },
79- {
80- "type": "SEQLEN",
81- "gears": [16, 32, 64]
82- }
83 ]83 ]
84 shape_handling = torch_npu._inductor.NPUShapeHandling(configs)84 shape_handling = torch_npu._inductor.NPUShapeHandling(configs)
85 self.assertNotEqual(shape_handling, None)85 self.assertNotEqual(shape_handling, None)
86- 86+ 
87 def test_init_with_policy(self):87 def test_init_with_policy(self):
88+ """TIMES policy with min/max should auto-generate gears."""
88 configs = [89 configs = [
89- {90+ {"type": "BATCHSIZE", "min_size": 2, "max_size": 8, "policy": "TIMES"},
90- "type": "BATCHSIZE",91+ {"type": "SEQLEN", "min_size": 2, "max_size": 8, "policy": "TIMES"},
91- "min_size": 2,
92- "max_size": 8,
93- "policy": "TIMES"
94- },
95- {
96- "type": "SEQLEN",
97- "min_size": 2,
98- "max_size": 8,
99- "policy": "TIMES"
100- }
101 ]92 ]
102 shape_handling = torch_npu._inductor.NPUShapeHandling(configs)93 shape_handling = torch_npu._inductor.NPUShapeHandling(configs)
103 self.assertNotEqual(shape_handling, None)94 self.assertNotEqual(shape_handling, None)
104- 95+ 
96+ def test_normalize_configs_edge_cases(self):
97+ """_normalize_configs handles empty input, dedup, float→int, TIMES expansion."""
98+ from torch_npu._inductor.adaptive_gears import AdaptiveGearRuntime
99+ from unittest.mock import MagicMock
100+ 
101+ runtime = MagicMock(spec=AdaptiveGearRuntime)
102+ runtime._expand_policy_gears = lambda config: AdaptiveGearRuntime._expand_policy_gears(runtime, config)
103+ 
104+ # Empty / None input → returns empty list without crashing
105+ self.assertEqual(AdaptiveGearRuntime._normalize_configs(runtime, []), [])
106+ self.assertEqual(AdaptiveGearRuntime._normalize_configs(runtime, None), [])
107+ 
108+ # Duplicate gears deduped and sorted ascending
109+ config = {"type": "BATCHSIZE", "gears": [64, 16, 32, 16]}
110+ result = AdaptiveGearRuntime._normalize_configs(runtime, [config])
111+ self.assertEqual(result[0]["gears"], [16, 32, 64])
112+ 
113+ # Float gears (e.g. from JSON) cast to int
114+ config = {"type": "BATCHSIZE", "gears": [16.0, 32.0]}
115+ result = AdaptiveGearRuntime._normalize_configs(runtime, [config])
116+ self.assertEqual(result[0]["gears"], [16, 32])
117+ # each value is a plain int
118+ self.assertIsInstance(result[0]["gears"][0], int)
119+ 
120+ # TIMES: min == max → single gear (no expansion needed)
121+ config = {"type": "BATCHSIZE", "min_size": 4, "max_size": 4}
122+ result = AdaptiveGearRuntime._normalize_configs(runtime, [config])
123+ self.assertEqual(result[0]["gears"], [4])
124+ 
125+ # TIMES: non-power-of-2 min_size → min preserved, then powers of 2, max appended
126+ config = {"type": "BATCHSIZE", "min_size": 3, "max_size": 100}
127+ result = AdaptiveGearRuntime._normalize_configs(runtime, [config])
128+ # min_size is the anchor
129+ self.assertEqual(result[0]["gears"][0], 3)
130+ # max_size is the cap
131+ self.assertEqual(result[0]["gears"][-1], 100)
132+ for g in result[0]["gears"][1:-1]:
133+ # power-of-2 check
134+ self.assertEqual(g & (g - 1), 0, f"{g} is not a power of 2")
135+ 
136+ # TIMES: power-of-2 min_size → next power of 2, then double
137+ config = {"type": "BATCHSIZE", "min_size": 8, "max_size": 128}
138+ result = AdaptiveGearRuntime._normalize_configs(runtime, [config])
139+ self.assertEqual(result[0]["gears"], [8, 16, 32, 64, 128])
140+ 
105 def test_transform_no_operation(self):141 def test_transform_no_operation(self):
106 configs = [142 configs = [
107 {143 {
@@ -187,9 +223,12 @@ class TestShapeHandling(TestCase):
187 input_tensor = torch.randn(200, 96) # 超过max_size223 input_tensor = torch.randn(200, 96) # 超过max_size
188 outputs = shape_handling.transform([input_tensor])224 outputs = shape_handling.transform([input_tensor])
189 # 验证分割结果：分割为两个组，第一段128，第二段72，再填充为128225 # 验证分割结果：分割为两个组，第一段128，第二段72，再填充为128
190- self.assertEqual(len(outputs), 2) # 分割为两个组226+ # 分割为两个组
191- self.assertEqual(outputs[0][0].shape, (128, 128)) # 第一段227+ self.assertEqual(len(outputs), 2)
192- self.assertEqual(outputs[1][0].shape, (128, 128)) # 第二段228+ # 第一段
229+ self.assertEqual(outputs[0][0].shape, (128, 128))
230+ # 第二段
231+ self.assertEqual(outputs[1][0].shape, (128, 128))
193 232
194 def test_recover_padding(self):233 def test_recover_padding(self):
195 """测试恢复填充的张量"""234 """测试恢复填充的张量"""
@@ -375,6 +414,40 @@ class TestShapeHandlingBranchCoverage(TestCase):
375 [{"type": "BATCHSIZE"}, {"type": "SEQLEN"}, {"type": "BATCHSIZE"}]414 [{"type": "BATCHSIZE"}, {"type": "SEQLEN"}, {"type": "BATCHSIZE"}]
376 )415 )
377 416 
417+ def test_validate_adaptive_configs(self):
418+ """Test adaptive config validation: unknown keys, range errors, type compatibility."""
419+ shape_handling = torch_npu._inductor.NPUShapeHandling()
420+ 
421+ # Unknown key
422+ with self.assertRaises(ValueError):
423+ shape_handling._validate_adaptive_configs({"unknown_key": 1})
424+ 
425+ # Ratio out of range
426+ with self.assertRaises(ValueError):
427+ shape_handling._validate_adaptive_configs({"weight_hit": 1.5})
428+ with self.assertRaises(ValueError):
429+ shape_handling._validate_adaptive_configs({"pad_add_threshold": -0.1})
430+ 
431+ # Seconds must be >= 0
432+ with self.assertRaises(ValueError):
433+ shape_handling._validate_adaptive_configs({"window_seconds": -1})
434+ # Zero is allowed (business layer guards against division-by-zero)
435+ shape_handling._validate_adaptive_configs({
436+ "window_seconds": 0,
437+ "recent_use_protect_seconds": 0,
438+ "update_interval_seconds": 0,
439+ })
440+ 
441+ # Int must be >= 1
442+ with self.assertRaises(ValueError):
443+ shape_handling._validate_adaptive_configs({"min_samples_per_gear": 0})
444+ 
445+ # Int accepted for float fields (ratio, seconds)
446+ shape_handling._validate_adaptive_configs({
447+ "weight_hit": 1, # int for ratio field
448+ "window_seconds": 300, # int for seconds field
449+ })
450+ 
378 def test_construct_indices_branches(self):451 def test_construct_indices_branches(self):
379 shape_handling = torch_npu._inductor.NPUShapeHandling()452 shape_handling = torch_npu._inductor.NPUShapeHandling()
380 tensors = [torch.randn(2, 3), torch.randn(2)]453 tensors = [torch.randn(2, 3), torch.randn(2)]
@@ -455,7 +528,7 @@ class TestShapeHandlingBranchCoverage(TestCase):
455 528 
456 def post_fn(outputs):529 def post_fn(outputs):
457 recorded["post"] = outputs530 recorded["post"] = outputs
458- return [("ok",)], [{"done": True}]531+ return [["ok"]], [{"done": True}]
459 532 
460 shape_handling_custom = torch_npu._inductor.NPUShapeHandling(533 shape_handling_custom = torch_npu._inductor.NPUShapeHandling(
461 transform_pre_fn=pre_fn,534 transform_pre_fn=pre_fn,
@@ -467,7 +540,7 @@ class TestShapeHandlingBranchCoverage(TestCase):
467 out_args, out_kwargs = shape_handling_custom.transform_hook(torch.tensor([7.0]))540 out_args, out_kwargs = shape_handling_custom.transform_hook(torch.tensor([7.0]))
468 self.assertIn("pre", recorded)541 self.assertIn("pre", recorded)
469 self.assertIn("post", recorded)542 self.assertIn("post", recorded)
470- self.assertEqual(out_args, [("ok",)])543+ self.assertEqual(out_args, [["ok"]])
471 self.assertEqual(out_kwargs, [{"done": True}])544 self.assertEqual(out_kwargs, [{"done": True}])
472 545 
473 shape_handling_none = torch_npu._inductor.NPUShapeHandling(transform_post_fn=lambda _: None)546 shape_handling_none = torch_npu._inductor.NPUShapeHandling(transform_post_fn=lambda _: None)
@@ -520,6 +593,643 @@ class TestShapeHandlingBranchCoverage(TestCase):
520 if hasattr(shape_handling_module.patch_shape_handling, "_is_patched"):593 if hasattr(shape_handling_module.patch_shape_handling, "_is_patched"):
521 delattr(shape_handling_module.patch_shape_handling, "_is_patched")594 delattr(shape_handling_module.patch_shape_handling, "_is_patched")
522 595 
596+ def test_transform_metadata_collection(self):
597+ configs = [
598+ {
599+ "type": "BATCHSIZE",
600+ "gears": [64],
601+ "dimensions": 0,
602+ "indices": [0],
603+ },
604+ {
605+ "type": "SEQLEN",
606+ "gears": [128],
607+ "dimensions": [1],
608+ "indices": [0],
609+ },
610+ ]
611+ shape_handling = torch_npu._inductor.NPUShapeHandling(configs)
612+ with mock.patch.object(
613+ torch_npu._inductor.NPUShapeHandling,
614+ "transform",
615+ return_value=[[torch.randn(64, 128)], [torch.randn(64, 128)]],
616+ ):
617+ _, metadata = shape_handling._transform_with_metadata(torch.randn(96, 96))
618+ self.assertEqual(len(metadata), 2)
619+ self.assertEqual(metadata[0]["raw_gear_values"], [[96], [96]])
620+ self.assertEqual(metadata[0]["mapped_gear_values"], [[64], [128]])
621+ self.assertGreater(metadata[0]["pad_ratios"][1][0], 0.0)
622+ self.assertGreater(metadata[0]["split_ratios"][0][0], 0.0)
623+ 
624+ def test_transform_metadata_collection_degrades_on_failure(self):
625+ """Metadata collection failure should not affect transform output."""
626+ configs = [{"type": "BATCHSIZE", "gears": [32], "dimensions": 0, "indices": [0]}]
627+ shape_handling = torch_npu._inductor.NPUShapeHandling(configs)
628+ 
629+ with mock.patch.object(
630+ torch_npu._inductor.NPUShapeHandling, "transform",
631+ return_value=[[torch.randn(32, 8)]],
632+ ):
633+ with mock.patch(
634+ "torch_npu._inductor.adaptive_gears.collect_transform_metadata",
635+ side_effect=RuntimeError("boom"),
636+ ):
637+ outputs, metadata = shape_handling._transform_with_metadata(torch.randn(16, 8))
638+ 
639+ # Transform output is intact
640+ self.assertEqual(len(outputs[0]), 1)
641+ # Metadata gracefully degraded to empty
642+ self.assertEqual(metadata, [])
643+ 
644+class TestAsyncWorkerAndConcurrency(TestCase):
645+ """测试异步Worker和并发控制机制"""
646+ 
647+ def test_async_worker_creation_and_execution(self):
648+ """测试后台Worker线程异步执行 run_once"""
649+ shape_handling = torch_npu._inductor.NPUShapeHandling(
650+ configs=[
651+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
652+ ],
653+ adaptive_configs={
654+ "update_interval_seconds": 0.1,
655+ "min_samples_per_gear": 1,
656+ "min_gear_count_per_type": 1,
657+ },
658+ )
659+ manager = shape_handling.adaptive_manager
660+ # adaptive manager created
661+ self.assertIsNotNone(manager)
662+ # worker attached
663+ self.assertIsNotNone(manager.worker)
664+ # daemon thread running
665+ self.assertTrue(manager._worker_thread.is_alive())
666+ 
667+ run_event = threading.Event()
668+ run_thread_id = None
669+ 
670+ def mock_run_once(ts):
671+ nonlocal run_thread_id
672+ run_thread_id = threading.current_thread().ident
673+ run_event.set()
674+ 
675+ with mock.patch.object(manager.worker, "run_once", side_effect=mock_run_once):
676+ # worker called run_once
677+ self.assertTrue(run_event.wait(timeout=5.0))
678+ 
679+ # thread recorded id
680+ self.assertIsNotNone(run_thread_id)
681+ # ran on background thread
682+ self.assertNotEqual(run_thread_id, threading.current_thread().ident)
683+ manager.shutdown()
684+ 
685+ def test_update_loop_survives_exception(self):
686+ """Worker thread should continue after run_once raises."""
687+ shape_handling = torch_npu._inductor.NPUShapeHandling(
688+ configs=[
689+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
690+ ],
691+ adaptive_configs={
692+ "update_interval_seconds": 0.1,
693+ "min_samples_per_gear": 1,
694+ "min_gear_count_per_type": 1,
695+ },
696+ )
697+ manager = shape_handling.adaptive_manager
698+ 
699+ call_count = 0
700+ run_event = threading.Event()
701+ 
702+ def mock_run_once(ts):
703+ nonlocal call_count
704+ call_count += 1
705+ if call_count == 1:
706+ raise RuntimeError("simulated failure")
707+ run_event.set()
708+ 
709+ with mock.patch.object(manager.worker, "run_once", side_effect=mock_run_once):
710+ # 2nd call succeeded
711+ self.assertTrue(run_event.wait(timeout=5.0))
712+ 
713+ # 1st raised, 2nd invoked → loop survived
714+ self.assertGreaterEqual(call_count, 2)
715+ manager.shutdown()
716+ 
717+ def test_snapshot_isolation_after_update(self):
718+ """Snapshot taken before an update must still reflect old gear set (clone-on-read)."""
719+ shape_handling = torch_npu._inductor.NPUShapeHandling(
720+ configs=[
721+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
722+ ],
723+ adaptive_configs={
724+ "recent_use_protect_seconds": 1.0,
725+ "min_samples_per_gear": 1,
726+ "min_gear_count_per_type": 1,
727+ },
728+ )
729+ manager = shape_handling.adaptive_manager
730+ 
731+ snapshot1 = manager.get_snapshot()
732+ original_gears = snapshot1.active_gears["BATCHSIZE"].copy()
733+ 
734+ manager.record_event([[20]], [[32]], [[0.375]], [[0.0]], 100.0)
735+ manager.worker.run_once(150.0)
736+ 
737+ # old snapshot unchanged
738+ self.assertEqual(snapshot1.active_gears["BATCHSIZE"], original_gears)
739+ manager.shutdown()
740+ 
741+ def test_high_concurrent_gear_update_scenarios(self):
742+ """测试高并发场景下的gear更新"""
743+ shape_handling = torch_npu._inductor.NPUShapeHandling(
744+ configs=[
745+ {"type": "BATCHSIZE", "gears": [16, 32, 64], "dimensions": 0, "indices": [0]},
746+ ],
747+ adaptive_configs={
748+ "recent_use_protect_seconds": 1.0,
749+ "min_samples_per_gear": 1,
750+ "add_min_samples": 1,
751+ "min_gear_count_per_type": 2,
752+ "recent_use_protect_seconds": 0.0,
753+ "replace_loss_threshold": 0.20,
754+ },
755+ )
756+ manager = shape_handling.adaptive_manager
757+ 
758+ def run_update(timestamp):
759+ manager.record_event([[16], [64]], [[32], [64]], [[0.5], [0.0]], [[0.0], [0.0]], 100.0 + timestamp)
760+ manager.worker.run_once(200.0 + timestamp)
761+ 
762+ threads = []
763+ for i in range(5):
764+ thread = threading.Thread(target=run_update, args=(i,))
765+ threads.append(thread)
766+ thread.start()
767+ 
768+ for thread in threads:
769+ thread.join()
770+ 
771+ latest_snapshot = manager.get_snapshot()
772+ self.assertIsNotNone(latest_snapshot)
773+ self.assertEqual(sorted(latest_snapshot.active_gears["BATCHSIZE"]), [32, 64])
774+ manager.shutdown()
775+ 
776+ 
777+class TestAdaptiveShapeHandling(TestCase):
778+ def test_npu_shape_handling_creates_snapshot_handler(self):
779+ shape_handling = torch_npu._inductor.NPUShapeHandling(
780+ configs=[
781+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
782+ {"type": "SEQLEN", "gears": [64, 128], "dimensions": [1], "indices": [0]},
783+ ],
784+ adaptive_configs={"recent_use_protect_seconds": 1.0},
785+ )
786+ manager = shape_handling.adaptive_manager
787+ snapshot = manager.get_snapshot()
788+ self.assertEqual(snapshot.active_gears["BATCHSIZE"], [16, 32])
789+ self.assertEqual(snapshot.active_gears["SEQLEN"], [64, 128])
790+ self.assertIsNotNone(snapshot.shape_handling)
791+ 
792+ def test_npu_shape_handling_records_event_and_builds_stats(self):
793+ shape_handling = torch_npu._inductor.NPUShapeHandling(
794+ configs=[
795+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
796+ {"type": "SEQLEN", "gears": [64, 128], "dimensions": [1], "indices": [0]},
797+ ],
798+ adaptive_configs={"window_seconds": 300.0},
799+ )
800+ manager = shape_handling.adaptive_manager
801+ snapshot = manager.get_snapshot()
802+ manager.record_event(
803+ raw_gear_values=[[20], [96]],
804+ mapped_gear_values=[[32], [128]],
805+ pad_ratios=[[0.375], [0.25]],
806+ split_ratios=[[0.0], [0.0]],
807+ event_ts=100.0,
808+ )
809+ 
810+ stats = manager.build_stats_snapshot(100.0)
811+ self.assertIn("BATCHSIZE:32", stats)
812+ self.assertIn("SEQLEN:128", stats)
813+ self.assertEqual(stats["BATCHSIZE:32"]["sample_count"], 1)
814+ self.assertEqual(stats["BATCHSIZE:32"]["pad_sample_count"], 1)
815+ self.assertEqual(stats["BATCHSIZE:32"]["split_sample_count"], 0)
816+ self.assertAlmostEqual(stats["BATCHSIZE:32"]["avg_pad_ratio"], 0.375)
817+ self.assertAlmostEqual(stats["SEQLEN:128"]["avg_pad_ratio"], 0.25)
818+ 
819+ def test_npu_shape_handling_records_cleanup_keys_per_gear(self):
820+ shape_handling = torch_npu._inductor.NPUShapeHandling(
821+ configs=[
822+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
823+ {"type": "SEQLEN", "gears": [64, 128], "dimensions": [1], "indices": [0]},
824+ ],
825+ adaptive_configs={"window_seconds": 300.0},
826+ )
827+ manager = shape_handling.adaptive_manager
828+ snapshot = manager.get_snapshot()
829+ fake_key = 42 # opaque integer key from pool.register()
830+ manager.record_event(
831+ raw_gear_values=[[20], [96]],
832+ mapped_gear_values=[[32], [128]],
833+ pad_ratios=[[0.375], [0.25]],
834+ split_ratios=[[0.0], [0.0]],
835+ event_ts=100.0,
836+ cleanup_key=fake_key,
837+ )
838+ 
839+ self.assertEqual(manager._states["BATCHSIZE:32"].cleanup_keys, {fake_key})
840+ self.assertEqual(manager._states["SEQLEN:128"].cleanup_keys, {fake_key})
841+ self.assertIsInstance(next(iter(manager._states["BATCHSIZE:32"].cleanup_keys)), int)
842+ 
843+ def test_npu_shape_handling_commit_update_uses_recorded_cleanup_keys(self):
844+ shape_handling = torch_npu._inductor.NPUShapeHandling(
845+ configs=[
846+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
847+ ],
848+ adaptive_configs={"window_seconds": 300.0},
849+ )
850+ manager = shape_handling.adaptive_manager
851+ snapshot = manager.get_snapshot()
852+ fake_key = 42 # opaque integer key
853+ manager.record_event(
854+ raw_gear_values=[[20]],
855+ mapped_gear_values=[[32]],
856+ pad_ratios=[[0.375]],
857+ split_ratios=[[0.0]],
858+ event_ts=100.0,
859+ cleanup_key=fake_key,
860+ )
861+ 
862+ _, removed_keys = manager.commit_update(
863+ configs=[{"type": "BATCHSIZE", "gears": [16], "dimensions": 0, "indices": [0], "policy": "CUSTOM"}],
864+ removed_gears=["BATCHSIZE:32"],
865+ now_ts=101.0,
866+ )
867+ 
868+ self.assertEqual(removed_keys, [fake_key])
869+ self.assertIsInstance(removed_keys[0], int)
870+ self.assertNotIn("BATCHSIZE:32", manager._states)
871+ 
872+ def test_npu_shape_handling_builds_separate_pad_and_split_sample_counts(self):
873+ shape_handling = torch_npu._inductor.NPUShapeHandling(
874+ configs=[
875+ {"type": "BATCHSIZE", "gears": [16, 32, 64], "dimensions": 0, "indices": [0]},
876+ ],
877+ adaptive_configs={"window_seconds": 300.0},
878+ )
879+ manager = shape_handling.adaptive_manager
880+ snapshot = manager.get_snapshot()
881+ manager.record_event([[20]], [[32]], [[0.375]], [[0.0]], 100.0)
882+ manager.record_event([[80]], [[64]], [[0.0]], [[0.20]], 101.0)
883+ 
884+ stats = manager.build_stats_snapshot(101.0)
885+ self.assertEqual(stats["BATCHSIZE:32"]["pad_sample_count"], 1)
886+ self.assertEqual(stats["BATCHSIZE:32"]["split_sample_count"], 0)
887+ self.assertEqual(stats["BATCHSIZE:64"]["pad_sample_count"], 0)
888+ self.assertEqual(stats["BATCHSIZE:64"]["split_sample_count"], 1)
889+ 
890+ def test_npu_shape_handling_builds_two_dim_stats_per_dimension(self):
891+ shape_handling = torch_npu._inductor.NPUShapeHandling(
892+ configs=[
893+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
894+ {"type": "SEQLEN", "gears": [64, 128], "dimensions": [1], "indices": [0]},
895+ ],
896+ adaptive_configs={"window_seconds": 300.0},
897+ )
898+ manager = shape_handling.adaptive_manager
899+ snapshot = manager.get_snapshot()
900+ manager.record_event([[20], [96]], [[32], [128]], [[0.375], [0.25]], [[0.0], [0.0]], 100.0)
901+ manager.record_event([[24], [112]], [[32], [128]], [[0.25], [0.125]], [[0.0], [0.0]], 101.0)
902+ 
903+ stats = manager.build_stats_snapshot(101.0)
904+ self.assertEqual(stats["BATCHSIZE:32"]["sample_count"], 2)
905+ self.assertEqual(stats["BATCHSIZE:32"]["raw_samples"], [20, 24])
906+ self.assertEqual(stats["SEQLEN:128"]["sample_count"], 2)
907+ self.assertEqual(stats["SEQLEN:128"]["raw_samples"], [96, 112])
908+ 
909+ def test_npu_shape_handling_recent_use_protect_skips_recent_hit_gear(self):
910+ shape_handling = torch_npu._inductor.NPUShapeHandling(
911+ configs=[
912+ {"type": "BATCHSIZE", "gears": [8, 16, 32], "dimensions": 0, "indices": [0]},
913+ ],
914+ adaptive_configs={
915+ "recent_use_protect_seconds": 60.0,
916+ "min_gear_count_per_type": 1,
917+ },
918+ )
919+ manager = shape_handling.adaptive_manager
920+ snapshot = manager.get_snapshot()
921+ manager.record_event([[20]], [[16]], [[0.375]], [[0.0]], 100.0)
922+ 
923+ stats = manager.build_stats_snapshot(120.0)
924+ breakdowns = manager.scorer.build_score_breakdown(manager.get_snapshot(), stats)
925+ candidates = manager.worker.build_eviction_candidates(
926+ breakdowns,
927+ manager.get_snapshot(),
928+ stats,
929+ 120.0,
930+ )
931+ self.assertEqual(candidates, {"BATCHSIZE": "BATCHSIZE:8"})
932+ 
933+ def test_npu_shape_handling_eviction_protection_skips_gear(self):
934+ shape_handling = torch_npu._inductor.NPUShapeHandling(
935+ configs=[
936+ {"type": "BATCHSIZE", "gears": [8, 16, 32], "dimensions": 0, "indices": [0]},
937+ ],
938+ adaptive_configs={
939+ "recent_use_protect_seconds": 300.0,
940+ "min_gear_count_per_type": 1,
941+ },
942+ )
943+ manager = shape_handling.adaptive_manager
944+ snapshot = manager.get_snapshot()
945+ manager.record_event([[20]], [[32]], [[0.375]], [[0.0]], 0.0)
946+ 
947+ manager.protect_gear_from_eviction("BATCHSIZE:16", 500.0)
948+ 
949+ stats = manager.build_stats_snapshot(500.0)
950+ breakdowns = manager.scorer.build_score_breakdown(manager.get_snapshot(), stats)
951+ candidates = manager.worker.build_eviction_candidates(
952+ breakdowns,
953+ manager.get_snapshot(),
954+ stats,
955+ 500.0,
956+ )
957+ self.assertEqual(candidates, {"BATCHSIZE": "BATCHSIZE:8"})
958+ 
959+ def test_npu_shape_handling_zero_usage_gears_are_eviction_candidates(self):
960+ shape_handling = torch_npu._inductor.NPUShapeHandling(
961+ configs=[
962+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
963+ ],
964+ adaptive_configs={
965+ "recent_use_protect_seconds": 0.0,
966+ "min_gear_count_per_type": 1,
967+ },
968+ )
969+ manager = shape_handling.adaptive_manager
970+ snapshot = manager.get_snapshot()
971+ manager.record_event([[30]], [[32]], [[0.0625]], [[0.0]], 100.0)
972+ 
973+ stats = manager.build_stats_snapshot(120.0)
974+ breakdowns = manager.scorer.build_score_breakdown(manager.get_snapshot(), stats)
975+ candidates = manager.worker.build_eviction_candidates(
976+ breakdowns,
977+ manager.get_snapshot(),
978+ stats,
979+ 120.0,
980+ )
981+ # only unused gear is candidate
982+ self.assertEqual(candidates, {"BATCHSIZE": "BATCHSIZE:16"})
983+ 
984+ def test_npu_shape_handling_update_adds_new_gear(self):
985+ shape_handling = torch_npu._inductor.NPUShapeHandling(
986+ configs=[
987+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
988+ {"type": "SEQLEN", "gears": [64], "dimensions": [1], "indices": [0]},
989+ ],
990+ adaptive_configs={
991+ "recent_use_protect_seconds": 1.0,
992+ "min_samples_per_gear": 2,
993+ "add_min_samples": 2,
994+ "pad_add_threshold": 0.20,
995+ "recent_use_protect_seconds": 0.0,
996+ },
997+ )
998+ manager = shape_handling.adaptive_manager
999+ snapshot = manager.get_snapshot()
1000+ manager.record_event([[20], [96]], [[32], [64]], [[0.375], [0.50]], [[0.0], [0.0]], 100.0)
1001+ manager.record_event([[20], [96]], [[32], [64]], [[0.375], [0.50]], [[0.0], [0.0]], 101.0)
1002+ 
1003+ manager.worker.run_once(150.0)
1004+ latest_snapshot = manager.get_snapshot()
1005+ # pad-driven gear added
1006+ self.assertIn(96, latest_snapshot.active_gears["SEQLEN"])
1007+ 
1008+ def test_npu_shape_handling_new_zero_usage_gear_respects_recent_create_protection(self):
1009+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1010+ configs=[
1011+ {"type": "BATCHSIZE", "gears": [16, 32, 64], "min_size": 1, "dimensions": 0, "indices": [0]},
1012+ ],
1013+ adaptive_configs={
1014+ "recent_use_protect_seconds": 1.0,
1015+ "add_min_samples": 2,
1016+ "pad_add_threshold": 0.20,
1017+ "recent_use_protect_seconds": 60.0,
1018+ "min_gear_count_per_type": 1,
1019+ },
1020+ )
1021+ manager = shape_handling.adaptive_manager
1022+ snapshot = manager.get_snapshot()
1023+ manager.record_event([[20]], [[32]], [[0.375]], [[0.0]], 10.0)
1024+ manager.record_event([[20]], [[32]], [[0.375]], [[0.0]], 11.0)
1025+ 
1026+ manager.worker.run_once(100.0)
1027+ 
1028+ latest_snapshot = manager.get_snapshot()
1029+ # gear 20 was added
1030+ self.assertIn(20, latest_snapshot.active_gears["BATCHSIZE"])
1031+ 
1032+ stats = manager.build_stats_snapshot(120.0)
1033+ breakdowns = manager.scorer.build_score_breakdown(latest_snapshot, stats)
1034+ candidates = manager.worker.build_eviction_candidates(
1035+ breakdowns, latest_snapshot, stats, 120.0,
1036+ )
1037+ # new gear has zero hits
1038+ self.assertEqual(stats["BATCHSIZE:20"]["sample_count"], 0)
1039+ # creation timestamp preserved
1040+ self.assertEqual(stats["BATCHSIZE:20"]["created_ts"], 100.0)
1041+ # gear 32 evictable, gear 20 protected
1042+ self.assertEqual(candidates, {"BATCHSIZE": "BATCHSIZE:32"})
1043+ 
1044+ def test_npu_shape_handling_update_adds_split_driven_gear(self):
1045+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1046+ configs=[
1047+ {"type": "BATCHSIZE", "gears": [64], "dimensions": 0, "indices": [0]},
1048+ ],
1049+ adaptive_configs={
1050+ "recent_use_protect_seconds": 1.0,
1051+ "add_min_samples": 2,
1052+ "pad_add_threshold": 0.90,
1053+ "split_add_threshold": 0.10,
1054+ "recent_use_protect_seconds": 0.0,
1055+ },
1056+ )
1057+ manager = shape_handling.adaptive_manager
1058+ snapshot = manager.get_snapshot()
1059+ manager.record_event([[96]], [[64]], [[0.0]], [[0.34]], 100.0)
1060+ manager.record_event([[96]], [[64]], [[0.0]], [[0.34]], 101.0)
1061+ 
1062+ manager.worker.run_once(150.0)
1063+ latest_snapshot = manager.get_snapshot()
1064+ # split-driven gear added at median
1065+ self.assertIn(96, latest_snapshot.active_gears["BATCHSIZE"])
1066+ 
1067+ def test_npu_shape_handling_two_dim_addition_candidates_are_built_per_dimension(self):
1068+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1069+ configs=[
1070+ {"type": "BATCHSIZE", "gears": [8], "dimensions": 0, "indices": [0]},
1071+ {"type": "SEQLEN", "gears": [32], "dimensions": [1], "indices": [0]},
1072+ ],
1073+ adaptive_configs={
1074+ "recent_use_protect_seconds": 1.0,
1075+ "add_min_samples": 2,
1076+ "split_add_threshold": 0.10,
1077+ "recent_use_protect_seconds": 0.0,
1078+ },
1079+ )
1080+ manager = shape_handling.adaptive_manager
1081+ snapshot = manager.get_snapshot()
1082+ manager.record_event([[16], [192]], [[8], [32]], [[0.0], [0.0]], [[0.50], [0.83]], 100.0)
1083+ manager.record_event([[16], [192]], [[8], [32]], [[0.0], [0.0]], [[0.50], [0.83]], 101.0)
1084+ manager.record_event([[48], [64]], [[8], [32]], [[0.0], [0.0]], [[0.83], [0.50]], 102.0)
1085+ manager.record_event([[48], [64]], [[8], [32]], [[0.0], [0.0]], [[0.83], [0.50]], 103.0)
1086+ 
1087+ stats = manager.build_stats_snapshot(103.0)
1088+ candidates = manager.worker.build_addition_candidates(
1089+ manager.get_snapshot(),
1090+ stats,
1091+ )
1092+ result = {shape_type: value for _, _, shape_type, value in candidates}
1093+ self.assertEqual(result, {"BATCHSIZE": 32, "SEQLEN": 128})
1094+ 
1095+ def test_npu_shape_handling_update_evicts_low_value_gear(self):
1096+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1097+ configs=[
1098+ {"type": "BATCHSIZE", "gears": [8, 16, 32], "dimensions": 0, "indices": [0]},
1099+ {"type": "SEQLEN", "gears": [64], "dimensions": [1], "indices": [0]},
1100+ ],
1101+ adaptive_configs={
1102+ "recent_use_protect_seconds": 1.0,
1103+ "min_samples_per_gear": 1,
1104+ "min_gear_count_per_type": 2,
1105+ "recent_use_protect_seconds": 0.0,
1106+ "replace_loss_threshold": 1.0,
1107+ },
1108+ )
1109+ manager = shape_handling.adaptive_manager
1110+ snapshot = manager.get_snapshot()
1111+ manager.record_event([[8], [64]], [[8], [64]], [[0.0], [0.0]], [[0.0], [0.0]], 10.0)
1112+ manager.record_event([[8], [64]], [[8], [64]], [[0.0], [0.0]], [[0.0], [0.0]], 11.0)
1113+ manager.record_event([[10], [64]], [[16], [64]], [[0.375], [0.0]], [[0.0], [0.0]], 12.0)
1114+ manager.record_event([[32], [64]], [[32], [64]], [[0.0], [0.0]], [[0.0], [0.0]], 13.0)
1115+ 
1116+ manager.worker.run_once(20.0)
1117+ latest_snapshot = manager.get_snapshot()
1118+ self.assertNotIn(16, latest_snapshot.active_gears["BATCHSIZE"])
1119+ self.assertIn(8, latest_snapshot.active_gears["BATCHSIZE"])
1120+ self.assertIn(32, latest_snapshot.active_gears["BATCHSIZE"])
1121+ 
1122+ def test_npu_shape_handling_two_dim_evicts_per_dimension(self):
1123+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1124+ configs=[
1125+ {"type": "BATCHSIZE", "gears": [8, 16, 32], "dimensions": 0, "indices": [0]},
1126+ {"type": "SEQLEN", "gears": [32, 64, 128], "dimensions": [1], "indices": [0]},
1127+ ],
1128+ adaptive_configs={
1129+ "recent_use_protect_seconds": 1.0,
1130+ "min_samples_per_gear": 1,
1131+ "min_gear_count_per_type": 2,
1132+ "add_min_samples": 10,
1133+ "recent_use_protect_seconds": 0.0,
1134+ "replace_loss_threshold": 1.0,
1135+ },
1136+ )
1137+ manager = shape_handling.adaptive_manager
1138+ snapshot = manager.get_snapshot()
1139+ manager.record_event([[8], [64]], [[8], [64]], [[0.0], [0.0]], [[0.0], [0.0]], 18.0)
1140+ manager.record_event([[10], [96]], [[16], [128]], [[0.375], [0.25]], [[0.0], [0.0]], 19.0)
1141+ manager.record_event([[32], [64]], [[32], [64]], [[0.0], [0.0]], [[0.0], [0.0]], 19.0)
1142+ 
1143+ manager.worker.run_once(20.0)
1144+ latest_snapshot = manager.get_snapshot()
1145+ self.assertEqual(latest_snapshot.active_gears["BATCHSIZE"], [8, 32])
1146+ self.assertEqual(latest_snapshot.active_gears["SEQLEN"], [64, 128])
1147+ 
1148+ def test_npu_shape_handling_resource_pressure_triggers_update(self):
1149+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1150+ configs=[
1151+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
1152+ ],
1153+ adaptive_configs={
1154+ "recent_use_protect_seconds": 300.0,
1155+ "device_memory_usage_threshold_ratio": 0.80,
1156+ },
1157+ )
1158+ manager = shape_handling.adaptive_manager
1159+ # Memory usage ratio = 1.0 - 20/100 = 0.80 >= threshold → high pressure
1160+ with mock.patch("torch_npu._inductor.adaptive_gears.torch.npu.mem_get_info", return_value=(20, 100)):
1161+ budget = manager.build_resource_budget()
1162+ self.assertTrue(budget["device_memory_usage_high"])
1163+ # Memory usage ratio = 1.0 - 70/100 = 0.30 < threshold → no pressure
1164+ with mock.patch("torch_npu._inductor.adaptive_gears.torch.npu.mem_get_info", return_value=(70, 100)):
1165+ budget = manager.build_resource_budget()
1166+ self.assertFalse(budget["device_memory_usage_high"])
1167+ 
1168+ def test_npu_shape_handling_resource_budget_uses_device_ratio(self):
1169+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1170+ configs=[
1171+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0]},
1172+ ],
1173+ adaptive_configs={
1174+ "device_memory_usage_threshold_ratio": 0.60,
1175+ },
1176+ )
1177+ manager = shape_handling.adaptive_manager
1178+ with mock.patch("torch_npu._inductor.adaptive_gears.torch.npu.mem_get_info", return_value=(40, 100)):
1179+ budget = manager.build_resource_budget()
1180+ 
1181+ self.assertAlmostEqual(budget["device_memory_usage_ratio"], 0.60)
1182+ self.assertEqual(budget["device_memory_usage_threshold_ratio"], 0.60)
1183+ self.assertTrue(budget["device_memory_usage_high"])
1184+ 
1185+ def test_npu_shape_handling_device_pressure_blocks_direct_add(self):
1186+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1187+ configs=[
1188+ {"type": "BATCHSIZE", "gears": [64], "dimensions": 0, "indices": [0]},
1189+ ],
1190+ adaptive_configs={
1191+ "recent_use_protect_seconds": 1.0,
1192+ "add_min_samples": 2,
1193+ "pad_add_threshold": 0.20,
1194+ "device_memory_usage_threshold_ratio": 0.50,
1195+ "min_gear_count_per_type": 1,
1196+ "recent_use_protect_seconds": 0.0,
1197+ },
1198+ )
1199+ manager = shape_handling.adaptive_manager
1200+ snapshot = manager.get_snapshot()
1201+ manager.record_event([[40]], [[64]], [[0.375]], [[0.0]], 100.0)
1202+ manager.record_event([[40]], [[64]], [[0.375]], [[0.0]], 101.0)
1203+ 
1204+ with mock.patch("torch_npu._inductor.adaptive_gears.torch.npu.mem_get_info", return_value=(40, 100)):
1205+ manager.worker.run_once(150.0)
1206+ latest_snapshot = manager.get_snapshot()
1207+ self.assertEqual(latest_snapshot.active_gears["BATCHSIZE"], [64])
1208+ 
1209+ def test_npu_shape_handling_max_gear_guard_blocks_unsafe_delete(self):
1210+ shape_handling = torch_npu._inductor.NPUShapeHandling(
1211+ configs=[
1212+ {"type": "BATCHSIZE", "gears": [32, 64], "dimensions": 0, "indices": [0]},
1213+ ],
1214+ adaptive_configs={
1215+ "recent_use_protect_seconds": 0.0,
1216+ "replace_loss_threshold": 1.0,
1217+ "min_gear_count_per_type": 1,
1218+ },
1219+ )
1220+ manager = shape_handling.adaptive_manager
1221+ snapshot = manager.get_snapshot()
1222+ manager.record_event([[40]], [[64]], [[0.375]], [[0.0]], 100.0)
1223+ 
1224+ stats = manager.build_stats_snapshot(100.0)
1225+ breakdowns = manager.scorer.build_score_breakdown(manager.get_snapshot(), stats)
1226+ candidates = manager.worker.build_eviction_candidates(
1227+ breakdowns,
1228+ manager.get_snapshot(),
1229+ stats,
1230+ 100.0,
1231+ )
1232+ self.assertNotIn("BATCHSIZE:64", candidates.values())
523 1233 
524class TestDynamicShapeCompile(TestCase):1234class TestDynamicShapeCompile(TestCase):
525 def test_npu_dynamic_shape_reuse_with_no_bucket(self):1235 def test_npu_dynamic_shape_reuse_with_no_bucket(self):
@@ -807,10 +1517,322 @@ class TestUnifiedCopy(TestCase):
807 self.assertIs(copied, obj)1517 self.assertIs(copied, obj)
808 1518 
809 1519 
1520+# ---------------------------------------------------------------------------
1521+# ST: Adaptive Gear end-to-end integration tests
1522+# ---------------------------------------------------------------------------
1523+ 
1524+_ST_ADAPTIVE_CONFIGS = {
1525+ "window_seconds": 300.0,
1526+ "recent_use_protect_seconds": 0,
1527+ "pad_add_threshold": 0.20,
1528+ "split_add_threshold": 0.20,
1529+ "add_min_samples": 2,
1530+ "min_samples_per_gear": 1,
1531+ "min_gear_count_per_type": 2,
1532+ "replace_loss_threshold": 1.0,
1533+ "update_interval_seconds": 9999.0,
luqichao
luqichaoluqichao6月8日

配置重复 同1533

likedislike
zhudada0120
zhudada0120
6月15日 评论：
1534+}
1535+ 
1536+_ST_DAEMON_ADAPTIVE_CONFIGS = dict(_ST_ADAPTIVE_CONFIGS)
1537+_ST_DAEMON_ADAPTIVE_CONFIGS["update_interval_seconds"] = 0.1
1538+ 
1539+ 
1540+def _make_adaptive_options(shape_configs, adaptive_configs=None):
1541+ opts = {
1542+ "enable_shape_handling": True,
1543+ "shape_handling_configs": shape_configs,
1544+ }
1545+ if adaptive_configs is not None:
1546+ opts["shape_handling_dict"] = {"adaptive_gears": adaptive_configs}
1547+ return opts
1548+ 
1549+ 
1550+def _run_model(compiled_fn, shape):
1551+ A = torch.randn(shape, device=device)
1552+ B = torch.randn(shape, device=device)
1553+ out = compiled_fn(A, B)
1554+ return out, A, B
1555+ 
1556+ 
1557+class TestAdaptiveGearsCompileST(TestCase):
1558+ 
1559+ def setUp(self):
1560+ torch._dynamo.reset()
1561+ self._captured_managers = []
1562+ self._original_sh_init = shape_handling_module.NPUShapeHandling.__init__
1563+ 
1564+ test_self = self
1565+ 
1566+ def capturing_init(sh_self, *args, **kwargs):
1567+ test_self._original_sh_init(sh_self, *args, **kwargs)
1568+ if sh_self.adaptive_manager is not None:
1569+ test_self._captured_managers.append(sh_self.adaptive_manager)
1570+ 
1571+ shape_handling_module.NPUShapeHandling.__init__ = capturing_init
1572+ 
1573+ def tearDown(self):
1574+ shape_handling_module.NPUShapeHandling.__init__ = self._original_sh_init
1575+ for mgr in self._captured_managers:
1576+ try:
1577+ mgr.shutdown()
1578+ except Exception:
1579+ pass
1580+ self._captured_managers.clear()
1581+ torch._dynamo.reset()
1582+ 
1583+ def _get_manager(self):
1584+ self.assertTrue(
1585+ len(self._captured_managers) > 0,
1586+ "No AdaptiveGearRuntime was captured — torch.compile did not create one",
1587+ )
1588+ return self._captured_managers[-1]
1589+ 
1590+ def _stop_daemon(self, manager):
1591+ manager._shutdown_event.set()
1592+ manager._worker_thread.join(timeout=2.0)
1593+ 
1594+ def test_adaptive_gears_basic_compile_and_compute(self):
1595+ shape_configs = [
1596+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0, 1]},
1597+ ]
1598+ options = _make_adaptive_options(shape_configs, _ST_ADAPTIVE_CONFIGS)
1599+ 
1600+ compiled_fn = torch.compile(
1601+ model_fn,
1602+ backend="inductor",
1603+ dynamic=False,
1604+ options=options,
1605+ )
1606+ 
1607+ for shape in [(8, 32), (24, 32), (10, 32)]:
1608+ out, A, B = _run_model(compiled_fn, shape)
1609+ self.assertTrue(
1610+ torch.allclose(out, A + B),
1611+ f"Output mismatch for shape {shape}",
1612+ )
1613+ 
1614+ manager = self._get_manager()
1615+ 
1616+ self.assertTrue(
1617+ manager._worker_thread.is_alive(),
1618+ "Daemon worker thread should be alive",
1619+ )
1620+ 
1621+ def test_adaptive_gears_event_recording_and_stats(self):
1622+ shape_configs = [
1623+ {"type": "BATCHSIZE", "gears": [16, 32], "dimensions": 0, "indices": [0, 1]},
1624+ ]
1625+ options = _make_adaptive_options(shape_configs, _ST_ADAPTIVE_CONFIGS)
1626+ 
1627+ compiled_fn = torch.compile(
1628+ model_fn,
1629+ backend="inductor",
1630+ dynamic=False,
1631+ options=options,
1632+ )
1633+ manager = self._get_manager()
1634+ 
1635+ for _ in range(3):
1636+ out, A, B = _run_model(compiled_fn, (24, 32))
1637+ self.assertTrue(torch.allclose(out, A + B))
1638+ 
1639+ stats = manager.build_stats_snapshot(time.time())
1640+ 
1641+ self.assertIn("BATCHSIZE:32", stats)
1642+ self.assertEqual(
1643+ stats["BATCHSIZE:32"]["sample_count"], 6,
1644+ "Exactly 6 samples should be recorded for gear 32 (3 calls × 2 tensors)",
1645+ )
1646+ self.assertEqual(
1647+ stats["BATCHSIZE:32"]["pad_sample_count"], 6,
1648+ "All 6 events should be classified as padding",
1649+ )
1650+ self.assertAlmostEqual(
1651+ stats["BATCHSIZE:32"]["avg_pad_ratio"], 0.25, places=2,
1652+ msg="avg_pad_ratio should be (32-24)/32 = 0.25",
1653+ )
1654+ self.assertEqual(
1655+ stats["BATCHSIZE:32"]["raw_samples"], [24, 24, 24, 24, 24, 24],
1656+ "raw_samples should have 6 entries (3 calls × 2 tensors)",
1657+ )
1658+ 
1659+ def test_adaptive_gears_gear_eviction_synchronous(self):
1660+ shape_configs = [
1661+ {"type": "BATCHSIZE", "gears": [16, 32, 64], "dimensions": 0, "indices": [0, 1]},
1662+ ]
1663+ options = _make_adaptive_options(shape_configs, _ST_ADAPTIVE_CONFIGS)
1664+ 
1665+ compiled_fn = torch.compile(
1666+ model_fn,
1667+ backend="inductor",
1668+ dynamic=False,
1669+ options=options,
1670+ )
1671+ manager = self._get_manager()
1672+ self._stop_daemon(manager)
1673+ 
1674+ for _ in range(5):
1675+ out, A, B = _run_model(compiled_fn, (24, 32))
1676+ self.assertTrue(torch.allclose(out, A + B))
1677+ 
1678+ manager.worker.run_once(time.time())
1679+ 
1680+ snapshot_after = manager.get_snapshot()
1681+ gears_after = sorted(snapshot_after.active_gears["BATCHSIZE"])
1682+ 
1683+ # Gear 16 evicted (zero hits) + gear 24 added (pad_ratio=0.25 > threshold 0.20)
1684+ self.assertEqual(
1685+ gears_after, [24, 32, 64],
1686+ f"Expected [24, 32, 64] after eviction+addition, got {gears_after}",
1687+ )
1688+ 
1689+ out, A, B = _run_model(compiled_fn, (8, 32))
1690+ self.assertTrue(
1691+ torch.allclose(out, A + B),
1692+ "Computation should be correct after gear eviction",
1693+ )
1694+ 
1695+ def test_adaptive_gears_gear_addition_synchronous(self):
1696+ shape_configs = [
1697+ {"type": "BATCHSIZE", "gears": [64], "dimensions": 0, "indices": [0, 1]},
1698+ ]
1699+ add_configs = dict(_ST_ADAPTIVE_CONFIGS)
1700+ add_configs["min_gear_count_per_type"] = 1
1701+ options = _make_adaptive_options(shape_configs, add_configs)
1702+ 
1703+ compiled_fn = torch.compile(
1704+ model_fn,
1705+ backend="inductor",
1706+ dynamic=False,
1707+ options=options,
1708+ )
1709+ manager = self._get_manager()
1710+ self._stop_daemon(manager)
1711+ 
1712+ for _ in range(3):
1713+ out, A, B = _run_model(compiled_fn, (20, 32))
1714+ self.assertTrue(torch.allclose(out, A + B))
1715+ 
1716+ manager.worker.run_once(time.time())
1717+ 
1718+ snapshot_after = manager.get_snapshot()
1719+ gears_after = snapshot_after.active_gears["BATCHSIZE"]
1720+ 
1721+ self.assertEqual(
1722+ len(gears_after), 2,
1723+ f"A new gear should be added, got {gears_after}",
1724+ )
1725+ self.assertIn(
1726+ 20, gears_after,
1727+ f"New gear 20 (median of raw samples) should be added, got {gears_after}",
1728+ )
1729+ 
1730+ out, A, B = _run_model(compiled_fn, (20, 32))
1731+ self.assertTrue(
1732+ torch.allclose(out, A + B),
1733+ "Computation should be correct after gear addition",
1734+ )
1735+ 
1736+ def test_adaptive_gears_graph_cleanup_after_eviction(self):
1737+ from torch_npu.npu._graph_resource_pool import GraphResourcePool
1738+ 
1739+ shape_configs = [
1740+ {"type": "BATCHSIZE", "gears": [16, 32, 64], "dimensions": 0, "indices": [0, 1]},
1741+ ]
1742+ options = _make_adaptive_options(shape_configs, _ST_ADAPTIVE_CONFIGS)
1743+ options["triton.cudagraphs"] = True
1744+ options["triton.cudagraph_trees"] = True
1745+ 
1746+ compiled_fn = torch.compile(
1747+ model_fn,
1748+ backend="inductor",
1749+ dynamic=False,
1750+ options=options,
1751+ )
1752+ manager = self._get_manager()
1753+ self._stop_daemon(manager)
1754+ 
1755+ for shape in [(8, 32), (24, 32), (48, 32)]:
1756+ for _ in range(2):
1757+ out, A, B = _run_model(compiled_fn, shape)
1758+ self.assertTrue(torch.allclose(out, A + B))
1759+ 
1760+ device_index = torch.npu.current_device()
1761+ pool = GraphResourcePool.get_pool(device_index)
1762+ 
1763+ self.assertGreater(
1764+ pool.entry_count, 0,
1765+ f"Pool should have entries after compilation, got {pool.entry_count}",
1766+ )
1767+ 
1768+ manager.worker.run_once(time.time())
1769+ 
1770+ snapshot_after = manager.get_snapshot()
1771+ gears_after = sorted(snapshot_after.active_gears["BATCHSIZE"])
1772+ 
1773+ self.assertEqual(
1774+ gears_after, [8, 24, 32, 48, 64],
1775+ f"Expected [8, 24, 32, 48, 64] after eviction+additions, got {gears_after}",
1776+ )
1777+ 
1778+ out, A, B = _run_model(compiled_fn, (24, 32))
1779+ self.assertTrue(
1780+ torch.allclose(out, A + B),
1781+ "Computation should be correct after graph cleanup",
1782+ )
1783+ 
1784+ def test_adaptive_gears_daemon_thread_triggers_update(self):
1785+ shape_configs = [
1786+ {"type": "BATCHSIZE", "gears": [16, 32, 64], "dimensions": 0, "indices": [0, 1]},
1787+ ]
1788+ options = _make_adaptive_options(shape_configs, _ST_DAEMON_ADAPTIVE_CONFIGS)
1789+ 
1790+ compiled_fn = torch.compile(
1791+ model_fn,
1792+ backend="inductor",
1793+ dynamic=False,
1794+ options=options,
1795+ )
1796+ manager = self._get_manager()
1797+ original_gears = sorted(manager.get_snapshot().active_gears["BATCHSIZE"])
1798+ 
1799+ for _ in range(5):
1800+ out, A, B = _run_model(compiled_fn, (24, 32))
1801+ self.assertTrue(torch.allclose(out, A + B))
1802+ 
1803+ self.assertTrue(
1804+ manager._worker_thread.is_alive(),
1805+ "Daemon thread should be alive",
1806+ )
1807+ 
1808+ deadline = time.time() + 5.0
1809+ current_gears = original_gears
1810+ while time.time() < deadline:
1811+ current_gears = sorted(manager.get_snapshot().active_gears["BATCHSIZE"])
1812+ if current_gears != original_gears:
1813+ break
1814+ time.sleep(0.1)
1815+ 
1816+ self.assertNotEqual(
1817+ current_gears, original_gears,
1818+ f"Daemon thread should trigger gear update within 5s. "
1819+ f"Original: {original_gears}, Current: {current_gears}",
1820+ )
1821+ 
1822+ out, A, B = _run_model(compiled_fn, (24, 32))
1823+ self.assertTrue(
1824+ torch.allclose(out, A + B),
1825+ "Computation should be correct after daemon-triggered update",
1826+ )
1827+ 
1828+ 
810instantiate_parametrized_tests(TestShapeHandling)1829instantiate_parametrized_tests(TestShapeHandling)
811instantiate_parametrized_tests(TestShapeHandlingBranchCoverage)1830instantiate_parametrized_tests(TestShapeHandlingBranchCoverage)
1831+instantiate_parametrized_tests(TestAsyncWorkerAndConcurrency)
1832+instantiate_parametrized_tests(TestAdaptiveShapeHandling)
812instantiate_parametrized_tests(TestUnifiedCopy)1833instantiate_parametrized_tests(TestUnifiedCopy)
813instantiate_parametrized_tests(TestDynamicShapeCompile)1834instantiate_parametrized_tests(TestDynamicShapeCompile)
1835+instantiate_parametrized_tests(TestAdaptiveGearsCompileST)
814 1836
815if __name__ == '__main__':1837if __name__ == '__main__':
816 run_tests()1838 run_tests()
Mtest/npu/test_graph_tree.py+319-1
@@ -1,4 +1,5 @@
1import os1import os
2+import threading
2 3 
3os.environ["ASCEND_LAUNCH_BLOCKING"] = "0"4os.environ["ASCEND_LAUNCH_BLOCKING"] = "0"
4 5 
@@ -7,6 +8,7 @@ import weakref
7import pytest8import pytest
8import torch9import torch
9import torch_npu10import torch_npu
11+from torch_npu.npu._graph_resource_pool import GraphResourcePool
10from torch_npu.npu._graph_tree import (12from torch_npu.npu._graph_tree import (
11 check_memory_pool,13 check_memory_pool,
12 clear_cublass_cache,14 clear_cublass_cache,
@@ -172,7 +174,8 @@ class TestTreeManagerContainer(TestCase):
172 ):174 ):
173 manager = self.container.get_tree_manager()175 manager = self.container.get_tree_manager()
174 self.assertIsNotNone(manager)176 self.assertIsNotNone(manager)
175- self.assertIs(manager, self.container.get_tree_manager()) # Same instance177+ # Same instance
178+ self.assertIs(manager, self.container.get_tree_manager())
176 179 
177 180 
178class TestStorageWeakRefWrapper(TestCase):181class TestStorageWeakRefWrapper(TestCase):
@@ -1205,6 +1208,321 @@ class TestNPUGraphTreeManager:
1205 )1208 )
1206 assert FunctionID(1) in manager.warned_functions1209 assert FunctionID(1) in manager.warned_functions
rain-666
rain-666rain-6666月1日

assert需要补充注释

likedislike
zhudada0120
zhudada0120
6月5日 评论：
1207 1210 
1211+ @patch('torch.npu.synchronize')
rain-666
rain-666rain-6666月1日

缺少边缘情况测试（如空输入、异常形状

likedislike
zhudada0120
zhudada0120
6月5日 评论：
1212+ @patch('torch_npu.npu._graph_tree.NPUGraphNode')
1213+ def test_record_function_pool_registration(self, mock_node, mock_synchronize):
1214+ """Both root and child nodes are registered in GraphResourcePool."""
1215+ GraphResourcePool.reset_all()
1216+ manager = NPUGraphTreeManager(0)
1217+ manager.ids_to_funcs[FunctionID(1)] = MagicMock()
1218+ manager.ids_to_stack_traces[FunctionID(1)] = "stack_trace"
1219+ manager.npu_graphs_thread_pool = "pool_handle"
1220+ manager.device_index = 0
1221+ manager.stream = MagicMock()
1222+ mock_node_instance = MagicMock()
1223+ mock_node.return_value = mock_node_instance
1224+ mock_node_instance.run_first_inputs.return_value = [torch.tensor([1.0])]
1225+ 
1226+ t1 = torch.tensor([1.0])
1227+ t2 = torch.tensor([2.0, 3.0])
1228+ manager.record_function([t1, t2], FunctionID(1))
1229+ 
1230+ # Root node: registered in the pool with an opaque integer key.
1231+ pool = GraphResourcePool.get_pool(0)
1232+ assert pool.entry_count == 1
1233+ root_keys = pool.consume_recent_keys()
1234+ assert len(root_keys) == 1
1235+ root_key = root_keys[0]
1236+ assert isinstance(root_key, int)
1237+ assert pool._entries[root_key] is mock_node_instance
1238+ # consume must drain the per-thread pending buffer.
1239+ assert pool.consume_recent_keys() == []
1240+ 
1241+ # Child node: also registered with its own key.
1242+ parent_node = MagicMock()
1243+ manager.current_node = parent_node
1244+ t3 = torch.tensor([4.0, 5.0, 6.0])
1245+ child_instance = MagicMock()
1246+ mock_node.return_value = child_instance
1247+ child_instance.run_first_inputs.return_value = [torch.tensor([1.0])]
1248+ manager.record_function([t3], FunctionID(1))
1249+ 
1250+ assert pool.entry_count == 2
1251+ child_keys = pool.consume_recent_keys()
1252+ assert len(child_keys) == 1
1253+ assert pool._entries[child_keys[0]] is child_instance
1254+ # consume must drain the buffer again.
1255+ assert pool.consume_recent_keys() == []
1256+ GraphResourcePool.reset_all()
1257+ 
1258+ 
1259+class TestGraphResourcePool:
1260+ """Tests for opaque-key registration / removal in GraphResourcePool."""
1261+ 
1262+ def setup_method(self):
1263+ GraphResourcePool.reset_all()
1264+ 
1265+ def teardown_method(self):
1266+ GraphResourcePool.reset_all()
1267+ 
1268+ # helpers ----------------------------------------------------------
1269+ 
1270+ @staticmethod
1271+ def _register_resource(pool, resource=None):
1272+ """Activate the pool, register *resource* (or a fresh MagicMock),
1273+ and return (key, resource)."""
1274+ pool.activate()
1275+ if resource is None:
1276+ resource = MagicMock()
1277+ key = pool.register(resource)
1278+ assert isinstance(key, int)
1279+ return key, resource
1280+ 
1281+ # tests ------------------------------------------------------------
1282+ 
1283+ def test_activate_deactivate(self):
1284+ """Default False, activate → True, deactivate → False."""
1285+ pool = GraphResourcePool.get_pool(0)
1286+ assert not pool.is_active()
1287+ pool.activate()
1288+ assert pool.is_active()
1289+ pool.deactivate()
1290+ assert not pool.is_active()
1291+ 
1292+ def test_activate_refcount(self):
1293+ """Nested activate/deactivate — outer activation survives inner
1294+ deactivate."""
1295+ pool = GraphResourcePool.get_pool(0)
1296+ pool.activate() # outer
1297+ pool.activate() # inner
1298+ pool.deactivate() # close inner — outer still active
1299+ assert pool.is_active()
1300+ key = pool.register(MagicMock())
1301+ assert key >= 0
1302+ assert pool.entry_count == 1
1303+ pool.deactivate() # close outer — now disabled
1304+ assert not pool.is_active()
1305+ 
1306+ @patch('torch.npu.synchronize')
1307+ def test_register_returns_int_key(self, _mock_sync):
1308+ """register() returns an opaque integer key when activated."""
1309+ pool = GraphResourcePool.get_pool(0)
1310+ key, _ = self._register_resource(pool)
1311+ assert isinstance(key, int)
1312+ assert pool.entry_count == 1
1313+ 
1314+ # -- consume_recent_keys -----------------------------------------------
1315+ 
1316+ def test_consume_recent_keys_returns_registered_keys(self):
1317+ """consume_recent_keys returns all keys registered since last consume."""
1318+ pool = GraphResourcePool.get_pool(0)
1319+ k1, _ = self._register_resource(pool)
1320+ k2, _ = self._register_resource(pool)
1321+ assert pool.entry_count == 2
1322+ keys = pool.consume_recent_keys()
1323+ assert keys == [k1, k2]
1324+ 
1325+ def test_consume_recent_keys_drains_buffer(self):
1326+ """A second consume_recent_keys returns [] (buffer was drained)."""
1327+ pool = GraphResourcePool.get_pool(0)
1328+ self._register_resource(pool)
1329+ self._register_resource(pool)
1330+ first = pool.consume_recent_keys()
1331+ assert len(first) == 2
1332+ # Buffer must be empty after consume.
1333+ assert pool.consume_recent_keys() == []
1334+ 
1335+ def test_consume_recent_keys_empty_when_no_registration(self):
1336+ """consume_recent_keys returns [] when nothing was registered."""
1337+ pool = GraphResourcePool.get_pool(0)
1338+ assert pool.consume_recent_keys() == []
1339+ 
1340+ def test_consume_recent_keys_thread_isolation(self):
1341+ """Each thread sees only its own registrations."""
1342+ pool = GraphResourcePool.get_pool(0)
1343+ results = {}
1344+ 
1345+ def register_and_consume(thread_id):
1346+ r1 = MagicMock()
1347+ r2 = MagicMock()
1348+ k1 = pool.register(r1)
1349+ k2 = pool.register(r2)
1350+ keys = pool.consume_recent_keys()
1351+ results[thread_id] = (k1, k2, keys)
1352+ 
1353+ t_a = threading.Thread(target=register_and_consume, args=(0,))
1354+ t_b = threading.Thread(target=register_and_consume, args=(1,))
1355+ t_a.start()
1356+ t_b.start()
1357+ t_a.join()
1358+ t_b.join()
1359+ 
1360+ # Each thread must receive exactly its own 2 keys.
1361+ k1_a, k2_a, keys_a = results[0]
1362+ k1_b, k2_b, keys_b = results[1]
1363+ assert keys_a == [k1_a, k2_a]
1364+ assert keys_b == [k1_b, k2_b]
1365+ # Keys across threads are globally unique.
1366+ assert len({k1_a, k2_a, k1_b, k2_b}) == 4
1367+ assert pool.entry_count == 4
1368+ 
1369+ # Consuming from the main thread must return nothing.
1370+ assert pool.consume_recent_keys() == []
1371+ 
1372+ # -- remove_by_keys ----------------------------------------------------
1373+ 
1374+ @patch('torch.npu.synchronize')
1375+ def test_register_and_remove_single(self, _mock_sync):
1376+ """Register one resource and remove it — basic happy path."""
1377+ pool = GraphResourcePool.get_pool(0)
1378+ key, resource = self._register_resource(pool)
1379+ assert pool.entry_count == 1
1380+ pool.remove_by_keys([key])
1381+ resource.release.assert_called_once()
1382+ assert pool.entry_count == 0
1383+ 
1384+ @patch('torch.npu.synchronize')
1385+ def test_remove_unknown_key_is_noop(self, _mock_sync):
1386+ """Removing a key that was never registered must not touch anything."""
1387+ pool = GraphResourcePool.get_pool(0)
1388+ key, resource = self._register_resource(pool)
1389+ # Use a key that is semantically impossible (negative).
1390+ pool.remove_by_keys([-1])
1391+ resource.release.assert_not_called()
1392+ assert pool.entry_count == 1
1393+ 
1394+ @patch('torch.npu.synchronize')
1395+ def test_multiple_resources_independent_keys(self, _mock_sync):
1396+ """Each register() returns a distinct key; removing one doesn't
1397+ affect the other."""
1398+ pool = GraphResourcePool.get_pool(0)
1399+ k1, r1 = self._register_resource(pool)
1400+ k2, r2 = self._register_resource(pool)
1401+ assert k1 != k2
1402+ assert pool.entry_count == 2
1403+ pool.remove_by_keys([k1])
1404+ r1.release.assert_called_once()
1405+ r2.release.assert_not_called()
1406+ assert pool.entry_count == 1
1407+ 
1408+ @patch('torch.npu.synchronize')
1409+ def test_pool_isolation_across_devices(self, _mock_sync):
1410+ """Per-device pools are independent."""
1411+ pool_0 = GraphResourcePool.get_pool(0)
1412+ pool_1 = GraphResourcePool.get_pool(1)
1413+ assert pool_0 is not pool_1
rain-666
rain-666rain-6666月1日

assert需要补充注释

likedislike
zhudada0120
zhudada0120
6月5日 评论：
1414+ k0, r0 = self._register_resource(pool_0)
1415+ k1, r1 = self._register_resource(pool_1)
1416+ pool_0.remove_by_keys([k0])
1417+ r0.release.assert_called_once()
1418+ r1.release.assert_not_called()
1419+ assert pool_1.entry_count == 1
1420+ 
1421+ @patch('torch.npu.synchronize')
1422+ def test_remove_empty_keys_is_noop(self, _mock_sync):
1423+ """Empty key list should not crash or release anything."""
1424+ pool = GraphResourcePool.get_pool(0)
1425+ key, resource = self._register_resource(pool)
1426+ pool.remove_by_keys([])
1427+ assert pool.entry_count == 1
1428+ resource.release.assert_not_called()
1429+ 
1430+ @patch('torch.npu.synchronize')
1431+ def test_remove_already_removed_key_is_noop(self, _mock_sync):
1432+ """Removing an already-removed key twice must not crash (idempotency)."""
1433+ pool = GraphResourcePool.get_pool(0)
1434+ key, resource = self._register_resource(pool)
1435+ pool.remove_by_keys([key])
1436+ assert pool.entry_count == 0
1437+ resource.release.assert_called_once()
1438+ pool.remove_by_keys([key]) # second removal — must not crash
1439+ assert pool.entry_count == 0
1440+ # release must NOT be called a second time.
1441+ resource.release.assert_called_once()
1442+ 
1443+ # -- lifecycle ---------------------------------------------------------
1444+ 
1445+ @patch('torch.npu.synchronize')
1446+ def test_register_consume_remove_lifecycle(self, _mock_sync):
1447+ """Full lifecycle: register → consume → remove, matching upper layer flow.
1448+ 
1449+ Simulates the new_fn loop in shape_handling.py where each src_fn
1450+ invocation produces N graph captures (N pool registrations) that are
1451+ then consumed, recorded as cleanup_keys on gears, and later removed
1452+ en masse on gear eviction.
1453+ """
1454+ pool = GraphResourcePool.get_pool(0)
1455+ 
1456+ # -- variant 0: 3 sub-graphs captured --
1457+ r_a, r_b, r_c = MagicMock(), MagicMock(), MagicMock()
1458+ k_a = pool.register(r_a)
1459+ k_b = pool.register(r_b)
1460+ k_c = pool.register(r_c)
1461+ assert pool.entry_count == 3
1462+ 
1463+ # Upper layer consumes keys for this variant.
1464+ keys_v0 = pool.consume_recent_keys()
1465+ assert keys_v0 == [k_a, k_b, k_c]
1466+ assert pool.consume_recent_keys() == [] # drained
1467+ 
1468+ # -- variant 1: 2 sub-graphs captured in a different shape --
1469+ r_d, r_e = MagicMock(), MagicMock()
1470+ k_d = pool.register(r_d)
1471+ k_e = pool.register(r_e)
1472+ assert pool.entry_count == 5
1473+ 
1474+ keys_v1 = pool.consume_recent_keys()
1475+ assert keys_v1 == [k_d, k_e]
1476+ 
1477+ # -- gear eviction removes variant 0 --
1478+ # All 3 keys for this variant are passed together (they may be
1479+ # distributed across several gears; e.g. BATCHSIZE:32 cleanup_keys
1480+ # also contains them).
1481+ pool.remove_by_keys(keys_v0)
1482+ r_a.release.assert_called_once()
1483+ r_b.release.assert_called_once()
1484+ r_c.release.assert_called_once()
1485+ # Variant 1 resources untouched.
1486+ assert pool.entry_count == 2
1487+ 
1488+ # -- gear eviction removes variant 1 later --
1489+ pool.remove_by_keys(keys_v1)
1490+ assert pool.entry_count == 0
1491+ 
1492+ # Re-removing any key is harmless (gears share keys → idempotent).
1493+ pool.remove_by_keys(keys_v0) # no crash
1494+ pool.remove_by_keys(keys_v1) # no crash
1495+ assert pool.entry_count == 0
1496+ 
1497+ @patch('torch.npu.synchronize')
1498+ def test_bulk_remove_mixed_keys(self, _mock_sync):
1499+ """Bulk remove handles a mix of existing and already-removed keys."""
1500+ pool = GraphResourcePool.get_pool(0)
1501+ k1, r1 = self._register_resource(pool)
1502+ k2, r2 = self._register_resource(pool)
1503+ 
1504+ # Remove k1 first.
1505+ pool.remove_by_keys([k1])
1506+ assert pool.entry_count == 1
1507+ 
1508+ # Bulk remove: k1 (already gone) + k2 (still present).
1509+ pool.remove_by_keys([k1, k2])
1510+ r1.release.assert_called_once() # only once, from first remove
1511+ r2.release.assert_called_once()
1512+ assert pool.entry_count == 0
1513+ 
1514+ @patch('torch.npu.synchronize')
1515+ def test_remove_calls_reset_when_no_release(self, _mock_sync):
1516+ """Simple-mode graph (torch.npu.NPUGraph) uses reset(), not release()."""
1517+ from unittest.mock import Mock
1518+ 
1519+ pool = GraphResourcePool.get_pool(0)
1520+ pool.activate()
1521+ resource = Mock(spec_set=['reset'])
1522+ key = pool.register(resource)
1523+ pool.remove_by_keys([key])
1524+ resource.reset.assert_called_once()
1525+ 
rain-666
rain-666rain-6666月1日

增加边界条件测试用例

likedislike
zhudada0120
zhudada0120
6月5日 评论：
1208 1526 
1209if __name__ == "__main__":1527if __name__ == "__main__":
1210 run_tests()1528 run_tests()
Atorch_npu/_inductor/adaptive_gears.py+895-0
@@ -0,0 +1,895 @@
1+from __future__ import annotations
2+ 
3+import copy
4+import logging
5+import math
6+import statistics
7+import threading
8+import time
9+from dataclasses import dataclass, field
10+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
11+ 
12+import torch
13+from torch_npu.npu._graph_resource_pool import GraphResourcePool
14+ 
15+logger = logging.getLogger(__name__)
16+ 
17+ 
18+# ---------------------------------------------------------------------------
19+# Models
20+# ---------------------------------------------------------------------------
21+ 
22+ 
23+@dataclass
24+class GearEventSample:
25+ """Single request event sample produced by shape handling, used as input for
26+ time-decayed window statistics (Section 5).
27+ 
28+ Fields:
29+ event_ts: Event timestamp in seconds, used for time-decay weight w_t(e) = exp(-λ · Δt).
30+ raw_value: Original dimension value of the request, used for raw length distribution
31+ to generate addition candidates via median.
32+ pad_ratio: Padding ratio of this request, used for gear weighted average pad cost.
33+ split_ratio: Split ratio of this request, used for gear weighted average split cost.
34+ """
35+ 
36+ event_ts: float
37+ raw_value: int
38+ pad_ratio: float
39+ split_ratio: float
40+ 
41+ 
42+@dataclass
43+class GearRuntimeState:
44+ """Per-gear runtime state storing decision data for eviction evaluation ("whether to evict"),
45+ not deletion safety ("whether safe to delete") which is handled by tree manager.
46+ 
47+ Fields:
48+ gear_id: Unique gear identifier in the format "{shape_type}:{gear_value}".
49+ shape_type: Dimension type this gear belongs to (e.g. "batch_size", "seq_len").
50+ gear_value: The gear threshold value.
51+ samples: Windowed event sample list, supporting hit_rate, avg_pad, avg_split computation.
52+ created_ts: Gear creation timestamp, used for recent-use protection window.
53+ last_hit_ts: Timestamp of the most recent hit, used for recent-use protection window.
54+ """
55+ 
56+ gear_id: str
57+ shape_type: str
58+ gear_value: int
59+ samples: List[GearEventSample] = field(default_factory=list)
60+ cleanup_keys: Set[int] = field(default_factory=set)
61+ created_ts: float = 0.0
62+ last_hit_ts: float = 0.0
63+ 
64+ 
65+@dataclass
66+class GearSnapshot:
67+ """Snapshot of the current gear set, enabling lock-free reads from the request thread
68+ via clone-on-read and atomic reference swaps on publish.
69+ 
70+ Fields:
71+ active_gears: Current active gears, keyed by shape_type, values are gear value lists.
72+ handler_configs: Per-dimension shape handler configs, used to rebuild config on update commit.
73+ shape_handling: Associated NPUShapeHandling instance for accessing graph manager and runtime config.
74+ version: Monotonically increasing version number, assigned at snapshot creation time.
75+ Useful for correlating logs across request and update threads.
76+ created_at: Unix timestamp (``time.time()``) when the snapshot was built.
77+ """
78+ 
79+ active_gears: Dict[str, List[int]]
80+ handler_configs: List[Dict[str, Any]]
81+ shape_handling: Any
82+ version: int = 0
83+ created_at: float = 0.0
84+ 
85+ def clone(self) -> "GearSnapshot":
86+ """Deep-copy the snapshot. active_gears and handler_configs are independent copies;
87+ shape_handling shares the reference. version and created_at are preserved."""
88+ return GearSnapshot(
89+ active_gears={shape_type: list(values) for shape_type, values in self.active_gears.items()},
90+ handler_configs=copy.deepcopy(self.handler_configs),
91+ shape_handling=self.shape_handling,
92+ version=self.version,
93+ created_at=self.created_at,
94+ )
95+ 
96+@dataclass
97+class ScoreBreakdown:
98+ """Scoring result for a single gear, used for eviction ranking, replacement loss recording,
99+ and observability logging (Section 12.2).
100+ 
101+ Fields:
102+ gear_id: Unique gear identifier.
103+ shape_type: Dimension type this gear belongs to.
104+ gear_value: The gear threshold value.
105+ hit_rate: Hit rate score component.
106+ avg_pad_ratio: Weighted average padding ratio component.
107+ avg_split_ratio: Weighted average split ratio component.
108+ score: Composite score = w_h * hit_rate - w_p * avg_pad - w_sp * avg_split.
109+ replace_loss: Replacement loss. Weighted average loss incurred when this gear's samples
110+ are remapped to the nearest remaining gear after deletion.
111+ replacement_gear_id: The best alternative gear_id identified during replacement loss computation.
112+ """
113+ 
114+ gear_id: str
115+ shape_type: str
116+ gear_value: int
117+ hit_rate: float
118+ avg_pad_ratio: float
119+ avg_split_ratio: float
120+ score: float
121+ replace_loss: float = 0.0
122+ replacement_gear_id: Optional[str] = None
123+ 
124+ 
125+# ---------------------------------------------------------------------------
126+# Snapshot
127+# ---------------------------------------------------------------------------
128+ 
129+ 
130+class GearSnapshotStore:
131+ """Thread-safe store for the current gear snapshot.
132+ 
133+ Reads are lock-free: a snapshot reference is read atomically (GIL) and
134+ cloned on-read, so the request thread never blocks on writes. Writes
135+ are serialized by a lock to prevent concurrent publishes from interleaving.
136+ """
137+ 
138+ def __init__(self, initial_state: GearSnapshot) -> None:
139+ self._lock = threading.Lock()
140+ self._current_snapshot = initial_state.clone()
141+ 
142+ def get_snapshot(self) -> GearSnapshot:
143+ return self._current_snapshot.clone()
144+ 
145+ def publish_snapshot(self, snapshot: GearSnapshot) -> GearSnapshot:
146+ with self._lock:
147+ self._current_snapshot = snapshot.clone()
148+ return self._current_snapshot
149+ 
150+ 
151+# ---------------------------------------------------------------------------
152+# Scorer
153+# ---------------------------------------------------------------------------
154+ 
155+ 
156+class GearScorer:
157+ """Scores gears based on time-decayed hit rate, padding cost, and split cost.
158+ 
159+ Produces a ``ScoreBreakdown`` per gear used for eviction ranking, and computes
160+ replacement loss to decide whether a specific gear is safe to remove.
161+ """
162+ 
163+ def __init__(self, adaptive_configs: Dict[str, float]) -> None:
164+ self.weight_hit = adaptive_configs.get("weight_hit", 0.60)
165+ self.weight_pad = adaptive_configs.get("weight_pad", 0.20)
166+ self.weight_split = adaptive_configs.get("weight_split", 0.20)
167+ 
168+ def build_score_breakdown(
169+ self,
170+ snapshot: GearSnapshot,
171+ stats_snapshot: Dict[str, Dict],
172+ ) -> Dict[str, ScoreBreakdown]:
173+ totals_by_type: Dict[str, float] = {}
174+ for stat in stats_snapshot.values():
175+ totals_by_type.setdefault(stat["shape_type"], 0.0)
176+ totals_by_type[stat["shape_type"]] += stat["weighted_hits"]
177+ 
178+ breakdowns: Dict[str, ScoreBreakdown] = {}
179+ for shape_type, gear_values in snapshot.active_gears.items():
180+ total_hits = max(totals_by_type.get(shape_type, 0.0), 1e-12)
181+ for gear_value in gear_values:
182+ gear_id = f"{shape_type}:{gear_value}"
183+ stat = stats_snapshot.get(gear_id)
184+ if stat is None:
185+ hit_rate = 0.0
186+ avg_pad_ratio = 0.0
187+ avg_split_ratio = 0.0
188+ else:
189+ hit_rate = stat["weighted_hits"] / total_hits
190+ avg_pad_ratio = stat["avg_pad_ratio"]
191+ avg_split_ratio = stat["avg_split_ratio"]
192+ score = (
193+ self.weight_hit * hit_rate
194+ - self.weight_pad * avg_pad_ratio
195+ - self.weight_split * avg_split_ratio
196+ )
197+ breakdowns[gear_id] = ScoreBreakdown(
198+ gear_id=gear_id,
199+ shape_type=shape_type,
200+ gear_value=gear_value,
201+ hit_rate=hit_rate,
202+ avg_pad_ratio=avg_pad_ratio,
203+ avg_split_ratio=avg_split_ratio,
204+ score=score,
205+ )
206+ return breakdowns
207+ 
208+ def compute_replace_loss(
209+ self,
210+ snapshot: GearSnapshot,
211+ stats_snapshot: Dict[str, Dict],
212+ gear_id: str,
213+ ) -> Tuple[float, Optional[str]]:
214+ stat = stats_snapshot.get(gear_id)
215+ if stat is None:
216+ return 0.0, None
217+ 
218+ shape_type = stat["shape_type"]
219+ gear_value = stat["gear_value"]
220+ alternatives = [
221+ candidate
222+ for candidate in snapshot.active_gears.get(shape_type, [])
223+ if candidate != gear_value
224+ ]
225+ if not alternatives:
226+ return float("inf"), None
227+ 
228+ raw_samples = stat.get("raw_samples", [])
229+ if not raw_samples:
230+ replacement = min(alternatives, key=lambda candidate: abs(candidate - gear_value))
231+ return 0.0, f"{shape_type}:{replacement}"
232+ 
233+ best_loss = float("inf")
234+ best_replacement = None
235+ for candidate in alternatives:
236+ total_loss = 0.0
237+ for raw_value in raw_samples:
238+ if candidate >= raw_value:
239+ total_loss += (candidate - raw_value) / max(candidate, 1)
240+ else:
241+ total_loss += (raw_value - candidate) / max(raw_value, 1)
242+ avg_loss = total_loss / max(len(raw_samples), 1)
243+ if avg_loss < best_loss:
244+ best_loss = avg_loss
245+ best_replacement = candidate
246+ 
247+ if best_replacement is None:
248+ return float("inf"), None
249+ return best_loss, f"{shape_type}:{best_replacement}"
250+ 
251+ 
252+# ---------------------------------------------------------------------------
253+# Runtime
254+# ---------------------------------------------------------------------------
255+ 
256+ 
257+class AdaptiveGearRuntime:
258+ """Central coordinator for the adaptive gear update mechanism.
259+ 
260+ Manages the main request-path logic (event recording, snapshot reads, update
261+ triggering) and owns the worker thread pool, scorer, and snapshot store.
262+ Graph cleanup after gear eviction is delegated to
263+ GraphResourcePool.remove_by_keys.
264+ """
265+ 
266+ DEFAULT_CONFIG = {
267+ "window_seconds": 300.0,
268+ "update_interval_seconds": 60.0,
269+ "pad_add_threshold": 0.35,
270+ "split_add_threshold": 0.20,
271+ "min_samples_per_gear": 5,
272+ "add_min_samples": 5,
273+ "min_gear_count_per_type": 1,
274+ "max_gears_per_type": 64,
275+ "recent_use_protect_seconds": 300.0,
276+ "replace_loss_threshold": 0.60,
277+ "weight_hit": 0.60,
278+ "weight_pad": 0.20,
279+ "weight_split": 0.20,
280+ "device_memory_usage_threshold_ratio": 0.90,
281+ }
282+ 
283+ def __init__(
284+ self,
285+ shape_handling_configs: List[Dict[str, Any]],
286+ adaptive_configs: Dict[str, Any],
287+ shape_handling_builder: Callable[[List[Dict[str, Any]]], Any],
288+ snapshot_store: Optional[GearSnapshotStore] = None,
289+ scorer: Optional[GearScorer] = None,
290+ ) -> None:
291+ self.device_index = torch.npu.current_device()
292+ self.shape_handling_builder = shape_handling_builder
293+ self.config = dict(self.DEFAULT_CONFIG)
294+ if adaptive_configs:
295+ self.config.update(adaptive_configs)
luqichao
luqichaoluqichao6月1日

这里直接使用用户输入的adaptive_configs进行update， 缺少合法性校验

likedislike
zhudada0120
zhudada0120
6月5日 评论：
296+ 
297+ self.scorer = scorer or GearScorer(self.config)
298+ self._sample_lock = threading.Lock()
299+ self._commit_lock = threading.Lock()
300+ self._snapshot_version = 0
301+ 
302+ normalized_configs = self._normalize_configs(shape_handling_configs)
303+ self.config_by_type = {config["type"]: config for config in normalized_configs}
304+ self._cached_ordered_configs = [
305+ self.config_by_type[shape_type] for shape_type in sorted(self.config_by_type.keys())
306+ ]
307+ initial_snapshot = self._build_snapshot(normalized_configs)
308+ self.snapshot_store = snapshot_store or GearSnapshotStore(initial_snapshot)
309+ self._states: Dict[str, GearRuntimeState] = {}
310+ self._ensure_runtime_states(initial_snapshot)
311+ 
312+ self.worker = GearUpdateWorker(self)
313+ 
314+ self._shutdown_event = threading.Event()
315+ self._worker_thread = threading.Thread(
316+ target=self._update_loop,
317+ daemon=True,
318+ name="adaptive-gears-loop",
319+ )
320+ self._worker_thread.start()
321+ 
322+ def shutdown(self) -> None:
323+ self._shutdown_event.set()
324+ if self._worker_thread is not None and self._worker_thread.is_alive():
325+ self._worker_thread.join(timeout=5.0)
326+ 
327+ def get_snapshot(self) -> GearSnapshot:
328+ return self.snapshot_store.get_snapshot()
329+ 
330+ def record_event(
331+ self,
332+ raw_gear_values: List[List[Optional[int]]],
333+ mapped_gear_values: List[List[Optional[int]]],
334+ pad_ratios: List[List[float]],
335+ split_ratios: List[List[float]],
336+ event_ts: float,
337+ cleanup_key: Optional[int] = None,
338+ ) -> None:
339+ if not raw_gear_values or not mapped_gear_values:
340+ return
341+ with self._sample_lock:
342+ for index, config in enumerate(self._ordered_configs()):
343+ raw_per_tensor = self._value_at(raw_gear_values, index, default=[])
344+ mapped_per_tensor = self._value_at(mapped_gear_values, index, default=[])
345+ cfg_pad = self._value_at(pad_ratios, index, default=[])
346+ cfg_split = self._value_at(split_ratios, index, default=[])
347+ for tensor_idx, (raw_value, mapped_value) in enumerate(
348+ zip(raw_per_tensor, mapped_per_tensor)
349+ ):
350+ if raw_value is None or mapped_value is None:
351+ continue
352+ gear_id = f"{config['type']}:{mapped_value}"
353+ state = self._states.get(gear_id)
354+ if state is None:
355+ state = GearRuntimeState(
356+ gear_id=gear_id,
357+ shape_type=config["type"],
358+ gear_value=mapped_value,
359+ )
360+ self._states[gear_id] = state
361+ sample = GearEventSample(
362+ event_ts=event_ts,
363+ raw_value=raw_value,
364+ pad_ratio=float(self._value_at(cfg_pad, tensor_idx, default=0.0) or 0.0),
365+ split_ratio=float(self._value_at(cfg_split, tensor_idx, default=0.0) or 0.0),
366+ )
367+ state.samples.append(sample)
368+ if cleanup_key is not None:
369+ state.cleanup_keys.add(cleanup_key)
370+ state.last_hit_ts = event_ts
371+ 
372+ def build_stats_snapshot(self, now_ts: float) -> Dict[str, Dict[str, Any]]:
373+ decay_lambda = math.log(2.0) / max(self.config["window_seconds"] / 2.0, 1e-6)
374+ window_seconds = self.config["window_seconds"]
375+ stats_snapshot: Dict[str, Dict[str, Any]] = {}
376+ with self._sample_lock:
377+ for gear_id, state in self._states.items():
378+ weighted_hits = 0.0
379+ weighted_pad_sum = 0.0
380+ weighted_split_sum = 0.0
381+ raw_samples = []
382+ kept_samples = []
383+ pad_sample_count = 0
384+ split_sample_count = 0
385+ for sample in state.samples:
386+ age = now_ts - sample.event_ts
387+ if age > window_seconds:
388+ continue
389+ weight = math.exp(-decay_lambda * max(age, 0.0))
390+ weighted_hits += weight
391+ weighted_pad_sum += sample.pad_ratio * weight
392+ weighted_split_sum += sample.split_ratio * weight
393+ raw_samples.append(sample.raw_value)
394+ if sample.pad_ratio > 0.0:
395+ pad_sample_count += 1
396+ if sample.split_ratio > 0.0:
397+ split_sample_count += 1
398+ kept_samples.append(sample)
399+ state.samples = kept_samples
400+ avg_pad_ratio = weighted_pad_sum / max(weighted_hits, 1e-12)
401+ avg_split_ratio = weighted_split_sum / max(weighted_hits, 1e-12)
402+ stats_snapshot[gear_id] = {
403+ "gear_id": gear_id,
404+ "shape_type": state.shape_type,
405+ "gear_value": state.gear_value,
406+ "weighted_hits": weighted_hits,
407+ "avg_pad_ratio": avg_pad_ratio,
408+ "avg_split_ratio": avg_split_ratio,
409+ "sample_count": len(raw_samples),
410+ "pad_sample_count": pad_sample_count,
411+ "split_sample_count": split_sample_count,
412+ "raw_samples": raw_samples,
413+ "created_ts": state.created_ts,
414+ "last_hit_ts": state.last_hit_ts,
415+ }
416+ current_snapshot = self.get_snapshot()
417+ for shape_type, values in current_snapshot.active_gears.items():
418+ for gear_value in values:
419+ gear_id = f"{shape_type}:{gear_value}"
420+ stats_snapshot.setdefault(
421+ gear_id,
422+ {
423+ "gear_id": gear_id,
424+ "shape_type": shape_type,
425+ "gear_value": gear_value,
426+ "weighted_hits": 0.0,
427+ "avg_pad_ratio": 0.0,
428+ "avg_split_ratio": 0.0,
429+ "sample_count": 0,
430+ "pad_sample_count": 0,
431+ "split_sample_count": 0,
432+ "raw_samples": [],
433+ "created_ts": 0.0,
434+ "last_hit_ts": 0.0,
435+ },
436+ )
437+ return stats_snapshot
438+ 
439+ def _update_loop(self) -> None:
440+ check_interval = self.config["update_interval_seconds"]
441+ while not self._shutdown_event.wait(timeout=check_interval):
442+ try:
443+ self.worker.run_once(time.time())
444+ except Exception:
445+ logger.warning("Adaptive gear update failed", exc_info=True)
446+ 
447+ def protect_gear_from_eviction(self, gear_id: str, now_ts: float) -> None:
448+ with self._sample_lock:
449+ state = self._states.get(gear_id)
450+ if state is not None:
451+ state.last_hit_ts = max(state.last_hit_ts, now_ts)
452+ 
453+ def commit_update(
454+ self,
455+ configs: List[Dict[str, Any]],
456+ removed_gears: List[str],
457+ now_ts: float = 0.0,
458+ ) -> Tuple[Optional[GearSnapshot], List[int]]:
459+ next_snapshot = self._build_snapshot(configs)
460+ published_snapshot = self.snapshot_store.publish_snapshot(next_snapshot)
461+ self._ensure_runtime_states(published_snapshot, created_ts=now_ts)
462+ removed_keys = self._collect_cleanup_keys(removed_gears)
463+ with self._sample_lock:
464+ for gear_id in removed_gears:
465+ self._states.pop(gear_id, None)
466+ return published_snapshot, removed_keys
467+ 
468+ def build_resource_budget(self) -> Dict[str, Any]:
469+ device_memory_threshold = self._normalize_ratio(self.config.get("device_memory_usage_threshold_ratio"))
470+ device_memory_usage_ratio = self._get_device_memory_usage_ratio()
471+ return {
472+ "device_memory_usage_ratio": device_memory_usage_ratio,
473+ "device_memory_usage_threshold_ratio": device_memory_threshold,
474+ "device_memory_usage_high": (
475+ device_memory_threshold is not None
476+ and device_memory_usage_ratio is not None
477+ and device_memory_usage_ratio >= device_memory_threshold
478+ ),
479+ }
480+ 
481+ def _ordered_configs(self) -> List[Dict[str, Any]]:
482+ return self._cached_ordered_configs
483+ 
484+ def _normalize_configs(self, configs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
rain-666rain-666
rain-666rain-6666月1日

_normalize_configs缺少对configs参数的完整性验证，建议添加必填字段检查和默认值设置

likedislike
zhudada0120
zhudada0120
6月5日 评论：
rain-666rain-6666月1日

建议补充配置验证规则和默认值设置的逻辑说明

likedislike
zhudada0120
zhudada0120
6月5日 评论：
485+ """Normalize user-supplied shape-handling configs into a canonical form.
486+ 
487+ Two paths:
488+ 1. *Explicit gears* — the user provides a ``gears`` list. Values are
489+ deduplicated, cast to int, and sorted.
490+ 2. *TIMES policy* — ``gears`` is empty; sizes are auto-generated from
491+ ``min_size``, ``max_size`` and the TIMES (power-of-two) strategy via
492+ ``_expand_policy_gears``. The expansion algorithm is kept consistent
493+ with ``NPUShapeHandling::GenerateGears`` in the C++ layer so that the
494+ initial gear set derived here matches what the C++ engine would
495+ produce internally.
496+ 
497+ The function modifies the config dicts in place and returns the same
498+ list (or ``[]`` for empty / None input).
499+ """
500+ if not configs:
501+ return []
502+ for config in configs:
503+ gears = config.get("gears") or self._expand_policy_gears(config)
504+ config["gears"] = sorted(set(int(gear) for gear in gears))
505+ return configs
506+ 
507+ def _expand_policy_gears(self, config: Dict[str, Any]) -> List[int]:
rain-666
rain-666rain-6666月1日

建议对这个算法补充注释说明

likedislike
zhudada0120
zhudada0120
6月5日 评论：
508+ """Expand ``min_size`` .. ``max_size`` using the TIMES (power-of-two) strategy.
509+ 
510+ The factor 2 is chosen because it keeps the gear count logarithmic
511+ relative to the size range — each gear is at most 2× the previous one,
512+ bounding the worst-case padding waste to ≤ 50 % while avoiding an
513+ explosion of graph variants. For example, range [1, 64] produces only
514+ 7 gears: [1, 2, 4, 8, 16, 32, 64].
515+ 
516+ This implementation mirrors ``NPUShapeHandling::GenerateGears`` in C++:
517+ it anchors at ``min_size``, then jumps to the next power of two and
518+ doubles from there, finally appending ``max_size`` if it was not already
519+ covered.
520+ """
521+ import math
522+ 
523+ min_size = int(config.get("min_size", 1))
524+ max_size = int(config.get("max_size", min_size))
525+ if max_size <= min_size:
526+ return [min_size]
527+ gears = [min_size]
528+ exp = math.ceil(math.log2(min_size)) if min_size > 0 else 0
529+ gear = 1 << exp
530+ if gear == min_size:
531+ gear <<= 1
532+ while gear > 0 and gear <= max_size:
533+ gears.append(gear)
534+ if gear > max_size // 2:
535+ break
536+ gear <<= 1
537+ if gears[-1] != max_size:
538+ gears.append(max_size)
539+ return gears
540+ 
541+ def _build_snapshot(self, configs: List[Dict[str, Any]]) -> GearSnapshot:
luqichao
luqichaoluqichao6月1日

snapshot是否可以增加版本号之类的信息，方便出现问题定位是哪个版本出的问题

likedislike
zhudada0120
zhudada0120
6月5日 评论：
542+ self._snapshot_version += 1
543+ active_gears = {config["type"]: list(config.get("gears", [])) for config in configs}
544+ shape_handling = self.shape_handling_builder(copy.deepcopy(configs))
545+ return GearSnapshot(
546+ active_gears=active_gears,
547+ handler_configs=copy.deepcopy(configs),
548+ shape_handling=shape_handling,
549+ version=self._snapshot_version,
550+ created_at=time.time(),
551+ )
552+ 
553+ def _ensure_runtime_states(self, snapshot: GearSnapshot, created_ts: float = 0.0) -> None:
554+ with self._sample_lock:
555+ for shape_type, gear_values in snapshot.active_gears.items():
556+ for gear_value in gear_values:
557+ gear_id = f"{shape_type}:{gear_value}"
558+ if gear_id not in self._states:
559+ self._states[gear_id] = GearRuntimeState(
560+ gear_id=gear_id,
561+ shape_type=shape_type,
562+ gear_value=gear_value,
563+ created_ts=created_ts,
564+ )
565+ 
566+ def _collect_cleanup_keys(self, gear_ids: List[str]) -> List[int]:
567+ keys: set[int] = set()
568+ with self._sample_lock:
569+ for gear_id in gear_ids:
570+ state = self._states.get(gear_id)
571+ if state is None:
572+ continue
573+ keys.update(state.cleanup_keys)
574+ state.cleanup_keys.clear()
575+ return list(keys)
576+ 
577+ def _normalize_ratio(self, ratio_value):
578+ if ratio_value is None:
579+ return None
580+ try:
581+ normalized = float(ratio_value)
582+ except (TypeError, ValueError):
583+ return None
584+ if normalized <= 0.0 or normalized > 1.0:
585+ return None
586+ return normalized
587+ 
588+ def _get_device_memory_usage_ratio(self):
589+ try:
590+ free_bytes, total_bytes = torch.npu.mem_get_info()
591+ except Exception:
592+ logger.debug("Failed to get device memory info", exc_info=True)
593+ return None
594+ if total_bytes <= 0:
595+ return None
596+ usage_ratio = 1.0 - (float(free_bytes) / float(total_bytes))
597+ return min(max(usage_ratio, 0.0), 1.0)
598+ 
599+ def _value_at(self, values, index: int, default=None):
600+ if isinstance(values, (list, tuple)) and index < len(values):
601+ return values[index]
602+ return default
603+ 
604+ 
605+# ---------------------------------------------------------------------------
606+# Update Worker
607+# ---------------------------------------------------------------------------
608+ 
609+ 
610+class GearUpdateWorker:
611+ """Asynchronous worker that executes gear update decisions.
612+ 
613+ Runs inside the worker thread pool. ``run_once`` performs scoring, eviction
614+ candidate selection, addition candidate generation, and commit -- all under
615+ the commit lock. After the lock is released, calls
616+ GraphResourcePool.remove_by_shapes to safely clean up evicted gear paths.
617+ 
618+ Parameters:
619+ runtime: The owning AdaptiveGearRuntime instance. Provides access to the
620+ snapshot store, sample lock, config, scorer, and device_index.
621+ """
622+ 
623+ def __init__(self, runtime: AdaptiveGearRuntime) -> None:
rain-666
rain-666rain-6666月1日

__init__方法参数缺少类型和含义说明，建议对每个参数添加注释说明

likedislike
zhudada0120
zhudada0120
6月5日 评论：
624+ self.runtime = runtime
625+ self.config = runtime.config
626+ self.scorer = runtime.scorer
627+ 
628+ def run_once(self, now_ts: float):
rain-666
rain-666rain-6666月1日

run_once方法注释过于简单，未说明执行周期和触发条件

likedislike
zhudada0120
zhudada0120
6月5日 评论：
629+ """Execute one cycle of the adaptive gear update algorithm.
630+ 
631+ Called periodically by the daemon thread at an interval controlled by
632+ ``update_interval_seconds`` (default: 60.0 s). Each cycle:
633+ 
634+ 1. Acquires the commit lock and snapshots the current state.
635+ 2. Computes time-decayed stats, scores each gear, and selects eviction
636+ and addition candidates.
637+ 3. Applies the update via ``commit_update``, which builds a new snapshot
638+ and publishes it atomically.
639+ 4. Releases the lock, then cleans up any evicted graph resources via
640+ ``GraphResourcePool.remove_by_shapes``.
641+ 
642+ The call is a no-op when there are no candidate changes to apply.
643+ """
644+ removed_keys: List[int] = []
645+ with self.runtime._commit_lock:
646+ snapshot = self.runtime.get_snapshot()
647+ stats_snapshot = self.runtime.build_stats_snapshot(now_ts)
648+ score_breakdowns = self.scorer.build_score_breakdown(snapshot, stats_snapshot)
649+ candidate_evictions = self.build_eviction_candidates(score_breakdowns, snapshot, stats_snapshot, now_ts)
650+ candidate_additions = self.build_addition_candidates(snapshot, stats_snapshot)
651+ resource_budget = self.runtime.build_resource_budget()
652+ result = self.commit_update(
653+ snapshot, stats_snapshot,
654+ candidate_evictions, candidate_additions, resource_budget, now_ts,
655+ )
656+ if result is not None:
657+ _, removed_keys = result
658+ 
659+ if removed_keys:
660+ self._cleanup_graphs(removed_keys, self.runtime.device_index)
661+ return result
662+ 
663+ @staticmethod
664+ def _cleanup_graphs(keys: List[int], device_index: int) -> None:
665+ """Release graph resources associated with evicted gears.
666+ 
667+ Pool keys (``(model_sig, tensor_sigs)`` tuples) are looked up in the
668+ per-device ``GraphResourcePool`` and removed. For each matched
669+ resource the pool calls ``release()`` (tree mode — ``NPUGraphNode``,
670+ which resets the underlying NPU graph and recycles block pointers) or
671+ ``reset()`` (simple mode — ``torch.npu.NPUGraph``). The
672+ ``torch.npu.synchronize()`` call before each release ensures all
673+ outstanding device work on those graphs has completed, preventing
674+ use-after-free.
675+ 
676+ Parameters:
677+ keys: Pool-key tuples previously registered via
678+ ``GraphResourcePool.register``.
679+ device_index: NPU device whose pool should be cleaned up. Must
680+ match the device on which the graphs were originally recorded.
681+ """
682+ GraphResourcePool.get_pool(device_index).remove_by_keys(keys)
683+ 
684+ def build_eviction_candidates(
685+ self,
686+ breakdowns: Dict[str, ScoreBreakdown],
687+ snapshot: GearSnapshot,
688+ stats_snapshot: Dict[str, Dict],
689+ now_ts: float,
690+ ) -> Dict[str, str]:
691+ max_gear_by_type = {
692+ shape_type: max(values)
693+ for shape_type, values in snapshot.active_gears.items()
694+ if values
695+ }
696+ candidates_by_type: Dict[str, List[Tuple[float, str]]] = {}
697+ for gear_id, breakdown in breakdowns.items():
698+ if breakdown.gear_value == max_gear_by_type.get(breakdown.shape_type):
699+ continue
700+ stat = stats_snapshot.get(
701+ gear_id,
702+ {
703+ "created_ts": 0.0,
704+ "last_hit_ts": 0.0,
705+ },
706+ )
707+ protect_anchor_ts = max(stat.get("created_ts", 0.0), stat.get("last_hit_ts", 0.0))
708+ if now_ts - protect_anchor_ts < self.config["recent_use_protect_seconds"]:
709+ continue
710+ if len(snapshot.active_gears.get(breakdown.shape_type, [])) <= self.config["min_gear_count_per_type"]:
711+ continue
712+ candidates_by_type.setdefault(breakdown.shape_type, []).append((breakdown.score, gear_id))
713+ selected_candidates = {}
714+ for shape_type, candidates in candidates_by_type.items():
715+ candidates.sort(key=lambda item: item[0])
716+ selected_candidates[shape_type] = candidates[0][1]
717+ return selected_candidates
718+ 
719+ def build_addition_candidates(
720+ self,
721+ snapshot: GearSnapshot,
722+ stats_snapshot: Dict[str, Dict],
723+ ) -> List[Tuple[float, int, str, int]]:
724+ max_gear_by_type = {
725+ shape_type: max(values)
726+ for shape_type, values in snapshot.active_gears.items()
727+ if values
728+ }
729+ candidates: List[Tuple[float, int, str, int]] = []
730+ 
731+ for stat in stats_snapshot.values():
732+ shape_type = stat["shape_type"]
733+ gear_value = stat["gear_value"]
734+ active_gears = set(snapshot.active_gears.get(shape_type, []))
735+ raw_samples = stat.get("raw_samples", [])
736+ if not raw_samples:
737+ continue
738+ is_max_gear = gear_value == max_gear_by_type.get(shape_type)
739+ 
740+ if is_max_gear:
741+ enough_split = stat.get("split_sample_count", 0) >= self.config["add_min_samples"]
742+ high_split = stat["avg_split_ratio"] >= self.config["split_add_threshold"]
743+ if enough_split and high_split:
744+ g_new = int(statistics.median(raw_samples))
745+ if g_new not in active_gears:
746+ pressure = stat["avg_split_ratio"] - self.config["split_add_threshold"]
747+ candidates.append((pressure, stat["sample_count"], shape_type, g_new))
748+ 
749+ enough_pad = stat.get("pad_sample_count", 0) >= self.config["add_min_samples"]
750+ high_pad = stat["avg_pad_ratio"] >= self.config["pad_add_threshold"]
751+ if enough_pad and high_pad:
752+ g_new = int(statistics.median(raw_samples))
753+ if g_new not in active_gears:
754+ pressure = stat["avg_pad_ratio"] - self.config["pad_add_threshold"]
755+ candidates.append((pressure, stat["sample_count"], shape_type, g_new))
756+ 
757+ candidates.sort(key=lambda x: (x[0], x[1]), reverse=True)
758+ return candidates
759+ 
760+ def commit_update(
761+ self,
762+ snapshot: GearSnapshot,
763+ stats_snapshot: Dict[str, Dict],
764+ candidate_evictions: Dict[str, str],
765+ candidate_additions: List[Tuple[float, int, str, int]],
766+ resource_budget: Dict[str, Any],
767+ now_ts: float,
768+ ) -> Optional[Tuple[Optional[GearSnapshot], List[Tuple[Tuple[int, ...], ...]]]]:
769+ resource_pressure_high = bool(resource_budget.get("device_memory_usage_high", False))
770+ next_configs = copy.deepcopy(snapshot.handler_configs)
771+ next_active_gears = {
772+ shape_type: set(values)
773+ for shape_type, values in snapshot.active_gears.items()
774+ }
775+ removed_gears: List[str] = []
776+ added = False
777+ 
778+ for shape_type, gear_id in candidate_evictions.items():
779+ gear_value = int(gear_id.split(":", 1)[1])
780+ replace_loss, _ = self.scorer.compute_replace_loss(
781+ snapshot,
782+ stats_snapshot,
783+ gear_id,
784+ )
785+ if replace_loss > self.config["replace_loss_threshold"]:
786+ self.runtime.protect_gear_from_eviction(gear_id, now_ts)
787+ continue
788+ next_active_gears[shape_type].discard(gear_value)
789+ removed_gears.append(gear_id)
790+ 
791+ addition_budget = None if not resource_pressure_high else len(removed_gears)
792+ for _, _, shape_type, candidate in candidate_additions:
793+ if addition_budget is not None and addition_budget <= 0:
794+ break
795+ if candidate in next_active_gears.setdefault(shape_type, set()):
796+ continue
797+ # Enforce per-type gear count cap to prevent unbounded graph variant growth
798+ # (each gear produces a separate compiled graph, consuming device memory).
799+ if len(next_active_gears.get(shape_type, set())) >= self.config["max_gears_per_type"]:
800+ continue
801+ next_active_gears[shape_type].add(candidate)
luqichao
luqichaoluqichao6月1日

这里gear总数是否可能超过MAX_GEARS_NUM(64)

likedislike
zhudada0120
zhudada0120
6月5日 评论：
802+ added = True
803+ if addition_budget is not None:
804+ addition_budget -= 1
805+ 
806+ if not removed_gears and not added:
807+ return None
808+ 
809+ for config in next_configs:
810+ active_values = sorted(next_active_gears.get(config["type"], set()))
811+ config["gears"] = active_values
812+ config["policy"] = "CUSTOM"
813+ 
814+ return self.runtime.commit_update(next_configs, removed_gears, now_ts=now_ts)
815+ 
816+ 
817+# ---------------------------------------------------------------------------
818+# Metadata collection utilities
819+# ---------------------------------------------------------------------------
820+ 
821+ 
822+def _resolve_dimension(config: Dict[str, Any], dimensions: List[int], position: int) -> int:
823+ if config.get("type") == "BATCHSIZE":
824+ if len(dimensions) == 0:
825+ return 0
826+ return dimensions[0]
827+ if len(dimensions) == 0:
828+ return 1
829+ if position < len(dimensions):
830+ return dimensions[position]
831+ return dimensions[0]
832+ 
833+ 
834+def _extract_gear_shapes(tensors: List[torch.Tensor], configs: List[Dict[str, Any]]) -> List[List[Optional[int]]]:
835+ all_shapes: List[List[Optional[int]]] = []
836+ for config in configs:
837+ indices = config.get("indices", [])
838+ dimensions = config.get("dimensions", [])
839+ target_indices = indices if len(indices) > 0 else list(range(len(tensors)))
840+ config_values: List[Optional[int]] = []
841+ for position, tensor_index in enumerate(target_indices):
842+ if tensor_index >= len(tensors):
843+ config_values.append(None)
844+ continue
845+ tensor = tensors[tensor_index]
846+ dimension = _resolve_dimension(config, dimensions, position)
847+ if tensor.ndim > dimension:
848+ config_values.append(tensor.shape[dimension])
849+ else:
850+ config_values.append(None)
851+ all_shapes.append(config_values)
852+ return all_shapes
853+ 
854+ 
855+def collect_transform_metadata(
luqichao
luqichaoluqichao6月1日

这里使用的configs应该是用户传入的顺序， 可能是[SEQLEN, BATCHSIZE]， AdaptiveGearRuntime中会按shape_type 进行排序，然后再记录， 看起来会导致错位

likedislike
zhudada0120
zhudada0120
6月5日 评论：
856+ inputs: List[torch.Tensor],
857+ trans_outputs: List[List[torch.Tensor]],
858+ configs: List[Dict[str, Any]],
859+) -> List[Dict[str, Any]]:
860+ """Collect per-variant pad / split metadata for adaptive gear recording.
861+ 
862+ The *cleanup_key* is no longer constructed here — it is obtained from
863+ ``GraphResourcePool.consume_recent_keys()`` after ``src_fn`` returns.
864+ """
865+ # Sort configs by type to match AdaptiveGearRuntime._ordered_configs().
866+ sorted_configs = sorted(configs, key=lambda c: c["type"])
867+ raw_gear_values = _extract_gear_shapes(inputs, sorted_configs)
868+ metadata = []
869+ for output_group in trans_outputs:
870+ mapped_gear_values = _extract_gear_shapes(output_group, sorted_configs)
871+ pad_ratios: List[List[float]] = []
872+ split_ratios: List[List[float]] = []
873+ for raw_per_cfg, mapped_per_cfg in zip(raw_gear_values, mapped_gear_values):
874+ cfg_pad: List[float] = []
875+ cfg_split: List[float] = []
876+ for raw_value, mapped_value in zip(raw_per_cfg, mapped_per_cfg):
877+ if mapped_value is None:
878+ cfg_pad.append(0.0)
879+ cfg_split.append(0.0)
880+ continue
881+ r = raw_value if raw_value is not None else mapped_value
882+ cfg_pad.append(max(mapped_value - r, 0) / max(mapped_value, 1))
883+ cfg_split.append(max(r - mapped_value, 0) / max(r, 1))
884+ pad_ratios.append(cfg_pad)
885+ split_ratios.append(cfg_split)
886+ 
887+ metadata.append(
888+ {
889+ "raw_gear_values": raw_gear_values,
890+ "mapped_gear_values": mapped_gear_values,
891+ "pad_ratios": pad_ratios,
892+ "split_ratios": split_ratios,
893+ }
894+ )
895+ return metadata
Mtorch_npu/_inductor/shape_handling.py+183-25
@@ -1,13 +1,17 @@
1-__all__ = ["NPUShapeHandling"]1+__all__ = ["NPUShapeHandling", "unified_copy"]
2 2 
3from typing import Any, Callable, Dict, List, Optional, Tuple3from typing import Any, Callable, Dict, List, Optional, Tuple
4import copy4import copy
5import logging5import logging
6+import sys
7+import time
6import warnings8import warnings
7from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec9from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec
8import torch10import torch
9import torch_npu._C11import torch_npu._C
10 12 
13+logger = logging.getLogger(__name__)
14+ 
11 15 
12class NPUShapeHandling(torch_npu._C._NPUShapeHandling):16class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
13 r"""Wrapper around a NPU shape handling configuration.17 r"""Wrapper around a NPU shape handling configuration.
@@ -47,6 +51,7 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
47 def __init__(51 def __init__(
48 self,52 self,
49 configs: List[Dict[str, Any]] = None,53 configs: List[Dict[str, Any]] = None,
54+ adaptive_configs: Optional[Dict[str, Any]] = None,
50 transform_pre_fn: Optional[Callable[..., List[torch.Tensor]]] = None,55 transform_pre_fn: Optional[Callable[..., List[torch.Tensor]]] = None,
51 transform_post_fn: Optional[Callable[[List[List[torch.Tensor]]], Tuple[List[Tuple], List[Dict]]]] = None,56 transform_post_fn: Optional[Callable[[List[List[torch.Tensor]]], Tuple[List[Tuple], List[Dict]]]] = None,
52 recover_pre_fn: Optional[Callable[[List[Any]], List[List[torch.Tensor]]]] = None,57 recover_pre_fn: Optional[Callable[[List[Any]], List[List[torch.Tensor]]]] = None,
@@ -69,6 +74,8 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
69 self.transform_post_fn = transform_post_fn74 self.transform_post_fn = transform_post_fn
70 self.recover_pre_fn = recover_pre_fn75 self.recover_pre_fn = recover_pre_fn
71 self.recover_post_fn = recover_post_fn76 self.recover_post_fn = recover_post_fn
77+ self.adaptive_configs = adaptive_configs
78+ self._adaptive_manager = None
72 if configs and len(configs) > 0:79 if configs and len(configs) > 0:
73 self._validate_configs(configs)80 self._validate_configs(configs)
74 self.configs = configs81 self.configs = configs
@@ -85,6 +92,71 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
85 "max_size": 1024,92 "max_size": 1024,
86 "policy": "TIMES"93 "policy": "TIMES"
87 }]94 }]
95+ if self.adaptive_configs is not None:
96+ self._validate_adaptive_configs(self.adaptive_configs)
97+ self._adaptive_manager = self._create_adaptive_manager()
98+ 
99+ @property
100+ def adaptive_manager(self):
101+ return self._adaptive_manager
102+ 
103+ def _validate_adaptive_configs(self, adaptive_configs: Dict[str, Any]) -> None:
104+ from .adaptive_gears import AdaptiveGearRuntime
105+ 
106+ _ratio_keys = {
107+ "pad_add_threshold", "split_add_threshold", "replace_loss_threshold",
108+ "weight_hit", "weight_pad", "weight_split",
109+ "device_memory_usage_threshold_ratio",
110+ }
111+ _seconds_keys = {
112+ "window_seconds", "recent_use_protect_seconds",
113+ "update_interval_seconds",
114+ }
115+ _int_keys = {
116+ "min_samples_per_gear", "min_gear_count_per_type",
117+ "add_min_samples", "max_gears_per_type",
118+ }
119+ 
120+ valid_keys = set(AdaptiveGearRuntime.DEFAULT_CONFIG.keys())
121+ for key, value in adaptive_configs.items():
122+ if key not in valid_keys:
123+ raise ValueError(
124+ f"Unknown adaptive config key: '{key}'. "
125+ f"Valid keys: {sorted(valid_keys)}"
126+ )
127+ if key in _ratio_keys:
128+ if not isinstance(value, (int, float)) or not (0.0 <= value <= 1.0):
129+ raise ValueError(
130+ f"adaptive config '{key}' must be a number in [0, 1], got {value!r}"
131+ )
132+ elif key in _seconds_keys:
133+ if not isinstance(value, (int, float)) or value < 0:
134+ raise ValueError(
135+ f"adaptive config '{key}' must be >= 0, got {value!r}"
136+ )
137+ elif key in _int_keys:
138+ if not isinstance(value, int) or value < 1:
139+ raise ValueError(
140+ f"adaptive config '{key}' must be a positive int, got {value!r}"
141+ )
142+ 
143+ def _create_adaptive_manager(self):
144+ from .adaptive_gears import AdaptiveGearRuntime
145+ 
146+ return AdaptiveGearRuntime(
147+ shape_handling_configs=self.configs,
148+ adaptive_configs=self.adaptive_configs,
149+ shape_handling_builder=self._build_adaptive_snapshot_handler,
150+ )
151+ 
152+ def _build_adaptive_snapshot_handler(self, shape_configs: List[Dict[str, Any]]):
153+ return NPUShapeHandling(
154+ configs=shape_configs,
155+ transform_pre_fn=self.transform_pre_fn,
156+ transform_post_fn=self.transform_post_fn,
157+ recover_pre_fn=self.recover_pre_fn,
158+ recover_post_fn=self.recover_post_fn,
159+ )
88 160 
89 def _validate_configs(self, configs: List[Dict[str, Any]]) -> None:161 def _validate_configs(self, configs: List[Dict[str, Any]]) -> None:
90 if not configs or len(configs) == 0:162 if not configs or len(configs) == 0:
@@ -237,37 +309,70 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
237 *args: Any,309 *args: Any,
238 **kwargs: Any310 **kwargs: Any
239 ) -> Tuple[List[Tuple], List[Dict]]:311 ) -> Tuple[List[Tuple], List[Dict]]:
240- # 获取 logger312+ outputs, _, _ = self._transform_with_context(*args, **kwargs)
241- logger = logging.getLogger(__name__)313+ return outputs
242- # 预处理阶段优化：统一使用预定义函数或默认逻辑314+ 
315+ def _transform_with_metadata(
rain-666rain-666
rain-666rain-6666月1日

_transform_with_metadata缺少异常处理，建议添加try-catch块，确保元数据收集失败不影响主流程

likedislike
zhudada0120
zhudada0120
6月5日 评论：
rain-666rain-6666月1日

新增的元数据收集功能，建议补充注释说明

likedislike
zhudada0120
zhudada0120
6月5日 评论：
316+ self,
317+ *args: Any,
318+ **kwargs: Any
319+ ) -> Tuple[Tuple[List[Tuple], List[Dict]], List[Dict[str, Any]]]:
320+ """Like :meth:`transform_hook`, but also collects per-gear metadata.
321+ 
322+ After shape transformation, this method extracts the raw (pre-transform)
323+ and mapped (post-transform) dimension values for every tensor / config
324+ pair, computes per-tensor pad and split ratios, and packages them
325+ together with ``cleanup_shapes`` (the post-transform tensor shapes used
326+ for graph resource pool indexing) into a list of metadata dicts — one
327+ per gear variant produced by the transform.
328+ 
329+ The metadata is consumed by :meth:`AdaptiveGearRuntime.record_event` to
330+ build the time-decayed statistics that drive gear eviction / addition
331+ decisions.
332+ 
333+ Failures in metadata collection are caught and logged; an empty metadata
334+ list is returned in that case so that the core inference path is never
335+ blocked.
336+ """
337+ from .adaptive_gears import collect_transform_metadata
338+ 
339+ outputs, inputs, trans_outputs = self._transform_with_context(*args, **kwargs)
340+ try:
341+ metadata = collect_transform_metadata(inputs, trans_outputs, self.configs)
342+ except Exception:
343+ logger.warning("Failed to collect transform metadata, adaptive gears degraded", exc_info=True)
344+ metadata = []
345+ return outputs, metadata
346+ 
347+ def _transform_with_context(
rain-666
rain-666rain-6666月1日

建议提取公共逻辑，减少代码重复

likedislike
zhudada0120
zhudada0120
6月5日 评论：
348+ self,
349+ *args: Any,
350+ **kwargs: Any
351+ ) -> Tuple[Tuple[List[Tuple], List[Dict]], List[torch.Tensor], List[List[torch.Tensor]]]:
243 if self.transform_pre_fn:352 if self.transform_pre_fn:
244 inputs = self.transform_pre_fn(*args, **kwargs)353 inputs = self.transform_pre_fn(*args, **kwargs)
245 else:354 else:
246 inputs, indices, leaves, spec = self._process_inputs(args, kwargs)355 inputs, indices, leaves, spec = self._process_inputs(args, kwargs)
247- 356+ 
248- # 提取转换前的形状 (inputs 通常是 Tensor 列表)
249 if logger.isEnabledFor(logging.INFO):357 if logger.isEnabledFor(logging.INFO):
250 pre_shapes = [self.get_shape_safe(t) for t in inputs]358 pre_shapes = [self.get_shape_safe(t) for t in inputs]
251 logger.info(f"[Transform] Starting. Input tensors: {len(inputs)}, Shapes: {pre_shapes}")359 logger.info(f"[Transform] Starting. Input tensors: {len(inputs)}, Shapes: {pre_shapes}")
252 360 
253- # 执行核心转换操作
254 trans_outputs = self.transform(tensors=inputs)361 trans_outputs = self.transform(tensors=inputs)
255 362 
256- # 提取转换后的形状
257 if logger.isEnabledFor(logging.INFO):363 if logger.isEnabledFor(logging.INFO):
258 post_shapes = [self.get_shape_safe(t) for t in trans_outputs]364 post_shapes = [self.get_shape_safe(t) for t in trans_outputs]
259 logger.info(f"> Post-transform content: {post_shapes}")365 logger.info(f"> Post-transform content: {post_shapes}")
260- 366+ 
261- # 后处理阶段优化：避免嵌套循环
262 if self.transform_post_fn:367 if self.transform_post_fn:
263 outputs = self.transform_post_fn(trans_outputs)368 outputs = self.transform_post_fn(trans_outputs)
264 else:369 else:
265 outputs = self._recover_inputs(trans_outputs, indices, leaves, spec)370 outputs = self._recover_inputs(trans_outputs, indices, leaves, spec)
266- 371+ 
267 if not outputs:372 if not outputs:
268 logger.error(f"CRITICAL: _recover_inputs returned NULL")373 logger.error(f"CRITICAL: _recover_inputs returned NULL")
269- 374+ 
270- return outputs375+ return outputs, inputs, trans_outputs
271 376 
272 def flatten_to_tensors(self, structure: Any) -> Tuple[List[torch.Tensor], List[int], List[Any], TreeSpec]:377 def flatten_to_tensors(self, structure: Any) -> Tuple[List[torch.Tensor], List[int], List[Any], TreeSpec]:
273 leaves, spec = tree_flatten(structure)378 leaves, spec = tree_flatten(structure)
@@ -302,7 +407,10 @@ class NPUShapeHandling(torch_npu._C._NPUShapeHandling):
302 res = []407 res = []
303 for processd_tensors in transform_res:408 for processd_tensors in transform_res:
304 res.append(self.unflatten_from_tensors(processd_tensors, indices, list(leaves), spec))409 res.append(self.unflatten_from_tensors(processd_tensors, indices, list(leaves), spec))
305- return zip(*res)410+ if not res:
411+ return [], []
412+ args_list, kwargs_list = zip(*res)
413+ return list(args_list), list(kwargs_list)
306 414
307 def _process_outputs(415 def _process_outputs(
308 self,416 self,
@@ -458,15 +566,18 @@ def patch_dynamo_context():
458 trans_post_fn = None566 trans_post_fn = None
459 re_pre_fn = None567 re_pre_fn = None
460 re_post_fn = None568 re_post_fn = None
569+ adaptive_configs = None
461 function_dict = compiler_config.get("shape_handling_dict")570 function_dict = compiler_config.get("shape_handling_dict")
462 if function_dict is not None:571 if function_dict is not None:
463 trans_pre_fn = function_dict.get("trans_pre_fn", None)572 trans_pre_fn = function_dict.get("trans_pre_fn", None)
464 trans_post_fn = function_dict.get("trans_post_fn", None)573 trans_post_fn = function_dict.get("trans_post_fn", None)
465 re_pre_fn = function_dict.get("re_pre_fn", None)574 re_pre_fn = function_dict.get("re_pre_fn", None)
466 re_post_fn = function_dict.get("re_post_fn", None)575 re_post_fn = function_dict.get("re_post_fn", None)
576+ adaptive_configs = function_dict.get("adaptive_gears", None)
467 577
468 self.shape_handling = NPUShapeHandling(578 self.shape_handling = NPUShapeHandling(
469 configs=compiler_config.get("shape_handling_configs"),579 configs=compiler_config.get("shape_handling_configs"),
580+ adaptive_configs=adaptive_configs,
470 transform_pre_fn=trans_pre_fn,581 transform_pre_fn=trans_pre_fn,
471 transform_post_fn=trans_post_fn,582 transform_post_fn=trans_post_fn,
472 recover_pre_fn=re_pre_fn,583 recover_pre_fn=re_pre_fn,
@@ -477,19 +588,66 @@ def patch_dynamo_context():
477 src_fn = src_call(self, fn)588 src_fn = src_call(self, fn)
478 if isinstance(fn, torch.nn.Module) or inspect.isclass(fn):589 if isinstance(fn, torch.nn.Module) or inspect.isclass(fn):
479 return src_fn590 return src_fn
480- 591+ 
481 def new_fn(*args, **kwargs):592 def new_fn(*args, **kwargs):
482 if (is_enable_shape_handling(self.callback, compiler_config=self.compiler_config)):593 if (is_enable_shape_handling(self.callback, compiler_config=self.compiler_config)):
483- new_args, new_kwargs = self.shape_handling.transform_hook(*args, **kwargs)594+ sh = self.shape_handling
484- args_is_split = len(args) != 0 and len(new_args) > 1595+ manager = sh.adaptive_manager
485- kwargs_is_split = len(kwargs) != 0 and len(new_kwargs) > 1596+ snapshot = manager.get_snapshot() if manager else None
486- zipped_params = zip(new_args, new_kwargs)597+ target_sh = snapshot.shape_handling if snapshot else sh
487- res = [598+ 
488- unified_copy(src_fn(*arg, **kwargs)) if args_is_split or kwargs_is_split599+ pool = None
489- else src_fn(*arg, **kwargs)600+ try:
490- for arg, kwargs in zipped_params601+ if manager:
491- ]602+ from torch_npu.npu._graph_resource_pool import GraphResourcePool
492- return self.shape_handling.recover_hook(res)603+ pool = GraphResourcePool.get_pool(torch.npu.current_device())
604+ pool.activate()
605+ (new_args, new_kwargs), metadata = target_sh._transform_with_metadata(*args, **kwargs)
606+ else:
607+ new_args, new_kwargs = target_sh.transform_hook(*args, **kwargs)
608+ metadata = []
609+ args_is_split = len(args) != 0 and len(new_args) > 1
610+ kwargs_is_split = len(kwargs) != 0 and len(new_kwargs) > 1
611+ res = []
612+ for index, (call_args, call_kwargs) in enumerate(zip(new_args, new_kwargs)):
613+ raised_exc = None
614+ try:
615+ result = (
616+ unified_copy(src_fn(*call_args, **call_kwargs))
617+ if args_is_split or kwargs_is_split
618+ else src_fn(*call_args, **call_kwargs)
619+ )
620+ except Exception:
621+ raised_exc = sys.exc_info()[0]
622+ raise
623+ finally:
624+ if pool is not None:
625+ # Consume graph keys registered during this
626+ # src_fn invocation (pool.activate was called).
627+ keys = pool.consume_recent_keys()
628+ if raised_exc is not None:
629+ if keys:
630+ pool.remove_by_keys(keys)
631+ elif manager and snapshot and index < len(metadata):
632+ meta = metadata[index]
633+ for key in (keys or [None]):
634+ try:
635+ manager.record_event(
636+ raw_gear_values=meta["raw_gear_values"],
637+ mapped_gear_values=meta["mapped_gear_values"],
638+ pad_ratios=meta["pad_ratios"],
639+ split_ratios=meta["split_ratios"],
640+ event_ts=time.time(),
641+ cleanup_key=key,
642+ )
643+ except Exception:
644+ logger.warning("Failed to record adaptive gear event", exc_info=True)
645+ res.append(result)
646+ output = target_sh.recover_hook(res)
647+ return output
648+ finally:
649+ if pool is not None:
650+ pool.deactivate()
493 return src_fn(*args, **kwargs)651 return src_fn(*args, **kwargs)
494 return new_fn652 return new_fn
495 _TorchDynamoContext.__call__ = new_call653 _TorchDynamoContext.__call__ = new_call
Atorch_npu/npu/_graph_resource_pool.py+166-0
@@ -0,0 +1,166 @@
1+from __future__ import annotations
2+ 
3+import threading
4+from typing import Any, Dict, List
5+ 
6+import torch
7+ 
8+ 
9+class GraphResourcePool:
10+ """Per-device graph resource registry with opaque integer keys.
11+ 
12+ Each graph resource (``NPUGraphNode`` in tree mode, ``torch.npu.NPUGraph``
13+ in simple mode) is registered with an auto-incrementing integer key. The
14+ key is a pure lookup handle — it carries no semantic information.
15+ 
16+ Two kinds of resources are supported, corresponding to the two graph
17+ compilation modes:
18+ 
19+ * tree mode — ``NPUGraphNode`` (has ``release()`` → block pointers)
20+ * simple mode — ``torch.npu.NPUGraph`` (has ``reset()``)
21+ 
22+ **Key lifecycle**::
23+ 
24+ # Lower layer (graph capture, inside src_fn):
25+ key = pool.register(node) # → opaque int
26+ 
27+ # Upper layer (shape_handling, after src_fn returns):
28+ keys = pool.consume_recent_keys() # → List[int]
29+ 
30+ # Background (gear eviction):
31+ pool.remove_by_keys([42, 43]) # idempotent
32+ 
33+ Thread safety
34+ -------------
35+ All mutations to ``_entries`` and ``_pending_by_thread`` are serialised
36+ under ``self._lock``. The heavy work — ``torch.npu.synchronize`` and
37+ ``resource.release() / reset()`` — is performed **outside** the lock.
38+ 
39+ ``register`` and ``consume_recent_keys`` form a per-thread
40+ producer-consumer pair: ``register`` appends keys to the calling thread's
41+ pending list, ``consume_recent_keys`` drains it. This guarantees that
42+ concurrent inferences on different threads never mix their keys.
43+ 
44+ ``remove_by_keys`` is idempotent: a key that has already been removed
45+ (by a prior gear eviction) simply results in ``_entries.pop(key, None)``
46+ returning ``None``, which is skipped.
47+ """
48+ 
49+ _pools: Dict[int, "GraphResourcePool"] = {}
rain-666
rain-666rain-6666月1日

资源池的线程安全保证未在注释中说明，补充说明线程安全范围和锁粒度

likedislike
zhudada0120
zhudada0120
6月5日 评论：
50+ 
51+ # ---- pool lifecycle -------------------------------------------------
52+ 
53+ @classmethod
54+ def get_pool(cls, device_index: int) -> "GraphResourcePool":
55+ if device_index not in cls._pools:
56+ cls._pools[device_index] = cls(device_index)
57+ return cls._pools[device_index]
58+ 
59+ @classmethod
60+ def reset_all(cls) -> None:
61+ cls._pools.clear()
62+ 
63+ def __init__(self, device_index: int) -> None:
64+ self.device_index = device_index
65+ self._lock = threading.Lock()
66+ self._next_key = 0
67+ # key → graph resource (1:1 mapping)
68+ self._entries: Dict[int, Any] = {}
69+ # thread-id → keys registered since last consume (producer→consumer)
70+ self._pending_by_thread: Dict[int, List[int]] = {}
71+ # thread-id → activation refcount (0 ≈ not active)
72+ self._active_threads: Dict[int, int] = {}
73+ 
74+ # ---- activation (per-thread, refcounted) ----------------------------
75+ 
76+ def activate(self) -> None:
77+ """Allow :meth:`register` on the calling thread.
78+ 
79+ Safe to call multiple times — each call must be balanced by a
80+ corresponding :meth:`deactivate`.
81+ """
82+ tid = threading.get_ident()
83+ with self._lock:
84+ self._active_threads[tid] = self._active_threads.get(tid, 0) + 1
85+ 
86+ def deactivate(self) -> None:
87+ """Revoke one :meth:`activate` on the calling thread.
88+ 
89+ When the refcount drops to zero, subsequent :meth:`register` calls
90+ on this thread become no-ops.
91+ """
92+ tid = threading.get_ident()
93+ with self._lock:
94+ v = self._active_threads.get(tid, 1) - 1
95+ if v <= 0:
96+ self._active_threads.pop(tid, None)
97+ else:
98+ self._active_threads[tid] = v
99+ 
100+ def is_active(self) -> bool:
101+ """Return ``True`` if the calling thread has been activated."""
102+ tid = threading.get_ident()
103+ with self._lock:
104+ return tid in self._active_threads
105+ 
106+ # ---- lower-layer API (graph capture) --------------------------------
107+ 
108+ def register(self, resource: Any) -> int:
109+ """Register a graph resource and return an opaque integer key.
110+ 
111+ Called during graph capture from ``record_function`` (tree mode) or
112+ ``npugraphify_impl`` (simple mode). Callers should guard the call
113+ with :meth:`is_active` when registration is conditional on the
114+ adaptive-gear feature being enabled.
115+ """
116+ tid = threading.get_ident()
117+ with self._lock:
118+ key = self._next_key
119+ self._next_key += 1
120+ self._entries[key] = resource
121+ self._pending_by_thread.setdefault(tid, []).append(key)
122+ return key
123+ 
124+ # ---- upper-layer API (after src_fn returns) -------------------------
125+ 
126+ def consume_recent_keys(self) -> List[int]:
127+ """Return and clear all keys registered on this thread since the
128+ last call to ``consume_recent_keys``.
129+ 
130+ Called in ``new_fn`` immediately after ``src_fn`` returns, so the
131+ returned keys correspond exactly to the graph resources captured
132+ during that ``src_fn`` invocation.
133+ """
134+ tid = threading.get_ident()
135+ with self._lock:
136+ return self._pending_by_thread.pop(tid, [])
137+ 
138+ # ---- cleanup API (gear eviction) ------------------------------------
139+ 
140+ def remove_by_keys(self, keys: List[int]) -> None:
141+ """Release graph resources identified by *keys*.
142+ 
143+ Idempotent — keys that have already been removed are silently
144+ skipped. Called by ``GearUpdateWorker`` on the background thread
145+ after gear eviction.
146+ """
147+ to_cleanup: List[Any] = []
148+ with self._lock:
149+ for key in keys:
150+ resource = self._entries.pop(key, None)
151+ if resource is not None:
152+ to_cleanup.append(resource)
153+ for resource in to_cleanup:
154+ torch.npu.synchronize()
155+ if hasattr(resource, "release"):
156+ resource.release()
157+ else:
158+ resource.reset()
159+ 
160+ # ---- debugging ------------------------------------------------------
161+ 
162+ @property
163+ def entry_count(self) -> int:
164+ """Number of currently registered graph resources (for tests)."""
165+ with self._lock:
166+ return len(self._entries)
Mtorch_npu/npu/_graph_tree.py+66-1
@@ -110,6 +110,7 @@ from torch_npu.npu._aclgraph_update_plan.resolver import (
110 validate_aclgraph_update_plan_for_graph,110 validate_aclgraph_update_plan_for_graph,
111)111)
112import torch_npu.npu.aclnn112import torch_npu.npu.aclnn
113+from torch_npu.npu._graph_resource_pool import GraphResourcePool
113 114 
114if TYPE_CHECKING:115if TYPE_CHECKING:
115 from torch._inductor.utils import InputType116 from torch._inductor.utils import InputType
@@ -778,6 +779,7 @@ class NPUGraphNode:
778 self.device = device_index779 self.device = device_index
779 self.stack_traces = stack_traces780 self.stack_traces = stack_traces
780 self.stream = stream781 self.stream = stream
782+ self.function_id: Optional[FunctionID] = None # set when attached to tree
781 783 
782 # Enable re-record a cudagraph when static tensor address changed.784 # Enable re-record a cudagraph when static tensor address changed.
783 # if not we should error when it changed.785 # if not we should error when it changed.
@@ -1507,6 +1509,7 @@ class NPUGraphNode:
1507 1509 
1508 def add_child(self, function_id: FunctionID, node: NPUGraphNode) -> None:1510 def add_child(self, function_id: FunctionID, node: NPUGraphNode) -> None:
1509 "Adds node as a a child of self"1511 "Adds node as a a child of self"
1512+ node.function_id = function_id
1510 self.children[function_id].append(node)1513 self.children[function_id].append(node)
1511 1514 
1512 @staticmethod1515 @staticmethod
@@ -1621,12 +1624,70 @@ class NPUGraphNode:
1621 self.cached_tensor_outputs.clear()1624 self.cached_tensor_outputs.clear()
1622 1625 
1623 for i, unaliased in enumerate(self.unaliased_in_all_paths):1626 for i, unaliased in enumerate(self.unaliased_in_all_paths):
1624- if unaliased:1627+ if unaliased and i < len(self.outputs_weakrefs):
1625 n = self.outputs_weakrefs[i]1628 n = self.outputs_weakrefs[i]
1626 if n is None:1629 if n is None:
1627 raise RuntimeError("check n is not None fail")1630 raise RuntimeError("check n is not None fail")
1628 n.remove_extra_reference()1631 n.remove_extra_reference()
1629 1632 
1633+ def release(self) -> None:
rain-666rain-666
rain-666rain-6666月1日

确保所有资源都被正确释放，添加验证逻辑

likedislike
zhudada0120
zhudada0120
6月5日 评论：
rain-666rain-6666月1日

release方法未说明释放的资源类型和顺序，建议添加资源释放顺序和依赖关系说明

likedislike
zhudada0120
zhudada0120
6月5日 评论：
1634+ """Release all resources held by this graph node and detach from the tree.
1635+ 
1636+ Release order:
1637+ 1. **Detach from tree** — remove self from ``parent.children`` (or from
1638+ ``tree_manager.roots`` for root nodes) so that subsequent replay
1639+ lookups do not encounter a freed node.
1640+ 2. **Cached tensors** — clear all cached tensor references held for
1641+ graph inputs / outputs (``remove_node_cached_tensors``).
1642+ 3. **Output blocks** — collect unaliased output tensor data pointers
1643+ for later recycling via the caching allocator.
1644+ 4. **NPU graph** — call ``self.graph.reset()`` to free the underlying
1645+ device graph resources, then drop the reference.
1646+ 5. **Weakrefs and recording state** — null out ``outputs_weakrefs``,
1647+ ``tensor_weakrefs``, ``recording_outputs``, and
1648+ ``checkpointed_caching_state`` to break reference cycles.
1649+ 6. **Block pointers** — for each collected unaliased output block, call
1650+ ``raw_delete`` on the caching allocator, returning the memory to the
1651+ free pool.
1652+ 
1653+ After this call the node is safe to drop; the NPU graph memory is
1654+ released and the tree structure no longer references this node.
1655+ """
1656+ # Detach from tree structure to avoid dangling references.
1657+ parent = self.parent
1658+ function_id = self.function_id
1659+ if parent is not None:
1660+ if function_id is not None:
1661+ siblings = parent.children.get(function_id, [])
1662+ if self in siblings:
1663+ siblings.remove(self)
1664+ else:
1665+ # Root node: remove from tree manager's roots dict.
1666+ container = get_container(self.device)
1667+ tree_manager = container.tree_manager if container is not None else None
1668+ if tree_manager is not None and tree_manager.roots is not None:
1669+ if function_id is not None:
1670+ root_list = tree_manager.roots.get(function_id, [])
1671+ if self in root_list:
1672+ root_list.remove(self)
1673+ self.remove_node_cached_tensors()
1674+ blocks_to_recycle = []
1675+ for i, unaliased in enumerate(self.unaliased_in_all_paths):
1676+ if unaliased and i < len(self.outputs_weakrefs):
1677+ ref = self.outputs_weakrefs[i]
1678+ if ref is not None and ref():
1679+ blocks_to_recycle.append(ref().data_ptr())
1680+ if self.graph is not None:
1681+ self.graph.reset()
1682+ self.graph = None
1683+ self.cached_tensor_outputs = []
1684+ self.outputs_weakrefs = []
1685+ self.tensor_weakrefs = []
1686+ self.recording_outputs = None
1687+ self.checkpointed_caching_state = None
1688+ for ptr in set(blocks_to_recycle):
1689+ torch_npu._C._npu_npuCachingAllocator_raw_delete(ptr)
1690+ 
1630 def remove_path_cached_tensors(self) -> None:1691 def remove_path_cached_tensors(self) -> None:
1631 for node in self._path_from_root:1692 for node in self._path_from_root:
1632 node.remove_node_cached_tensors()1693 node.remove_node_cached_tensors()
@@ -2254,6 +2315,7 @@ class NPUGraphTreeManager:
2254 self.stream,2315 self.stream,
2255 )2316 )
2256 if self.current_node is None:2317 if self.current_node is None:
2318+ node.function_id = function_id
2257 self.roots[function_id].append(node)2319 self.roots[function_id].append(node)
2258 else:2320 else:
2259 self.current_node.add_child(function_id, node)2321 self.current_node.add_child(function_id, node)
@@ -2262,6 +2324,9 @@ class NPUGraphTreeManager:
2262 self.update_generation()2324 self.update_generation()
2263 log.debug("NPUGRAPH-TREE State state=RECORDING, node=%s, gen=%d",2325 log.debug("NPUGRAPH-TREE State state=RECORDING, node=%s, gen=%d",
2264 graph_id.id, self.current_gen)2326 graph_id.id, self.current_gen)
2327+ pool = GraphResourcePool.get_pool(self.device_index)
L
Lluochao606月22日

该功能默认不开启，注册加入map的动作理论上也是默认不开启才对，保证资源清理的一致性。

likedislike
zhudada0120
zhudada0120
6月23日 评论：
2328+ if pool.is_active():
2329+ pool.register(node)
2265 torch.npu.synchronize()2330 torch.npu.synchronize()
2266 return node.run_first_inputs(new_inputs)2331 return node.run_first_inputs(new_inputs)
2267 2332 
Mtorch_npu/utils/_graph_tree.py+10-0
@@ -197,6 +197,16 @@ def npugraphify_impl(
197 if not isinstance(static_outputs, (list, tuple)):197 if not isinstance(static_outputs, (list, tuple)):
198 static_outputs = (static_outputs,)198 static_outputs = (static_outputs,)
199 199 
200+ # Register the underlying NPUGraph in the device-level resource pool.
201+ from torch_npu.npu._graph_resource_pool import GraphResourcePool
202+ device_index = next(
203+ (inp.device.index for inp in inputs if isinstance(inp, torch.Tensor)),
204+ 0,
205+ )
206+ pool = GraphResourcePool.get_pool(device_index)
207+ if pool.is_active():
208+ pool.register(graph)
209+ 
200 if config.size_asserts:210 if config.size_asserts:
201 211 
202 def run(new_inputs: List[InputType]) -> Callable[[List[InputType]], Any]:212 def run(new_inputs: List[InputType]) -> Callable[[List[InputType]], Any]: