From eb7ce9663f8d3cb039bfa896b59ab8a4dfc772eb Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 16:32:28 -0800 Subject: [PATCH 01/17] Write each CUDA weight once across methods CudaWeightCollector skipped a later method's weight storage only when it compared equal, on the device, to the storage the first method wrote. Every method's preprocess ends with release_moved_tensors, which frees those GPU copies, so by the next method the remembered storage is empty, the size check fails, and the same weights are copied to the host, hashed and written again for every method (about 4 GB per method for a 5-layer Muse Glimmer). Record which CPU tensor each tensor was moved from when the program is moved to the GPU, and remember the source view each storage key was written from. A later method's copy of the same source view is reused without being read or written. The source storage is held weakly, so nothing matches once it is gone, and any other storage still goes through the existing on-device comparison and content-hash check. --- backends/cuda/cuda_backend.py | 13 ++++ backends/cuda/cuda_weight_collector.py | 76 ++++++++++++++++++- backends/cuda/tests/test_cuda_weight_reuse.py | 68 +++++++++++++++++ 3 files changed, 155 insertions(+), 2 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 64f0866a642..0abe0d5d711 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1548,11 +1548,24 @@ def move_program_to_device( """ from torch.export.passes import move_to_device_pass + sources = { + **dict(edge_program.state_dict), + **{ + name: value + for name, value in edge_program.constants.items() + if isinstance(value, torch.Tensor) + }, + } if not cls._is_low_memory_mode(compile_specs): moved = move_to_device_pass(edge_program, device) else: moved = _move_to_device_resize_kv(edge_program, device) _track_moved_tensors(moved) + capture = CudaWeightCollector.active_capture() + if capture is not None: + capture.collector.track_moved_tensors( + sources, {**dict(moved.state_dict), **dict(moved.constants)} + ) return moved @classmethod diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index f826a34d718..546cec1f17f 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -12,6 +12,7 @@ import struct import tempfile import threading +import weakref from dataclasses import dataclass, field, replace from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple @@ -145,6 +146,16 @@ def _same_cuda_storage(first: torch.Tensor, second: torch.Tensor) -> bool: return torch.equal(_storage_bytes(first), _storage_bytes(second)) +def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: + return ( + source.untyped_storage()._cdata, + source.dtype, + source.storage_offset(), + tuple(source.shape), + tuple(source.stride()), + ) + + def _required_view_nbytes( fqn: str, sizes: Tuple[int, ...], @@ -391,6 +402,46 @@ def __init__(self) -> None: self._stored_storages: Dict[str, torch.Tensor] = {} # The key each storage was written under, which ends in its SHA-256. self._content_keys: Dict[str, str] = {} + # The CPU storage each storage key was written from. Each method moves + # the same source weights to the GPU and frees the copies afterwards, + # so a later method's copy of the same source view is recognized by its + # source and not written again. The storage is held weakly: once it is + # gone nothing matches, even if its address is reused. + self._written_sources: Dict[ + str, Tuple["weakref.ref[torch.UntypedStorage]", Tuple[Any, ...]] + ] = {} + # This method's moved storages -> the CPU tensor each was copied from. + self._moved_sources: Dict[int, torch.Tensor] = {} + + def track_moved_tensors( + self, sources: Dict[str, torch.Tensor], moved: Dict[str, torch.Tensor] + ) -> None: + """Record which CPU tensor each tensor of this method was moved from.""" + self._moved_sources = {} + for name, tensor in moved.items(): + source = sources.get(name) + if ( + isinstance(tensor, torch.Tensor) + and isinstance(source, torch.Tensor) + and source is not tensor + and source.device.type == "cpu" + and source.untyped_storage().nbytes() > 0 + and type(source) in (torch.Tensor, torch.nn.Parameter) + and type(tensor) in (torch.Tensor, torch.nn.Parameter) + ): + self._moved_sources[tensor.untyped_storage()._cdata] = source + + def _written_from_same_source(self, storage_key: str, tensor: torch.Tensor) -> bool: + source = self._moved_sources.get(tensor.untyped_storage()._cdata) + written = self._written_sources.get(storage_key) + if source is None or written is None: + return False + written_storage = written[0]() + return ( + written_storage is not None + and written_storage._cdata == source.untyped_storage()._cdata + and written[1] == _source_identity(source) + ) @contextlib.contextmanager def capture( @@ -404,6 +455,10 @@ def capture( finally: self._active.current = previous + @classmethod + def active_capture(cls) -> Optional[_CudaWeightCapture]: + return getattr(cls._active, "current", None) + @classmethod def current_capture(cls) -> _CudaWeightCapture: capture = getattr(cls._active, "current", None) @@ -460,9 +515,14 @@ def materialize( raise RuntimeError(f"Duplicate CUDA FQN weight key for {fqn!r}") # Library-local constants are keyed per shared library, so # only FQN weights can match an earlier method's storage. - shareable = tensor.is_cuda and not _is_aoti_library_local_fqn(fqn) + library_local = _is_aoti_library_local_fqn(fqn) + shareable = tensor.is_cuda and not library_local first = self._stored_storages.get(storage_key) - if ( + if not library_local and self._written_from_same_source( + storage_key, tensor + ): + stored.add(storage_key) + elif ( shareable and first is not None and _same_cuda_storage(first, tensor) @@ -472,6 +532,15 @@ def materialize( storages[storage_key] = _materialize_storage(tensor, directory) if shareable: self._stored_storages.setdefault(storage_key, tensor) + source = self._moved_sources.get(tensor.untyped_storage()._cdata) + if source is not None and not library_local: + self._written_sources.setdefault( + storage_key, + ( + weakref.ref(source.untyped_storage()), + _source_identity(source), + ), + ) sizes = tuple( int(size) for size in getattr(properties, "shape", tensor.shape) @@ -515,6 +584,7 @@ def materialize( ) ) + self._moved_sources = {} trim_host_memory() return CudaWeightArtifact(entries=entries, storages=storages, stored=stored) @@ -639,3 +709,5 @@ def finish(self) -> None: result.data_store_output = shared_output self._stored_storages.clear() self._content_keys.clear() + self._written_sources.clear() + self._moved_sources.clear() diff --git a/backends/cuda/tests/test_cuda_weight_reuse.py b/backends/cuda/tests/test_cuda_weight_reuse.py index 4413ecfd0f5..08c4dcf64bd 100644 --- a/backends/cuda/tests/test_cuda_weight_reuse.py +++ b/backends/cuda/tests/test_cuda_weight_reuse.py @@ -169,3 +169,71 @@ def test_same_cuda_storage_rejects_other_sizes_and_empty_storages(self) -> None: self.assertFalse(same(base, base[:8].clone())) empty = torch.empty(0, dtype=torch.uint8, device="cuda") self.assertTrue(same(empty, empty.clone())) + + def _moved_method(self, collector, so_key, sources): + """A method whose weights are moved from ``sources`` and freed after it.""" + moved = {name: tensor.to("cuda") for name, tensor in sources.items()} + collector.track_moved_tensors(sources, moved) + artifact, result = self._method(collector, so_key, _weights(**moved)) + for tensor in moved.values(): + tensor.untyped_storage().resize_(0) + return artifact, result + + def test_freed_copies_of_one_source_are_written_once(self) -> None: + sources = { + "w": torch.randn(1024, 33), + "b": torch.arange(7, dtype=torch.int8), + } + collector = CudaWeightCollector() + writes = [] + materialize_storage = cuda_weight_collector._materialize_storage + + def counting(tensor, directory): + writes.append(tensor.untyped_storage().nbytes()) + return materialize_storage(tensor, directory) + + with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + first, first_result = self._moved_method(collector, "so-0", sources) + second, second_result = self._moved_method(collector, "so-1", sources) + third, _ = self._moved_method(collector, "so-2", sources) + collector.finish() + self.assertEqual(len(writes), 2) + self.assertEqual(second.storages, {}) + self.assertEqual(third.storages, {}) + self.assertEqual(second.stored, {e.storage_key for e in second.entries}) + self.assertEqual( + decode_cuda_aoti_metadata(second_result.processed_bytes).entries, + decode_cuda_aoti_metadata(first_result.processed_bytes).entries, + ) + stored = { + key: _as_bytes(collector._store.buffers[index]) + for key, index in collector._store.key_to_buffer_idx.items() + } + for entry in decode_cuda_aoti_metadata(first_result.processed_bytes).entries: + name = entry.fqn + self.assertEqual( + stored[entry.storage_key], + sources[name].contiguous().view(torch.uint8).numpy().tobytes(), + ) + + def test_another_source_with_the_same_key_is_written_and_checked(self) -> None: + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", {"w": torch.randn(256, 8)}) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", {"w": torch.randn(256, 8)}) + + def test_a_source_view_with_other_metadata_is_written_again(self) -> None: + base = torch.randn(64, 16) + collector = CudaWeightCollector() + writes = [] + materialize_storage = cuda_weight_collector._materialize_storage + + def counting(tensor, directory): + writes.append(tensor.untyped_storage().nbytes()) + return materialize_storage(tensor, directory) + + with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + self._moved_method(collector, "so-0", {"w": base}) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", {"w": base[:32]}) + self.assertEqual(len(writes), 2) From 29208941f65eb016d2e15038e24d3a6c3e3f0a6b Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 17:17:01 -0800 Subject: [PATCH 02/17] Address review: weak moved-storage refs, low-memory only, per-capture state, e2e test --- backends/cuda/cuda_backend.py | 22 ++- backends/cuda/cuda_weight_collector.py | 161 ++++++++++-------- backends/cuda/tests/test_cuda_weight_reuse.py | 145 +++++++++++++--- 3 files changed, 235 insertions(+), 93 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 0abe0d5d711..ac80e84190b 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1166,7 +1166,10 @@ def materialize_weights_blob( blob_path = os.path.splitext(so_path)[0] + "_weights.blob" capture = CudaWeightCollector.current_capture() capture.artifact = capture.collector.materialize( - weights[0], os.path.dirname(blob_path), capture.device_type_for_weight + weights[0], + os.path.dirname(blob_path), + capture.device_type_for_weight, + capture, ) # Keep AotiBackend's existing path contract intact. The compatibility @@ -1545,9 +1548,19 @@ def move_program_to_device( check on the cache update. They are re-synthesized as zeros for the lifted graph and the serialized blob. This activates automatically with low-memory mode. Other (non-low-memory) exports use the stock pass. + + Low-memory methods share their CPU weights, so the CPU source of each + moved tensor is recorded on the active weight capture. ``materialize`` + then reuses a storage an earlier method already wrote from the same + source view. The moved storages must keep their source's bytes until + ``materialize`` runs. """ from torch.export.passes import move_to_device_pass + if not cls._is_low_memory_mode(compile_specs): + moved = move_to_device_pass(edge_program, device) + _track_moved_tensors(moved) + return moved sources = { **dict(edge_program.state_dict), **{ @@ -1556,14 +1569,11 @@ def move_program_to_device( if isinstance(value, torch.Tensor) }, } - if not cls._is_low_memory_mode(compile_specs): - moved = move_to_device_pass(edge_program, device) - else: - moved = _move_to_device_resize_kv(edge_program, device) + moved = _move_to_device_resize_kv(edge_program, device) _track_moved_tensors(moved) capture = CudaWeightCollector.active_capture() if capture is not None: - capture.collector.track_moved_tensors( + capture.track_moved_tensors( sources, {**dict(moved.state_dict), **dict(moved.constants)} ) return moved diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 546cec1f17f..74a8e58eaa6 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -12,7 +12,6 @@ import struct import tempfile import threading -import weakref from dataclasses import dataclass, field, replace from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple @@ -22,6 +21,7 @@ from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import PreprocessResult from executorch.exir.tensor import scalar_type_enum +from torch.multiprocessing.reductions import StorageWeakRef CUDA_AOTI_METADATA_MAGIC = b"ETCUDAFQN0" @@ -69,11 +69,67 @@ class CudaAotiMetadata: entries: List[CudaWeightEntry] +def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: + return ( + source.untyped_storage()._cdata, + source.dtype, + source.storage_offset(), + tuple(source.shape), + tuple(source.stride()), + ) + + +@dataclass(frozen=True) +class _MovedSource: + moved: StorageWeakRef + source: StorageWeakRef + identity: Tuple[Any, ...] + + @dataclass class _CudaWeightCapture: collector: "CudaWeightCollector" device_type_for_weight: Callable[[torch.Tensor], int] artifact: Optional[CudaWeightArtifact] = None + # This method's moved storage -> the CPU source it was copied from. Both + # are held weakly: a moved storage freed before materialize (and any + # storage later allocated at its address) or a freed source matches nothing. + moved_sources: Dict[int, "_MovedSource"] = field(default_factory=dict) + + def track_moved_tensors( + self, sources: Dict[str, torch.Tensor], moved: Dict[str, torch.Tensor] + ) -> None: + """Record the CPU tensor each tensor of this method was moved from. + + Must run inside ``capture()`` before ``materialize``, and the moved + storages must still hold their source's bytes when it runs: a matched + storage is reused without being read. + """ + self.moved_sources = {} + for name, tensor in moved.items(): + source = sources.get(name) + if ( + isinstance(tensor, torch.Tensor) + and isinstance(source, torch.Tensor) + and source is not tensor + and tensor.is_cuda + and source.device.type == "cpu" + and source.untyped_storage().nbytes() > 0 + and type(source) in (torch.Tensor, torch.nn.Parameter) + and type(tensor) in (torch.Tensor, torch.nn.Parameter) + ): + storage = tensor.untyped_storage() + self.moved_sources[storage._cdata] = _MovedSource( + moved=StorageWeakRef(storage), + source=StorageWeakRef(source.untyped_storage()), + identity=_source_identity(source), + ) + + def source_of(self, tensor: torch.Tensor) -> Optional["_MovedSource"]: + moved = self.moved_sources.get(tensor.untyped_storage()._cdata) + if moved is None or moved.moved.expired() or moved.source.expired(): + return None + return moved def trim_host_memory() -> None: @@ -146,16 +202,6 @@ def _same_cuda_storage(first: torch.Tensor, second: torch.Tensor) -> bool: return torch.equal(_storage_bytes(first), _storage_bytes(second)) -def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: - return ( - source.untyped_storage()._cdata, - source.dtype, - source.storage_offset(), - tuple(source.shape), - tuple(source.stride()), - ) - - def _required_view_nbytes( fqn: str, sizes: Tuple[int, ...], @@ -395,52 +441,29 @@ class CudaWeightCollector: def __init__(self) -> None: self._store = NamedDataStore() self._results: List[PreprocessResult] = [] - # The CUDA storages already written to the store, by storage key. Methods - # share their weights, so a later method's identical storage is compared - # with the first one on the device and not written again. The weights - # must not change until finish() releases them. + # The CUDA storages already written to the store, by storage key. A later + # method's identical storage is compared with the first one on the + # device and not written again while the first one is still allocated. self._stored_storages: Dict[str, torch.Tensor] = {} # The key each storage was written under, which ends in its SHA-256. self._content_keys: Dict[str, str] = {} - # The CPU storage each storage key was written from. Each method moves - # the same source weights to the GPU and frees the copies afterwards, - # so a later method's copy of the same source view is recognized by its - # source and not written again. The storage is held weakly: once it is - # gone nothing matches, even if its address is reused. - self._written_sources: Dict[ - str, Tuple["weakref.ref[torch.UntypedStorage]", Tuple[Any, ...]] - ] = {} - # This method's moved storages -> the CPU tensor each was copied from. - self._moved_sources: Dict[int, torch.Tensor] = {} - - def track_moved_tensors( - self, sources: Dict[str, torch.Tensor], moved: Dict[str, torch.Tensor] - ) -> None: - """Record which CPU tensor each tensor of this method was moved from.""" - self._moved_sources = {} - for name, tensor in moved.items(): - source = sources.get(name) - if ( - isinstance(tensor, torch.Tensor) - and isinstance(source, torch.Tensor) - and source is not tensor - and source.device.type == "cpu" - and source.untyped_storage().nbytes() > 0 - and type(source) in (torch.Tensor, torch.nn.Parameter) - and type(tensor) in (torch.Tensor, torch.nn.Parameter) - ): - self._moved_sources[tensor.untyped_storage()._cdata] = source - - def _written_from_same_source(self, storage_key: str, tensor: torch.Tensor) -> bool: - source = self._moved_sources.get(tensor.untyped_storage()._cdata) + # The CPU source view each storage key was written from. When methods + # share their CPU weights (low_memory_mode), each moves the same source + # to the GPU and frees the copy afterwards, so a later method's copy of + # the same source view is recognized by its source and not written + # again. The source is held weakly: once it is gone nothing matches, + # even if its address is reused. + self._written_sources: Dict[str, Tuple[StorageWeakRef, Tuple[Any, ...]]] = {} + + def _written_from_same_source( + self, storage_key: str, moved: Optional[_MovedSource] + ) -> bool: written = self._written_sources.get(storage_key) - if source is None or written is None: - return False - written_storage = written[0]() return ( - written_storage is not None - and written_storage._cdata == source.untyped_storage()._cdata - and written[1] == _source_identity(source) + moved is not None + and written is not None + and not written[0].expired() + and written[1] == moved.identity ) @contextlib.contextmanager @@ -461,7 +484,7 @@ def active_capture(cls) -> Optional[_CudaWeightCapture]: @classmethod def current_capture(cls) -> _CudaWeightCapture: - capture = getattr(cls._active, "current", None) + capture = cls.active_capture() if capture is None: raise RuntimeError( "CUDA structured weights must be materialized inside preprocess" @@ -473,8 +496,14 @@ def materialize( weights: Any, directory: str, device_type_for_weight: Callable[[torch.Tensor], int], + capture: Optional[_CudaWeightCapture] = None, ) -> CudaWeightArtifact: - """Turn AOTI ``Weights`` into one independently named blob per FQN.""" + """Turn AOTI ``Weights`` into one independently named blob per FQN. + + ``capture`` holds the sources ``move_program_to_device`` recorded for + this method; storages copied from an already written source view are + reused without being read. + """ trim_host_memory() entries: List[CudaWeightEntry] = [] storages: Dict[str, FileBackedData] = {} @@ -518,9 +547,12 @@ def materialize( library_local = _is_aoti_library_local_fqn(fqn) shareable = tensor.is_cuda and not library_local first = self._stored_storages.get(storage_key) - if not library_local and self._written_from_same_source( - storage_key, tensor - ): + moved = ( + capture.source_of(tensor) + if capture is not None and shareable + else None + ) + if self._written_from_same_source(storage_key, moved): stored.add(storage_key) elif ( shareable @@ -532,14 +564,11 @@ def materialize( storages[storage_key] = _materialize_storage(tensor, directory) if shareable: self._stored_storages.setdefault(storage_key, tensor) - source = self._moved_sources.get(tensor.untyped_storage()._cdata) - if source is not None and not library_local: - self._written_sources.setdefault( - storage_key, - ( - weakref.ref(source.untyped_storage()), - _source_identity(source), - ), + written = self._written_sources.get(storage_key) + if moved is not None and (written is None or written[0].expired()): + self._written_sources[storage_key] = ( + moved.source, + moved.identity, ) sizes = tuple( @@ -584,7 +613,6 @@ def materialize( ) ) - self._moved_sources = {} trim_host_memory() return CudaWeightArtifact(entries=entries, storages=storages, stored=stored) @@ -710,4 +738,3 @@ def finish(self) -> None: self._stored_storages.clear() self._content_keys.clear() self._written_sources.clear() - self._moved_sources.clear() diff --git a/backends/cuda/tests/test_cuda_weight_reuse.py b/backends/cuda/tests/test_cuda_weight_reuse.py index 08c4dcf64bd..2fba0893714 100644 --- a/backends/cuda/tests/test_cuda_weight_reuse.py +++ b/backends/cuda/tests/test_cuda_weight_reuse.py @@ -170,21 +170,30 @@ def test_same_cuda_storage_rejects_other_sizes_and_empty_storages(self) -> None: empty = torch.empty(0, dtype=torch.uint8, device="cuda") self.assertTrue(same(empty, empty.clone())) - def _moved_method(self, collector, so_key, sources): + def _moved_method(self, collector, so_key, sources, free_moved_first=False): """A method whose weights are moved from ``sources`` and freed after it.""" moved = {name: tensor.to("cuda") for name, tensor in sources.items()} - collector.track_moved_tensors(sources, moved) - artifact, result = self._method(collector, so_key, _weights(**moved)) - for tensor in moved.values(): + with collector.capture(lambda _: AOTI_DEVICE_TYPE_CUDA) as capture: + capture.track_moved_tensors(sources, moved) + weights = moved + if free_moved_first: + # A pass replaced the moved tensors before AOTI saw them. + weights = {name: tensor.clone() for name, tensor in moved.items()} + del moved + moved = {} + artifact = collector.materialize( + _weights(**weights), + self.directory.name, + lambda _: AOTI_DEVICE_TYPE_CUDA, + capture, + ) + result = _parent_result(so_key) + collector.add_preprocess_result(result, artifact, "cuda") + for tensor in [*moved.values(), *weights.values()]: tensor.untyped_storage().resize_(0) return artifact, result - def test_freed_copies_of_one_source_are_written_once(self) -> None: - sources = { - "w": torch.randn(1024, 33), - "b": torch.arange(7, dtype=torch.int8), - } - collector = CudaWeightCollector() + def _count_writes(self): writes = [] materialize_storage = cuda_weight_collector._materialize_storage @@ -192,10 +201,22 @@ def counting(tensor, directory): writes.append(tensor.untyped_storage().nbytes()) return materialize_storage(tensor, directory) - with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): - first, first_result = self._moved_method(collector, "so-0", sources) - second, second_result = self._moved_method(collector, "so-1", sources) - third, _ = self._moved_method(collector, "so-2", sources) + self.addCleanup(mock.patch.stopall) + mock.patch.object( + cuda_weight_collector, "_materialize_storage", counting + ).start() + return writes + + def test_freed_copies_of_one_source_are_written_once(self) -> None: + sources = { + "w": torch.randn(1024, 33), + "b": torch.arange(7, dtype=torch.int8), + } + collector = CudaWeightCollector() + writes = self._count_writes() + first, first_result = self._moved_method(collector, "so-0", sources) + second, second_result = self._moved_method(collector, "so-1", sources) + third, _ = self._moved_method(collector, "so-2", sources) collector.finish() self.assertEqual(len(writes), 2) self.assertEqual(second.storages, {}) @@ -210,10 +231,9 @@ def counting(tensor, directory): for key, index in collector._store.key_to_buffer_idx.items() } for entry in decode_cuda_aoti_metadata(first_result.processed_bytes).entries: - name = entry.fqn self.assertEqual( stored[entry.storage_key], - sources[name].contiguous().view(torch.uint8).numpy().tobytes(), + sources[entry.fqn].contiguous().view(torch.uint8).numpy().tobytes(), ) def test_another_source_with_the_same_key_is_written_and_checked(self) -> None: @@ -224,7 +244,80 @@ def test_another_source_with_the_same_key_is_written_and_checked(self) -> None: def test_a_source_view_with_other_metadata_is_written_again(self) -> None: base = torch.randn(64, 16) + for first, second in ( + (base, base[:32]), # shape + (base[:32], base[32:]), # storage offset + (base[:, :8], base[:, ::2]), # strides + ): + with self.subTest(second=(second.shape, second.storage_offset())): + collector = CudaWeightCollector() + writes = self._count_writes() + self._moved_method(collector, "so-0", {"w": first}) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", {"w": second}) + self.assertEqual(len(writes), 2) + mock.patch.stopall() + + def test_a_storage_freed_before_materialize_is_not_matched(self) -> None: + sources = {"w": torch.randn(256, 8)} collector = CudaWeightCollector() + writes = self._count_writes() + self._moved_method(collector, "so-0", sources) + # The moved copy is gone before materialize, so whatever storage the + # weights now use must be read and checked, not matched by address. + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method( + collector, "so-1", {"w": torch.randn(256, 8)}, free_moved_first=True + ) + self._moved_method(collector, "so-2", sources, free_moved_first=True) + self.assertEqual(len(writes), 3) + + def test_a_freed_written_source_lets_a_live_source_take_over(self) -> None: + collector = CudaWeightCollector() + writes = self._count_writes() + first = torch.randn(256, 8) + self._moved_method(collector, "so-0", {"w": first}) + shared = first.clone() + del first + self._moved_method(collector, "so-1", {"w": shared}) + self._moved_method(collector, "so-2", {"w": shared}) + self.assertEqual(len(writes), 2) + + +class _TwoMethodLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(64, 48) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + +class CudaWeightReuseExportTest(unittest.TestCase): + """Lowers two methods through CudaBackend, as a real export does.""" + + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA required") + + def _lower(self, low_memory: bool): + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + from executorch.exir.backend.compile_spec_schema import CompileSpec + + module = _TwoMethodLinear().eval() + programs = { + name: torch.export.export(module, (torch.randn(rows, 64),), strict=True) + for name, rows in (("forward_1", 1), ("forward_4", 4)) + } + + def partitioner(name): + specs = [CudaBackend.generate_method_name_compile_spec(name)] + if low_memory: + specs.append(CompileSpec("low_memory_mode", b"ON")) + return [CudaPartitioner(specs)] + writes = [] materialize_storage = cuda_weight_collector._materialize_storage @@ -233,7 +326,19 @@ def counting(tensor, directory): return materialize_storage(tensor, directory) with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): - self._moved_method(collector, "so-0", {"w": base}) - with self.assertRaisesRegex(ValueError, "with different data"): - self._moved_method(collector, "so-1", {"w": base[:32]}) - self.assertEqual(len(writes), 2) + to_edge_transform_and_lower( + programs, + partitioner={name: partitioner(name) for name in programs}, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + return writes + + def test_low_memory_methods_write_shared_weights_once(self) -> None: + weight_bytes = (64 * 48 + 48) * 4 + writes = self._lower(low_memory=True) + self.assertEqual(sum(writes), weight_bytes) + + def test_default_mode_writes_each_methods_own_copy(self) -> None: + weight_bytes = (64 * 48 + 48) * 4 + writes = self._lower(low_memory=False) + self.assertEqual(sum(writes), 2 * weight_bytes) From ed56bc3dd4f2789885d33af226f55cf0c6678f55 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 18:01:57 -0800 Subject: [PATCH 03/17] Address review: version in source identity, pinned weak refs, capture-only materialize, buck dep --- backends/cuda/cuda_backend.py | 24 +-- backends/cuda/cuda_weight_collector.py | 52 ++++--- backends/cuda/tests/targets.bzl | 1 + backends/cuda/tests/test_cuda_partitioner.py | 12 +- backends/cuda/tests/test_cuda_weight_reuse.py | 137 +++++++++--------- 5 files changed, 107 insertions(+), 119 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index ac80e84190b..825c2f7bd3d 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1166,10 +1166,7 @@ def materialize_weights_blob( blob_path = os.path.splitext(so_path)[0] + "_weights.blob" capture = CudaWeightCollector.current_capture() capture.artifact = capture.collector.materialize( - weights[0], - os.path.dirname(blob_path), - capture.device_type_for_weight, - capture, + weights[0], os.path.dirname(blob_path), capture ) # Keep AotiBackend's existing path contract intact. The compatibility @@ -1561,21 +1558,14 @@ def move_program_to_device( moved = move_to_device_pass(edge_program, device) _track_moved_tensors(moved) return moved - sources = { - **dict(edge_program.state_dict), - **{ - name: value - for name, value in edge_program.constants.items() - if isinstance(value, torch.Tensor) - }, - } + # Only parameters reach AOTI's weights as the moved storage: low-memory + # compilation hands AOTI CPU clones of buffers (_compile_time_cpu_clones). + sources = dict(edge_program.state_dict) moved = _move_to_device_resize_kv(edge_program, device) _track_moved_tensors(moved) - capture = CudaWeightCollector.active_capture() - if capture is not None: - capture.track_moved_tensors( - sources, {**dict(moved.state_dict), **dict(moved.constants)} - ) + CudaWeightCollector.current_capture().track_moved_tensors( + sources, dict(moved.state_dict) + ) return moved @classmethod diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 74a8e58eaa6..9380d847283 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -47,7 +47,9 @@ class CudaWeightArtifact: entries: List[CudaWeightEntry] storages: Dict[str, FileBackedData] # Storage keys an earlier method of this collector already stored with the - # same bytes; they are not written again. + # same bytes; they are not written again. A key reused because it was moved + # from the same CPU source view is not re-read: that relies on the shared + # source and the moved storage not changing between methods. stored: Set[str] = field(default_factory=set) @@ -70,8 +72,15 @@ class CudaAotiMetadata: def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: + try: + # Bumped by in-place updates, so a changed source no longer matches. + version = source._version + except RuntimeError: + # Inference tensors cannot be updated in place and have no counter. + version = None return ( source.untyped_storage()._cdata, + version, source.dtype, source.storage_offset(), tuple(source.shape), @@ -91,9 +100,10 @@ class _CudaWeightCapture: collector: "CudaWeightCollector" device_type_for_weight: Callable[[torch.Tensor], int] artifact: Optional[CudaWeightArtifact] = None - # This method's moved storage -> the CPU source it was copied from. Both - # are held weakly: a moved storage freed before materialize (and any - # storage later allocated at its address) or a freed source matches nothing. + # This method's moved storage -> the CPU source it was copied from, keyed by + # storage address. The StorageWeakRefs keep both storage objects (not their + # data) allocated, so no other storage can take either address while this + # capture lives. moved_sources: Dict[int, "_MovedSource"] = field(default_factory=dict) def track_moved_tensors( @@ -115,6 +125,8 @@ def track_moved_tensors( and tensor.is_cuda and source.device.type == "cpu" and source.untyped_storage().nbytes() > 0 + # Low-memory KV placeholders are emptied on the device. + and tensor.untyped_storage().nbytes() > 0 and type(source) in (torch.Tensor, torch.nn.Parameter) and type(tensor) in (torch.Tensor, torch.nn.Parameter) ): @@ -126,10 +138,7 @@ def track_moved_tensors( ) def source_of(self, tensor: torch.Tensor) -> Optional["_MovedSource"]: - moved = self.moved_sources.get(tensor.untyped_storage()._cdata) - if moved is None or moved.moved.expired() or moved.source.expired(): - return None - return moved + return self.moved_sources.get(tensor.untyped_storage()._cdata) def trim_host_memory() -> None: @@ -451,8 +460,9 @@ def __init__(self) -> None: # share their CPU weights (low_memory_mode), each moves the same source # to the GPU and frees the copy afterwards, so a later method's copy of # the same source view is recognized by its source and not written - # again. The source is held weakly: once it is gone nothing matches, - # even if its address is reused. + # again. The StorageWeakRef keeps the source's storage object allocated, + # so its address, part of the identity, is not reused; once its data is + # freed, a later live source takes the key over. self._written_sources: Dict[str, Tuple[StorageWeakRef, Tuple[Any, ...]]] = {} def _written_from_same_source( @@ -460,10 +470,7 @@ def _written_from_same_source( ) -> bool: written = self._written_sources.get(storage_key) return ( - moved is not None - and written is not None - and not written[0].expired() - and written[1] == moved.identity + moved is not None and written is not None and written[1] == moved.identity ) @contextlib.contextmanager @@ -478,13 +485,9 @@ def capture( finally: self._active.current = previous - @classmethod - def active_capture(cls) -> Optional[_CudaWeightCapture]: - return getattr(cls._active, "current", None) - @classmethod def current_capture(cls) -> _CudaWeightCapture: - capture = cls.active_capture() + capture = getattr(cls._active, "current", None) if capture is None: raise RuntimeError( "CUDA structured weights must be materialized inside preprocess" @@ -495,8 +498,7 @@ def materialize( self, weights: Any, directory: str, - device_type_for_weight: Callable[[torch.Tensor], int], - capture: Optional[_CudaWeightCapture] = None, + capture: _CudaWeightCapture, ) -> CudaWeightArtifact: """Turn AOTI ``Weights`` into one independently named blob per FQN. @@ -524,7 +526,7 @@ def materialize( and tuple(tensor.stride()) == tuple(getattr(properties, "stride", ())) ) del storage - device_type = device_type_for_weight(tensor) + device_type = capture.device_type_for_weight(tensor) expected_storage_nbytes = int( getattr(properties, "storage_size", None) or 0 ) @@ -547,11 +549,7 @@ def materialize( library_local = _is_aoti_library_local_fqn(fqn) shareable = tensor.is_cuda and not library_local first = self._stored_storages.get(storage_key) - moved = ( - capture.source_of(tensor) - if capture is not None and shareable - else None - ) + moved = capture.source_of(tensor) if shareable else None if self._written_from_same_source(storage_key, moved): stored.add(storage_key) elif ( diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index f403dfb60ad..dac423e5954 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -383,6 +383,7 @@ def define_common_targets(is_fbcode = False): deps = [ "//caffe2:torch", "//executorch/backends/cuda:cuda_backend", + "//executorch/backends/cuda:cuda_partitioner", "//executorch/exir:lib", ], keep_gpu_sections = True, diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 159fdeeaa1f..337e1f53d72 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -50,9 +50,9 @@ class TestCudaLowMemoryExport(unittest.TestCase): @staticmethod def _materialize(weights, directory, device_type=AOTI_DEVICE_TYPE_CPU): - return CudaWeightCollector().materialize( - weights, directory, lambda _: device_type - ) + collector = CudaWeightCollector() + with collector.capture(lambda _: device_type) as capture: + return collector.materialize(weights, directory, capture) @staticmethod def _parent_result(so_key: str) -> PreprocessResult: @@ -222,9 +222,9 @@ def test_low_memory_cpu_clones_keep_cuda_device_type(self, _) -> None: weights = Weights({"weight": (tensor, TensorProperties(tensor))}) with tempfile.TemporaryDirectory() as directory: - artifact = CudaWeightCollector().materialize( - weights, directory, _aoti_device_type_for_weight - ) + collector = CudaWeightCollector() + with collector.capture(_aoti_device_type_for_weight) as capture: + artifact = collector.materialize(weights, directory, capture) self.assertEqual(artifact.entries[0].device_type, AOTI_DEVICE_TYPE_CUDA) for storage in artifact.storages.values(): storage.close() diff --git a/backends/cuda/tests/test_cuda_weight_reuse.py b/backends/cuda/tests/test_cuda_weight_reuse.py index 2fba0893714..d1a4ba395b4 100644 --- a/backends/cuda/tests/test_cuda_weight_reuse.py +++ b/backends/cuda/tests/test_cuda_weight_reuse.py @@ -6,6 +6,7 @@ """CudaWeightCollector writes a weight storage the methods share only once.""" +import contextlib import tempfile import unittest from unittest import mock @@ -40,6 +41,20 @@ def _as_bytes(buffer) -> bytes: return buffer if isinstance(buffer, bytes) else buffer.to_bytes() +@contextlib.contextmanager +def _counting_writes(): + """Yields the sizes of the storages the collector writes.""" + writes = [] + materialize_storage = cuda_weight_collector._materialize_storage + + def counting(tensor, directory): + writes.append(tensor.untyped_storage().nbytes()) + return materialize_storage(tensor, directory) + + with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + yield writes + + class CudaWeightReuseTest(unittest.TestCase): def setUp(self) -> None: if not torch.cuda.is_available(): @@ -49,9 +64,8 @@ def setUp(self) -> None: def _method(self, collector, so_key, weights): """One method's preprocess: materialize, then add to the shared store.""" - artifact = collector.materialize( - weights, self.directory.name, lambda _: AOTI_DEVICE_TYPE_CUDA - ) + with collector.capture(lambda _: AOTI_DEVICE_TYPE_CUDA) as capture: + artifact = collector.materialize(weights, self.directory.name, capture) result = _parent_result(so_key) collector.add_preprocess_result(result, artifact, "cuda") return artifact, result @@ -60,14 +74,7 @@ def test_a_storage_shared_by_methods_is_written_once(self) -> None: weight = torch.randn(1024, 33, device="cuda") bias = torch.arange(7, dtype=torch.int8, device="cuda") collector = CudaWeightCollector() - writes = [] - materialize_storage = cuda_weight_collector._materialize_storage - - def counting(tensor, directory): - writes.append(tensor.untyped_storage().nbytes()) - return materialize_storage(tensor, directory) - - with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + with _counting_writes() as writes: first, first_result = self._method( collector, "so-0", _weights(w=weight, b=bias) ) @@ -170,22 +177,23 @@ def test_same_cuda_storage_rejects_other_sizes_and_empty_storages(self) -> None: empty = torch.empty(0, dtype=torch.uint8, device="cuda") self.assertTrue(same(empty, empty.clone())) - def _moved_method(self, collector, so_key, sources, free_moved_first=False): - """A method whose weights are moved from ``sources`` and freed after it.""" + def _moved_method(self, collector, so_key, sources, replace=None): + """A method whose weights are moved from ``sources`` and freed after it. + + ``replace`` frees the moved copies before AOTI sees them (as a pass + replacing them would) and returns the tensors to materialize instead. + """ moved = {name: tensor.to("cuda") for name, tensor in sources.items()} with collector.capture(lambda _: AOTI_DEVICE_TYPE_CUDA) as capture: capture.track_moved_tensors(sources, moved) weights = moved - if free_moved_first: - # A pass replaced the moved tensors before AOTI saw them. - weights = {name: tensor.clone() for name, tensor in moved.items()} - del moved - moved = {} + if replace is not None: + shapes = {name: tensor.shape for name, tensor in moved.items()} + for tensor in moved.values(): + tensor.untyped_storage().resize_(0) + weights = replace(shapes) artifact = collector.materialize( - _weights(**weights), - self.directory.name, - lambda _: AOTI_DEVICE_TYPE_CUDA, - capture, + _weights(**weights), self.directory.name, capture ) result = _parent_result(so_key) collector.add_preprocess_result(result, artifact, "cuda") @@ -193,30 +201,16 @@ def _moved_method(self, collector, so_key, sources, free_moved_first=False): tensor.untyped_storage().resize_(0) return artifact, result - def _count_writes(self): - writes = [] - materialize_storage = cuda_weight_collector._materialize_storage - - def counting(tensor, directory): - writes.append(tensor.untyped_storage().nbytes()) - return materialize_storage(tensor, directory) - - self.addCleanup(mock.patch.stopall) - mock.patch.object( - cuda_weight_collector, "_materialize_storage", counting - ).start() - return writes - def test_freed_copies_of_one_source_are_written_once(self) -> None: sources = { "w": torch.randn(1024, 33), "b": torch.arange(7, dtype=torch.int8), } collector = CudaWeightCollector() - writes = self._count_writes() - first, first_result = self._moved_method(collector, "so-0", sources) - second, second_result = self._moved_method(collector, "so-1", sources) - third, _ = self._moved_method(collector, "so-2", sources) + with _counting_writes() as writes: + first, first_result = self._moved_method(collector, "so-0", sources) + second, second_result = self._moved_method(collector, "so-1", sources) + third, _ = self._moved_method(collector, "so-2", sources) collector.finish() self.assertEqual(len(writes), 2) self.assertEqual(second.storages, {}) @@ -251,36 +245,47 @@ def test_a_source_view_with_other_metadata_is_written_again(self) -> None: ): with self.subTest(second=(second.shape, second.storage_offset())): collector = CudaWeightCollector() - writes = self._count_writes() - self._moved_method(collector, "so-0", {"w": first}) - with self.assertRaisesRegex(ValueError, "with different data"): - self._moved_method(collector, "so-1", {"w": second}) + with _counting_writes() as writes: + self._moved_method(collector, "so-0", {"w": first}) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", {"w": second}) self.assertEqual(len(writes), 2) - mock.patch.stopall() + + def test_a_source_updated_in_place_is_written_and_checked(self) -> None: + sources = {"w": torch.nn.Parameter(torch.randn(256, 8))} + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", sources) + with torch.no_grad(): + sources["w"].add_(1) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", sources) def test_a_storage_freed_before_materialize_is_not_matched(self) -> None: sources = {"w": torch.randn(256, 8)} collector = CudaWeightCollector() - writes = self._count_writes() self._moved_method(collector, "so-0", sources) - # The moved copy is gone before materialize, so whatever storage the - # weights now use must be read and checked, not matched by address. + # Same source, but the moved copy is freed before AOTI sees it and a + # storage with other bytes takes its place: it must be read and checked. with self.assertRaisesRegex(ValueError, "with different data"): self._moved_method( - collector, "so-1", {"w": torch.randn(256, 8)}, free_moved_first=True + collector, + "so-1", + sources, + replace=lambda shapes: { + name: torch.randn(shape, device="cuda") + for name, shape in shapes.items() + }, ) - self._moved_method(collector, "so-2", sources, free_moved_first=True) - self.assertEqual(len(writes), 3) def test_a_freed_written_source_lets_a_live_source_take_over(self) -> None: collector = CudaWeightCollector() - writes = self._count_writes() first = torch.randn(256, 8) - self._moved_method(collector, "so-0", {"w": first}) - shared = first.clone() - del first - self._moved_method(collector, "so-1", {"w": shared}) - self._moved_method(collector, "so-2", {"w": shared}) + with _counting_writes() as writes: + self._moved_method(collector, "so-0", {"w": first}) + shared = first.clone() + del first + self._moved_method(collector, "so-1", {"w": shared}) + self._moved_method(collector, "so-2", {"w": shared}) self.assertEqual(len(writes), 2) @@ -318,14 +323,7 @@ def partitioner(name): specs.append(CompileSpec("low_memory_mode", b"ON")) return [CudaPartitioner(specs)] - writes = [] - materialize_storage = cuda_weight_collector._materialize_storage - - def counting(tensor, directory): - writes.append(tensor.untyped_storage().nbytes()) - return materialize_storage(tensor, directory) - - with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + with _counting_writes() as writes: to_edge_transform_and_lower( programs, partitioner={name: partitioner(name) for name in programs}, @@ -333,12 +331,13 @@ def counting(tensor, directory): ) return writes + # The linear weight and its bias. + _SIZES = (64 * 48 * 4, 48 * 4) + def test_low_memory_methods_write_shared_weights_once(self) -> None: - weight_bytes = (64 * 48 + 48) * 4 writes = self._lower(low_memory=True) - self.assertEqual(sum(writes), weight_bytes) + self.assertEqual([writes.count(size) for size in self._SIZES], [1, 1]) def test_default_mode_writes_each_methods_own_copy(self) -> None: - weight_bytes = (64 * 48 + 48) * 4 writes = self._lower(low_memory=False) - self.assertEqual(sum(writes), 2 * weight_bytes) + self.assertEqual([writes.count(size) for size in self._SIZES], [2, 2]) From 5efe1bf074764f80cc6b4ce95659eb4f56f87437 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 19:17:37 -0800 Subject: [PATCH 04/17] Address review: check the moved copy's version, parameters only, simpler bookkeeping --- backends/cuda/cuda_backend.py | 11 ++- backends/cuda/cuda_weight_collector.py | 96 +++++++++---------- backends/cuda/tests/test_cuda_weight_reuse.py | 48 +++++++--- 3 files changed, 87 insertions(+), 68 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 825c2f7bd3d..3e9097838a8 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1558,12 +1558,15 @@ def move_program_to_device( moved = move_to_device_pass(edge_program, device) _track_moved_tensors(moved) return moved - # Only parameters reach AOTI's weights as the moved storage: low-memory - # compilation hands AOTI CPU clones of buffers (_compile_time_cpu_clones). - sources = dict(edge_program.state_dict) + # Only parameters are recorded: low-memory compilation can hand AOTI CPU + # clones of buffers (_compile_time_cpu_clones). + sources = { + name: edge_program.state_dict[name] + for name in edge_program.graph_signature.parameters + } moved = _move_to_device_resize_kv(edge_program, device) _track_moved_tensors(moved) - CudaWeightCollector.current_capture().track_moved_tensors( + CudaWeightCollector.current_capture().record_moved_sources( sources, dict(moved.state_dict) ) return moved diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 9380d847283..c36676a33da 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -48,8 +48,8 @@ class CudaWeightArtifact: storages: Dict[str, FileBackedData] # Storage keys an earlier method of this collector already stored with the # same bytes; they are not written again. A key reused because it was moved - # from the same CPU source view is not re-read: that relies on the shared - # source and the moved storage not changing between methods. + # from the same CPU source view is not re-read: that relies on the source + # not changing between methods in ways its version counter misses. stored: Set[str] = field(default_factory=set) @@ -72,15 +72,11 @@ class CudaAotiMetadata: def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: - try: - # Bumped by in-place updates, so a changed source no longer matches. - version = source._version - except RuntimeError: - # Inference tensors cannot be updated in place and have no counter. - version = None + # Bumped by in-place ops autograd tracks; writes through .data, .numpy() or + # the storage do not bump it. return ( source.untyped_storage()._cdata, - version, + source._version, source.dtype, source.storage_offset(), tuple(source.shape), @@ -90,9 +86,11 @@ def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: @dataclass(frozen=True) class _MovedSource: - moved: StorageWeakRef source: StorageWeakRef identity: Tuple[Any, ...] + # The moved tensor's version when it was recorded; a write to it before + # materialize makes it be read and checked again. + moved_version: int @dataclass @@ -100,46 +98,45 @@ class _CudaWeightCapture: collector: "CudaWeightCollector" device_type_for_weight: Callable[[torch.Tensor], int] artifact: Optional[CudaWeightArtifact] = None - # This method's moved storage -> the CPU source it was copied from, keyed by - # storage address. The StorageWeakRefs keep both storage objects (not their - # data) allocated, so no other storage can take either address while this - # capture lives. - moved_sources: Dict[int, "_MovedSource"] = field(default_factory=dict) + # This method's moved storage -> the CPU source it was copied from. The + # StorageWeakRefs (which hash and compare by address) keep both storage + # objects, not their data, allocated, so no other storage can take either + # address while this capture lives. + moved_sources: Dict[StorageWeakRef, "_MovedSource"] = field(default_factory=dict) - def track_moved_tensors( + def record_moved_sources( self, sources: Dict[str, torch.Tensor], moved: Dict[str, torch.Tensor] ) -> None: """Record the CPU tensor each tensor of this method was moved from. - Must run inside ``capture()`` before ``materialize``, and the moved - storages must still hold their source's bytes when it runs: a matched - storage is reused without being read. + Must run inside ``capture()`` before ``materialize``. A matched storage + is reused without being read, so nothing may write the moved tensors + between this call and ``materialize`` except through in-place ops + their version counter records. """ self.moved_sources = {} for name, tensor in moved.items(): source = sources.get(name) if ( - isinstance(tensor, torch.Tensor) - and isinstance(source, torch.Tensor) - and source is not tensor + type(source) in (torch.Tensor, torch.nn.Parameter) + and type(tensor) in (torch.Tensor, torch.nn.Parameter) and tensor.is_cuda and source.device.type == "cpu" + # Weights loaded under inference_mode can change without a + # version bump. + and not source.is_inference() and source.untyped_storage().nbytes() > 0 # Low-memory KV placeholders are emptied on the device. and tensor.untyped_storage().nbytes() > 0 - and type(source) in (torch.Tensor, torch.nn.Parameter) - and type(tensor) in (torch.Tensor, torch.nn.Parameter) ): - storage = tensor.untyped_storage() - self.moved_sources[storage._cdata] = _MovedSource( - moved=StorageWeakRef(storage), - source=StorageWeakRef(source.untyped_storage()), - identity=_source_identity(source), + self.moved_sources[StorageWeakRef(tensor.untyped_storage())] = ( + _MovedSource( + source=StorageWeakRef(source.untyped_storage()), + identity=_source_identity(source), + moved_version=tensor._version, + ) ) - def source_of(self, tensor: torch.Tensor) -> Optional["_MovedSource"]: - return self.moved_sources.get(tensor.untyped_storage()._cdata) - def trim_host_memory() -> None: gc.collect() @@ -463,15 +460,7 @@ def __init__(self) -> None: # again. The StorageWeakRef keeps the source's storage object allocated, # so its address, part of the identity, is not reused; once its data is # freed, a later live source takes the key over. - self._written_sources: Dict[str, Tuple[StorageWeakRef, Tuple[Any, ...]]] = {} - - def _written_from_same_source( - self, storage_key: str, moved: Optional[_MovedSource] - ) -> bool: - written = self._written_sources.get(storage_key) - return ( - moved is not None and written is not None and written[1] == moved.identity - ) + self._written_sources: Dict[str, _MovedSource] = {} @contextlib.contextmanager def capture( @@ -549,8 +538,19 @@ def materialize( library_local = _is_aoti_library_local_fqn(fqn) shareable = tensor.is_cuda and not library_local first = self._stored_storages.get(storage_key) - moved = capture.source_of(tensor) if shareable else None - if self._written_from_same_source(storage_key, moved): + moved = ( + capture.moved_sources.get(StorageWeakRef(tensor.untyped_storage())) + if shareable + else None + ) + if moved is not None and tensor._version != moved.moved_version: + moved = None + written = self._written_sources.get(storage_key) + if ( + moved is not None + and written is not None + and written.identity == moved.identity + ): stored.add(storage_key) elif ( shareable @@ -562,12 +562,10 @@ def materialize( storages[storage_key] = _materialize_storage(tensor, directory) if shareable: self._stored_storages.setdefault(storage_key, tensor) - written = self._written_sources.get(storage_key) - if moved is not None and (written is None or written[0].expired()): - self._written_sources[storage_key] = ( - moved.source, - moved.identity, - ) + if moved is not None and ( + written is None or written.source.expired() + ): + self._written_sources[storage_key] = moved sizes = tuple( int(size) for size in getattr(properties, "shape", tensor.shape) diff --git a/backends/cuda/tests/test_cuda_weight_reuse.py b/backends/cuda/tests/test_cuda_weight_reuse.py index d1a4ba395b4..7e1fadc0af8 100644 --- a/backends/cuda/tests/test_cuda_weight_reuse.py +++ b/backends/cuda/tests/test_cuda_weight_reuse.py @@ -7,6 +7,7 @@ """CudaWeightCollector writes a weight storage the methods share only once.""" import contextlib +import gc import tempfile import unittest from unittest import mock @@ -177,27 +178,32 @@ def test_same_cuda_storage_rejects_other_sizes_and_empty_storages(self) -> None: empty = torch.empty(0, dtype=torch.uint8, device="cuda") self.assertTrue(same(empty, empty.clone())) - def _moved_method(self, collector, so_key, sources, replace=None): + def _moved_method(self, collector, so_key, sources, replace=None, update=None): """A method whose weights are moved from ``sources`` and freed after it. - ``replace`` frees the moved copies before AOTI sees them (as a pass + ``replace`` drops the moved copies before AOTI sees them (as a pass replacing them would) and returns the tensors to materialize instead. + ``update`` writes to the moved copies before AOTI sees them. """ moved = {name: tensor.to("cuda") for name, tensor in sources.items()} with collector.capture(lambda _: AOTI_DEVICE_TYPE_CUDA) as capture: - capture.track_moved_tensors(sources, moved) - weights = moved + capture.record_moved_sources(sources, moved) + if update is not None: + with torch.no_grad(): + for tensor in moved.values(): + update(tensor) if replace is not None: shapes = {name: tensor.shape for name, tensor in moved.items()} - for tensor in moved.values(): - tensor.untyped_storage().resize_(0) - weights = replace(shapes) + # Only the capture's records may keep the storages alive. + moved.clear() + gc.collect() + moved = replace(shapes) artifact = collector.materialize( - _weights(**weights), self.directory.name, capture + _weights(**moved), self.directory.name, capture ) result = _parent_result(so_key) collector.add_preprocess_result(result, artifact, "cuda") - for tensor in [*moved.values(), *weights.values()]: + for tensor in moved.values(): tensor.untyped_storage().resize_(0) return artifact, result @@ -260,12 +266,22 @@ def test_a_source_updated_in_place_is_written_and_checked(self) -> None: with self.assertRaisesRegex(ValueError, "with different data"): self._moved_method(collector, "so-1", sources) + def test_a_moved_copy_written_before_materialize_is_checked(self) -> None: + sources = {"w": torch.randn(256, 8)} + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", sources) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method( + collector, "so-1", sources, update=lambda tensor: tensor.add_(1) + ) + def test_a_storage_freed_before_materialize_is_not_matched(self) -> None: sources = {"w": torch.randn(256, 8)} collector = CudaWeightCollector() self._moved_method(collector, "so-0", sources) # Same source, but the moved copy is freed before AOTI sees it and a - # storage with other bytes takes its place: it must be read and checked. + # storage with other bytes is allocated, possibly at its old address: + # it must be read and checked. with self.assertRaisesRegex(ValueError, "with different data"): self._moved_method( collector, @@ -293,9 +309,11 @@ class _TwoMethodLinear(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear = torch.nn.Linear(64, 48) + # A persistent buffer keeps the per-method handling. + self.register_buffer("scale", (torch.rand(48) + 0.5).to(torch.bfloat16)) def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.linear(x) + return self.linear(x) * self.scale class CudaWeightReuseExportTest(unittest.TestCase): @@ -331,13 +349,13 @@ def partitioner(name): ) return writes - # The linear weight and its bias. - _SIZES = (64 * 48 * 4, 48 * 4) + # The linear weight, its bias and the bf16 buffer. + _SIZES = (64 * 48 * 4, 48 * 4, 48 * 2) def test_low_memory_methods_write_shared_weights_once(self) -> None: writes = self._lower(low_memory=True) - self.assertEqual([writes.count(size) for size in self._SIZES], [1, 1]) + self.assertEqual([writes.count(size) for size in self._SIZES], [1, 1, 2]) def test_default_mode_writes_each_methods_own_copy(self) -> None: writes = self._lower(low_memory=False) - self.assertEqual([writes.count(size) for size in self._SIZES], [2, 2]) + self.assertEqual([writes.count(size) for size in self._SIZES[:2]], [2, 2]) From 9e61b7a2d12a97766e086615286f09b41f567c91 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 16:32:28 -0800 Subject: [PATCH 05/17] Read the CUDA cache counters once per autotune capture After every CUDA-graph timing capture, _release_cached_memory_if_large called torch.cuda.memory_reserved() and torch.cuda.memory_allocated(). Each flattens the whole allocator stat tree into a Python dict, and that tree keeps one entry per CUDA-graph private pool ever captured, so the check gets slower with every capture of the export. Profiling a 5-layer Muse Glimmer batching export put about 80% of the candidate-benchmark time of forward_8..forward_64 (214 s of 267 s for forward_64) and nearly all of draft_forward's in this check. Read reserved_bytes.all.current and allocated_bytes.all.current from one nested stats call. The cached-bytes value and the empty_cache threshold are unchanged, so the timed candidates and the picks are the same; a probe with 2048 captured pools gave the same bytes about 20x faster. --- backends/cuda/autotune/cuda_graph_timing.py | 13 +++++++++---- backends/cuda/tests/test_cuda_graph_timing.py | 13 +++++++++++++ 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/backends/cuda/autotune/cuda_graph_timing.py b/backends/cuda/autotune/cuda_graph_timing.py index 66a969ea8f2..a50612dd538 100644 --- a/backends/cuda/autotune/cuda_graph_timing.py +++ b/backends/cuda/autotune/cuda_graph_timing.py @@ -117,10 +117,15 @@ def _default_cuda_bench( def _release_cached_memory_if_large() -> None: - if ( - torch.cuda.memory_reserved() - torch.cuda.memory_allocated() - > _EMPTY_CACHE_THRESHOLD_BYTES - ): + # memory_reserved() and memory_allocated() each flatten every allocator + # stat, including one entry per CUDA-graph private pool ever captured, so + # their cost grows with every capture. Read the two counters once instead. + stats = torch._C._cuda_memoryStats(torch.cuda.current_device()) + cached = ( + stats["reserved_bytes"]["all"]["current"] + - stats["allocated_bytes"]["all"]["current"] + ) + if cached > _EMPTY_CACHE_THRESHOLD_BYTES: torch.cuda.empty_cache() diff --git a/backends/cuda/tests/test_cuda_graph_timing.py b/backends/cuda/tests/test_cuda_graph_timing.py index 5cf2f4b5f75..217e5e8d39b 100644 --- a/backends/cuda/tests/test_cuda_graph_timing.py +++ b/backends/cuda/tests/test_cuda_graph_timing.py @@ -153,6 +153,19 @@ def test_compile_spec_turns_it_off(self) -> None: self.assertIs(benchmarking._BENCHMARK_DISPATCH.get("cuda"), before) self.assertIs(benchmarking._BENCHMARK_DISPATCH.get("cuda"), before) + def test_cached_memory_check_matches_the_flattened_counters(self) -> None: + x = torch.zeros(1 << 16, device="cuda") + with cgt.cuda_graph_timing(): + for _ in range(8): + benchmarking.benchmarker.benchmark(lambda: x.add_(1), device="cuda") + flattened = torch.cuda.memory_reserved() - torch.cuda.memory_allocated() + for threshold, empties in ((flattened - 1, 1), (flattened, 0)): + with self.subTest(threshold=threshold), mock.patch.object( + cgt, "_EMPTY_CACHE_THRESHOLD_BYTES", threshold + ), mock.patch.object(torch.cuda, "empty_cache") as empty_cache: + cgt._release_cached_memory_if_large() + self.assertEqual(empty_cache.call_count, empties) + # --------------------------------------------------------------------------- # End to end: export while the CPU is saturated, and check every pick. From cac43b8357ded7804def8303c737f401f3f4349c Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 17:17:02 -0800 Subject: [PATCH 06/17] Address review: public nested stats API, test fails on the flattening path --- backends/cuda/autotune/cuda_graph_timing.py | 17 +++++++++++------ backends/cuda/tests/test_cuda_graph_timing.py | 13 ++++++++++++- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/backends/cuda/autotune/cuda_graph_timing.py b/backends/cuda/autotune/cuda_graph_timing.py index a50612dd538..dbf138f8cd3 100644 --- a/backends/cuda/autotune/cuda_graph_timing.py +++ b/backends/cuda/autotune/cuda_graph_timing.py @@ -31,8 +31,9 @@ captured). To keep captures succeeding: -- each capture's graph and its memory pool are released right after timing, - and the allocator's cache is emptied once it holds a lot of unused memory; +- each capture's graph is released right after timing, which leaves its + memory pool to be freed the next time the allocator's cache is emptied; the + cache is emptied once it holds a lot of unused memory; - on CUDA OOM the cache is emptied and the capture retried; if that fails too, the tensors given as ``offload`` (the model's weights moved to the GPU for the compile) are parked on the CPU while the candidate is timed and restored @@ -117,10 +118,14 @@ def _default_cuda_bench( def _release_cached_memory_if_large() -> None: - # memory_reserved() and memory_allocated() each flatten every allocator - # stat, including one entry per CUDA-graph private pool ever captured, so - # their cost grows with every capture. Read the two counters once instead. - stats = torch._C._cuda_memoryStats(torch.cuda.current_device()) + # The allocator stats keep an entry for every CUDA-graph private pool not + # yet destroyed, and memory_reserved() and memory_allocated() each flatten + # and sort all of them in Python. Read the two counters from one nested + # call instead; it still builds the per-pool entries, so it grows with the + # pool count too, but far more slowly. + stats = torch.cuda.memory_stats_as_nested_dict() + if not stats: + return cached = ( stats["reserved_bytes"]["all"]["current"] - stats["allocated_bytes"]["all"]["current"] diff --git a/backends/cuda/tests/test_cuda_graph_timing.py b/backends/cuda/tests/test_cuda_graph_timing.py index 217e5e8d39b..4fef9189adc 100644 --- a/backends/cuda/tests/test_cuda_graph_timing.py +++ b/backends/cuda/tests/test_cuda_graph_timing.py @@ -9,6 +9,7 @@ python -m pytest backends/cuda/tests/test_cuda_graph_timing.py -v """ +import gc import statistics import struct import unittest @@ -158,11 +159,21 @@ def test_cached_memory_check_matches_the_flattened_counters(self) -> None: with cgt.cuda_graph_timing(): for _ in range(8): benchmarking.benchmarker.benchmark(lambda: x.add_(1), device="cuda") + # Free what earlier tests' mock reference cycles still hold, so no block + # is released between reading the counters and the zero-margin check. + gc.collect() flattened = torch.cuda.memory_reserved() - torch.cuda.memory_allocated() for threshold, empties in ((flattened - 1, 1), (flattened, 0)): with self.subTest(threshold=threshold), mock.patch.object( cgt, "_EMPTY_CACHE_THRESHOLD_BYTES", threshold - ), mock.patch.object(torch.cuda, "empty_cache") as empty_cache: + ), mock.patch.object( + torch.cuda, "empty_cache" + ) as empty_cache, mock.patch.object( + # memory_reserved() and memory_allocated() flatten through this. + torch.cuda.memory, + "memory_stats", + side_effect=AssertionError("flattened the allocator stats"), + ): cgt._release_cached_memory_if_large() self.assertEqual(empty_cache.call_count, empties) From 8d0488cc2fc32c5c8a717998f10c3331376fda53 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 18:39:57 -0800 Subject: [PATCH 07/17] Address review: per-pool entries outlive their pools; tighten the flattening guard --- backends/cuda/autotune/cuda_graph_timing.py | 19 +++++++++---------- backends/cuda/tests/test_cuda_graph_timing.py | 7 ++++++- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/backends/cuda/autotune/cuda_graph_timing.py b/backends/cuda/autotune/cuda_graph_timing.py index dbf138f8cd3..0ca6783c6ef 100644 --- a/backends/cuda/autotune/cuda_graph_timing.py +++ b/backends/cuda/autotune/cuda_graph_timing.py @@ -31,9 +31,9 @@ captured). To keep captures succeeding: -- each capture's graph is released right after timing, which leaves its - memory pool to be freed the next time the allocator's cache is emptied; the - cache is emptied once it holds a lot of unused memory; +- each capture's graph is released right after timing (``torch.cuda.graph`` + also empties the allocator's cache before every capture), and the cache is + emptied once it holds a lot of unused memory; - on CUDA OOM the cache is emptied and the capture retried; if that fails too, the tensors given as ``offload`` (the model's weights moved to the GPU for the compile) are parked on the CPU while the candidate is timed and restored @@ -118,14 +118,13 @@ def _default_cuda_bench( def _release_cached_memory_if_large() -> None: - # The allocator stats keep an entry for every CUDA-graph private pool not - # yet destroyed, and memory_reserved() and memory_allocated() each flatten - # and sort all of them in Python. Read the two counters from one nested - # call instead; it still builds the per-pool entries, so it grows with the - # pool count too, but far more slowly. + # The allocator stats keep an entry for every CUDA-graph private pool ever + # captured, even after the pool is released and the cache emptied, and + # memory_reserved() and memory_allocated() each flatten and sort all of + # them in Python. Read the two counters from one nested call instead; it + # still builds the per-pool entries, so it grows with the capture count + # too, but far more slowly. stats = torch.cuda.memory_stats_as_nested_dict() - if not stats: - return cached = ( stats["reserved_bytes"]["all"]["current"] - stats["allocated_bytes"]["all"]["current"] diff --git a/backends/cuda/tests/test_cuda_graph_timing.py b/backends/cuda/tests/test_cuda_graph_timing.py index 4fef9189adc..a3cc46f02a2 100644 --- a/backends/cuda/tests/test_cuda_graph_timing.py +++ b/backends/cuda/tests/test_cuda_graph_timing.py @@ -156,9 +156,10 @@ def test_compile_spec_turns_it_off(self) -> None: def test_cached_memory_check_matches_the_flattened_counters(self) -> None: x = torch.zeros(1 << 16, device="cuda") - with cgt.cuda_graph_timing(): + with cgt.cuda_graph_timing() as stats: for _ in range(8): benchmarking.benchmarker.benchmark(lambda: x.add_(1), device="cuda") + self.assertEqual(stats.captured, 8) # Free what earlier tests' mock reference cycles still hold, so no block # is released between reading the counters and the zero-margin check. gc.collect() @@ -173,6 +174,10 @@ def test_cached_memory_check_matches_the_flattened_counters(self) -> None: torch.cuda.memory, "memory_stats", side_effect=AssertionError("flattened the allocator stats"), + ), mock.patch.object( + torch.cuda, + "memory_stats", + side_effect=AssertionError("flattened the allocator stats"), ): cgt._release_cached_memory_if_large() self.assertEqual(empty_cache.call_count, empties) From 88b61efc8c018da32c34b52aa0e6694a51267870 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Fri, 13 Nov 2026 19:17:37 -0800 Subject: [PATCH 08/17] Address review: count the nested stats reads in the test --- backends/cuda/tests/test_cuda_graph_timing.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/backends/cuda/tests/test_cuda_graph_timing.py b/backends/cuda/tests/test_cuda_graph_timing.py index a3cc46f02a2..9dba3fa9e57 100644 --- a/backends/cuda/tests/test_cuda_graph_timing.py +++ b/backends/cuda/tests/test_cuda_graph_timing.py @@ -170,16 +170,14 @@ def test_cached_memory_check_matches_the_flattened_counters(self) -> None: ), mock.patch.object( torch.cuda, "empty_cache" ) as empty_cache, mock.patch.object( - # memory_reserved() and memory_allocated() flatten through this. - torch.cuda.memory, - "memory_stats", - side_effect=AssertionError("flattened the allocator stats"), - ), mock.patch.object( + # The flattening APIs reach the nested stats through + # torch.cuda.memory's own global, so they never count here. torch.cuda, - "memory_stats", - side_effect=AssertionError("flattened the allocator stats"), - ): + "memory_stats_as_nested_dict", + wraps=torch.cuda.memory_stats_as_nested_dict, + ) as nested: cgt._release_cached_memory_if_large() + self.assertEqual(nested.call_count, 1) self.assertEqual(empty_cache.call_count, empties) From 9624e444ffc3b18b06e1a9f0e200918e4aabc887 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Sat, 14 Nov 2026 00:04:09 -0800 Subject: [PATCH 09/17] Address review: empty the cache after every capture instead of reading the stats --- backends/cuda/autotune/cuda_graph_timing.py | 27 ++++----------- backends/cuda/tests/test_cuda_graph_timing.py | 33 +++++++------------ 2 files changed, 17 insertions(+), 43 deletions(-) diff --git a/backends/cuda/autotune/cuda_graph_timing.py b/backends/cuda/autotune/cuda_graph_timing.py index 0ca6783c6ef..bba332a48ae 100644 --- a/backends/cuda/autotune/cuda_graph_timing.py +++ b/backends/cuda/autotune/cuda_graph_timing.py @@ -31,9 +31,11 @@ captured). To keep captures succeeding: -- each capture's graph is released right after timing (``torch.cuda.graph`` - also empties the allocator's cache before every capture), and the cache is - emptied once it holds a lot of unused memory; +- each capture's graph is released and the allocator's cache emptied right + after timing. The cache is not checked first: the allocator stats keep an + entry for every private pool ever captured, so reading them costs more with + each capture, while emptying a cache that holds little more than the + released pool is cheap; - on CUDA OOM the cache is emptied and the capture retried; if that fails too, the tensors given as ``offload`` (the model's weights moved to the GPU for the compile) are parked on the CPU while the candidate is timed and restored @@ -60,7 +62,6 @@ _TARGET_GRAPH_US = 2000.0 _MIN_CALLS = 4 _MAX_CALLS = 32 -_EMPTY_CACHE_THRESHOLD_BYTES = 1 << 30 @dataclass @@ -117,22 +118,6 @@ def _default_cuda_bench( return bench.benchmark_gpu(fn, warmup=warmup, rep=rep, **kwargs) -def _release_cached_memory_if_large() -> None: - # The allocator stats keep an entry for every CUDA-graph private pool ever - # captured, even after the pool is released and the cache emptied, and - # memory_reserved() and memory_allocated() each flatten and sort all of - # them in Python. Read the two counters from one nested call instead; it - # still builds the per-pool entries, so it grows with the capture count - # too, but far more slowly. - stats = torch.cuda.memory_stats_as_nested_dict() - cached = ( - stats["reserved_bytes"]["all"]["current"] - - stats["allocated_bytes"]["all"]["current"] - ) - if cached > _EMPTY_CACHE_THRESHOLD_BYTES: - torch.cuda.empty_cache() - - _l2_flush_buffers: dict[int, torch.Tensor] = {} @@ -189,7 +174,7 @@ def time_with_cuda_graph_us(fn: Callable[[], Any]) -> float: with contextlib.suppress(RuntimeError): graph.reset() del graph - _release_cached_memory_if_large() + torch.cuda.empty_cache() class ParkedTensorsNotRestored(RuntimeError): diff --git a/backends/cuda/tests/test_cuda_graph_timing.py b/backends/cuda/tests/test_cuda_graph_timing.py index 9dba3fa9e57..f0edc318307 100644 --- a/backends/cuda/tests/test_cuda_graph_timing.py +++ b/backends/cuda/tests/test_cuda_graph_timing.py @@ -9,7 +9,6 @@ python -m pytest backends/cuda/tests/test_cuda_graph_timing.py -v """ -import gc import statistics import struct import unittest @@ -154,31 +153,21 @@ def test_compile_spec_turns_it_off(self) -> None: self.assertIs(benchmarking._BENCHMARK_DISPATCH.get("cuda"), before) self.assertIs(benchmarking._BENCHMARK_DISPATCH.get("cuda"), before) - def test_cached_memory_check_matches_the_flattened_counters(self) -> None: + def test_captures_do_not_read_the_allocator_stats(self) -> None: x = torch.zeros(1 << 16, device="cuda") - with cgt.cuda_graph_timing() as stats: + # The torch.cuda and torch.accelerator stats APIs all end in these. + with mock.patch.object( + torch._C, "_cuda_memoryStats", wraps=torch._C._cuda_memoryStats + ) as cuda_stats, mock.patch.object( + torch._C, + "_accelerator_getDeviceStats", + wraps=torch._C._accelerator_getDeviceStats, + ) as accelerator_stats, cgt.cuda_graph_timing() as stats: for _ in range(8): benchmarking.benchmarker.benchmark(lambda: x.add_(1), device="cuda") self.assertEqual(stats.captured, 8) - # Free what earlier tests' mock reference cycles still hold, so no block - # is released between reading the counters and the zero-margin check. - gc.collect() - flattened = torch.cuda.memory_reserved() - torch.cuda.memory_allocated() - for threshold, empties in ((flattened - 1, 1), (flattened, 0)): - with self.subTest(threshold=threshold), mock.patch.object( - cgt, "_EMPTY_CACHE_THRESHOLD_BYTES", threshold - ), mock.patch.object( - torch.cuda, "empty_cache" - ) as empty_cache, mock.patch.object( - # The flattening APIs reach the nested stats through - # torch.cuda.memory's own global, so they never count here. - torch.cuda, - "memory_stats_as_nested_dict", - wraps=torch.cuda.memory_stats_as_nested_dict, - ) as nested: - cgt._release_cached_memory_if_large() - self.assertEqual(nested.call_count, 1) - self.assertEqual(empty_cache.call_count, empties) + self.assertEqual(cuda_stats.call_count, 0) + self.assertEqual(accelerator_stats.call_count, 0) # --------------------------------------------------------------------------- From 456fa4e1c07b0ae57643bc5dc7c23157bea01b89 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 10 Nov 2026 23:39:55 -0800 Subject: [PATCH 10/17] Respect CPU quotas when sizing quantized GEMM compiler workers Use the tightest visible cgroup CPU quota together with process affinity and host CPU count when sizing the ExecuTorch quantized GEMM compiler pool. Handle both cgroup versions and retain the existing 32-worker maximum on unrestricted hosts. CPU-only regression tests cover quota, affinity, namespace mount roots, parent limits and fallback behavior. Authored with OpenAI Codex. --- backends/cuda/BUCK | 12 ++ backends/cuda/_cpu_count.py | 88 +++++++++++ backends/cuda/tests/targets.bzl | 10 ++ backends/cuda/tests/test_cpu_count.py | 147 ++++++++++++++++++ .../triton/kernels/quantized_gemm_family.py | 3 +- 5 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 backends/cuda/_cpu_count.py create mode 100644 backends/cuda/tests/test_cpu_count.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 5ab15c33016..32e299aa467 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -3,6 +3,17 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target( + _kind = runtime.python_library, + name = "cpu_count", + srcs = [ + "_cpu_count.py", + ], + visibility = [ + "//executorch/backends/cuda/...", + ], +) + fbcode_target( _kind = runtime.python_library, name = "coalesced_int4_tensor", @@ -188,6 +199,7 @@ fbcode_target( ], deps = [ ":autotune", + ":cpu_count", "//caffe2:torch", ], ) diff --git a/backends/cuda/_cpu_count.py b/backends/cuda/_cpu_count.py new file mode 100644 index 00000000000..9112efd9461 --- /dev/null +++ b/backends/cuda/_cpu_count.py @@ -0,0 +1,88 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import os +import re +import sys +from pathlib import Path, PurePosixPath + + +def _cgroup_cpu_count() -> int | None: + if sys.platform != "linux": + return None + try: + memberships = Path("/proc/self/cgroup").read_text().splitlines() + mounts = Path("/proc/self/mountinfo").read_text().splitlines() + except OSError: + return None + + groups = {} + for membership in memberships: + fields = membership.split(":", 2) + if len(fields) == 3: + for controller in fields[1].split(","): + groups[controller] = PurePosixPath(fields[2]) + + limits = [] + for mount in mounts: + before, separator, after = mount.partition(" - ") + fields, filesystem = before.split(), after.split() + if not separator or len(fields) < 5 or len(filesystem) < 3: + continue + unified = filesystem[0] == "cgroup2" + if not unified and not ( + filesystem[0] == "cgroup" and "cpu" in filesystem[2].split(",") + ): + continue + group = groups.get("" if unified else "cpu") + if group is None: + continue + root, mountpoint = ( + Path(re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), field)) + for field in fields[3:5] + ) + try: + relative = ( + group.relative_to(root) + if group != PurePosixPath("/") + else PurePosixPath(".") + ) + except ValueError: + continue + if ".." in relative.parts: + continue + directory = mountpoint / relative + while True: + try: + if unified: + quota_text, period_text = ( + (directory / "cpu.max").read_text().split() + ) + quota, period = int(quota_text), int(period_text) + else: + quota = int((directory / "cpu.cfs_quota_us").read_text()) + period = int((directory / "cpu.cfs_period_us").read_text()) + if quota > 0 and period > 0: + limits.append((quota + period - 1) // period) + except (OSError, ValueError): + pass + if directory == mountpoint: + break + directory = directory.parent + return min(limits) if limits else None + + +def available_cpu_count() -> int: + """Respect affinity and CPU time quotas, including visible parent cgroups.""" + count = os.cpu_count() or 1 + try: + count = min(count, len(os.sched_getaffinity(0))) + except (AttributeError, OSError): + pass + quota_count = _cgroup_cpu_count() + if quota_count is not None: + count = min(count, quota_count) + return max(1, count) diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index dac423e5954..339f01ecda6 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -7,6 +7,16 @@ def define_common_targets(is_fbcode = False): if not is_fbcode: return + python_unittest( + name = "test_cpu_count", + srcs = [ + "test_cpu_count.py", + ], + deps = [ + "//executorch/backends/cuda:cpu_count", + ], + ) + python_unittest_remote_gpu( name = "test_cuda_export", srcs = [ diff --git a/backends/cuda/tests/test_cpu_count.py b/backends/cuda/tests/test_cpu_count.py new file mode 100644 index 00000000000..8d019d7be33 --- /dev/null +++ b/backends/cuda/tests/test_cpu_count.py @@ -0,0 +1,147 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from executorch.backends.cuda import _cpu_count + + +class AvailableCpuCountTest(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.mount = Path(self.temporary.name) + self.proc_cgroup = "0::/job\n" + self.proc_mountinfo = f"1 0 0:1 / {self.mount} rw - cgroup2 cgroup rw\n" + original_read = Path.read_text + + def read_text(path, *args, **kwargs): + if path == Path("/proc/self/cgroup"): + return self.proc_cgroup + if path == Path("/proc/self/mountinfo"): + return self.proc_mountinfo + return original_read(path, *args, **kwargs) + + for patch in ( + mock.patch.object(_cpu_count.sys, "platform", "linux"), + mock.patch.object(_cpu_count.os, "cpu_count", return_value=44), + mock.patch.object( + _cpu_count.os, + "sched_getaffinity", + return_value=set(range(44)), + create=True, + ), + mock.patch.object(Path, "read_text", read_text), + ): + patch.start() + self.addCleanup(patch.stop) + + def quota(self, group, quota, period=100000, unified=True): + directory = self.mount / group + directory.mkdir(parents=True, exist_ok=True) + if unified: + (directory / "cpu.max").write_text(f"{quota} {period}\n") + else: + (directory / "cpu.cfs_quota_us").write_text(str(quota)) + (directory / "cpu.cfs_period_us").write_text(str(period)) + + def test_unlimited_host_preserves_cpu_count(self): + self.quota("job", "max") + self.assertEqual(_cpu_count.available_cpu_count(), 44) + + def test_affinity_restricts_available_cpus(self): + self.quota("job", "max") + with mock.patch.object( + _cpu_count.os, "sched_getaffinity", return_value=set(range(7)) + ): + self.assertEqual(_cpu_count.available_cpu_count(), 7) + + def test_cgroup_v2_quota_without_affinity_limit(self): + self.quota("job", 1100000) + self.assertEqual(_cpu_count.available_cpu_count(), 11) + + def test_tightest_visible_parent_quota(self): + self.proc_cgroup = "0::/parent/job\n" + self.quota("parent/job", "max") + self.quota("parent", 1100000) + self.quota("", 800000) + self.assertEqual(_cpu_count.available_cpu_count(), 8) + + def test_mount_root_maps_host_cgroup_path(self): + self.proc_cgroup = "0::/containers/container/job\n" + self.proc_mountinfo = ( + f"1 0 0:1 /containers/container {self.mount} rw - cgroup2 cgroup rw\n" + ) + self.quota("job", 1100000) + self.assertEqual(_cpu_count.available_cpu_count(), 11) + + def test_cgroup_namespace_root_maps_mountpoint(self): + self.proc_cgroup = "0::/\n" + self.proc_mountinfo = ( + f"1 0 0:1 /containers/container {self.mount} rw - cgroup2 cgroup rw\n" + ) + self.quota("", 1100000) + self.assertEqual(_cpu_count.available_cpu_count(), 11) + + def test_cgroup_v1_combined_controllers(self): + self.proc_cgroup = "4:cpu,cpuacct:/job\n2:memory:/other\n" + self.proc_mountinfo = ( + f"1 0 0:1 / {self.mount} rw - cgroup cgroup rw,cpu,cpuacct\n" + ) + self.quota("job", 1100000, unified=False) + self.assertEqual(_cpu_count.available_cpu_count(), 11) + + def test_cgroup_v1_unlimited_parent_can_restrict(self): + self.proc_cgroup = "4:cpu:/parent/job\n" + self.proc_mountinfo = f"1 0 0:1 / {self.mount} rw - cgroup cgroup rw,cpu\n" + self.quota("parent/job", -1, unified=False) + self.quota("parent", 1100000, unified=False) + self.assertEqual(_cpu_count.available_cpu_count(), 11) + + def test_fractional_quota_rounds_up_to_a_worker(self): + for quota, expected in ((50000, 1), (150000, 2)): + with self.subTest(quota=quota): + self.quota("job", quota) + self.assertEqual(_cpu_count.available_cpu_count(), expected) + + def test_affinity_tighter_than_quota(self): + self.quota("job", 1100000) + with mock.patch.object(_cpu_count.os, "sched_getaffinity", return_value={0, 1}): + self.assertEqual(_cpu_count.available_cpu_count(), 2) + + def test_missing_proc_files_falls_back(self): + with mock.patch.object(Path, "read_text", side_effect=PermissionError): + self.assertEqual(_cpu_count.available_cpu_count(), 44) + + def test_invalid_or_missing_quota_does_not_hide_parent(self): + self.proc_cgroup = "0::/parent/job\n" + self.quota("parent", 1100000) + for value in (None, "invalid", "0 100000", "1100000 0"): + with self.subTest(value=value): + if value is not None: + self.quota("parent/job", 1) + (self.mount / "parent/job/cpu.max").write_text(value) + self.assertEqual(_cpu_count.available_cpu_count(), 11) + + def test_non_linux_without_affinity(self): + with ( + mock.patch.object(_cpu_count.sys, "platform", "win32"), + mock.patch.object( + _cpu_count.os, "sched_getaffinity", side_effect=AttributeError + ), + ): + self.assertEqual(_cpu_count.available_cpu_count(), 44) + + def test_missing_cpu_count_still_returns_one(self): + with mock.patch.object(_cpu_count.os, "cpu_count", return_value=None): + self.assertEqual(_cpu_count.available_cpu_count(), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/triton/kernels/quantized_gemm_family.py b/backends/cuda/triton/kernels/quantized_gemm_family.py index aa58b17699c..cb2e46b551c 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -41,6 +41,7 @@ import torch import triton +from executorch.backends.cuda._cpu_count import available_cpu_count from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import splitk_reduce from torch.library import triton_op, wrap_triton @@ -180,7 +181,7 @@ def _has_untuned_configs(kernel, args: Sequence, kwargs: dict) -> bool: def _compile_workers() -> int: - return min(_MAX_COMPILE_THREADS, os.cpu_count() or 1) + return min(_MAX_COMPILE_THREADS, available_cpu_count()) _PROCESS_POOL_LOCK = threading.Lock() From 687cd9a78e6ebb06cffb8a09b1bfe23919e989e5 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 11 Nov 2026 00:22:24 -0800 Subject: [PATCH 11/17] Move CPU count helper to backends/cuda/utils and reduce its complexity Rename backends/cuda/_cpu_count.py to backends/cuda/utils/cpu_count.py and split cgroup parsing into small helpers so flake8 C901 passes. Behavior is unchanged. --- backends/cuda/BUCK | 3 +- backends/cuda/_cpu_count.py | 88 -------------- backends/cuda/tests/test_cpu_count.py | 46 ++++---- .../triton/kernels/quantized_gemm_family.py | 2 +- backends/cuda/utils/__init__.py | 5 + backends/cuda/utils/cpu_count.py | 108 ++++++++++++++++++ 6 files changed, 139 insertions(+), 113 deletions(-) delete mode 100644 backends/cuda/_cpu_count.py create mode 100644 backends/cuda/utils/__init__.py create mode 100644 backends/cuda/utils/cpu_count.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 32e299aa467..68e3777a06f 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -7,7 +7,8 @@ fbcode_target( _kind = runtime.python_library, name = "cpu_count", srcs = [ - "_cpu_count.py", + "utils/__init__.py", + "utils/cpu_count.py", ], visibility = [ "//executorch/backends/cuda/...", diff --git a/backends/cuda/_cpu_count.py b/backends/cuda/_cpu_count.py deleted file mode 100644 index 9112efd9461..00000000000 --- a/backends/cuda/_cpu_count.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import os -import re -import sys -from pathlib import Path, PurePosixPath - - -def _cgroup_cpu_count() -> int | None: - if sys.platform != "linux": - return None - try: - memberships = Path("/proc/self/cgroup").read_text().splitlines() - mounts = Path("/proc/self/mountinfo").read_text().splitlines() - except OSError: - return None - - groups = {} - for membership in memberships: - fields = membership.split(":", 2) - if len(fields) == 3: - for controller in fields[1].split(","): - groups[controller] = PurePosixPath(fields[2]) - - limits = [] - for mount in mounts: - before, separator, after = mount.partition(" - ") - fields, filesystem = before.split(), after.split() - if not separator or len(fields) < 5 or len(filesystem) < 3: - continue - unified = filesystem[0] == "cgroup2" - if not unified and not ( - filesystem[0] == "cgroup" and "cpu" in filesystem[2].split(",") - ): - continue - group = groups.get("" if unified else "cpu") - if group is None: - continue - root, mountpoint = ( - Path(re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), field)) - for field in fields[3:5] - ) - try: - relative = ( - group.relative_to(root) - if group != PurePosixPath("/") - else PurePosixPath(".") - ) - except ValueError: - continue - if ".." in relative.parts: - continue - directory = mountpoint / relative - while True: - try: - if unified: - quota_text, period_text = ( - (directory / "cpu.max").read_text().split() - ) - quota, period = int(quota_text), int(period_text) - else: - quota = int((directory / "cpu.cfs_quota_us").read_text()) - period = int((directory / "cpu.cfs_period_us").read_text()) - if quota > 0 and period > 0: - limits.append((quota + period - 1) // period) - except (OSError, ValueError): - pass - if directory == mountpoint: - break - directory = directory.parent - return min(limits) if limits else None - - -def available_cpu_count() -> int: - """Respect affinity and CPU time quotas, including visible parent cgroups.""" - count = os.cpu_count() or 1 - try: - count = min(count, len(os.sched_getaffinity(0))) - except (AttributeError, OSError): - pass - quota_count = _cgroup_cpu_count() - if quota_count is not None: - count = min(count, quota_count) - return max(1, count) diff --git a/backends/cuda/tests/test_cpu_count.py b/backends/cuda/tests/test_cpu_count.py index 8d019d7be33..50f040abf43 100644 --- a/backends/cuda/tests/test_cpu_count.py +++ b/backends/cuda/tests/test_cpu_count.py @@ -9,7 +9,7 @@ from pathlib import Path from unittest import mock -from executorch.backends.cuda import _cpu_count +from executorch.backends.cuda.utils import cpu_count class AvailableCpuCountTest(unittest.TestCase): @@ -29,10 +29,10 @@ def read_text(path, *args, **kwargs): return original_read(path, *args, **kwargs) for patch in ( - mock.patch.object(_cpu_count.sys, "platform", "linux"), - mock.patch.object(_cpu_count.os, "cpu_count", return_value=44), + mock.patch.object(cpu_count.sys, "platform", "linux"), + mock.patch.object(cpu_count.os, "cpu_count", return_value=44), mock.patch.object( - _cpu_count.os, + cpu_count.os, "sched_getaffinity", return_value=set(range(44)), create=True, @@ -53,25 +53,25 @@ def quota(self, group, quota, period=100000, unified=True): def test_unlimited_host_preserves_cpu_count(self): self.quota("job", "max") - self.assertEqual(_cpu_count.available_cpu_count(), 44) + self.assertEqual(cpu_count.available_cpu_count(), 44) def test_affinity_restricts_available_cpus(self): self.quota("job", "max") with mock.patch.object( - _cpu_count.os, "sched_getaffinity", return_value=set(range(7)) + cpu_count.os, "sched_getaffinity", return_value=set(range(7)) ): - self.assertEqual(_cpu_count.available_cpu_count(), 7) + self.assertEqual(cpu_count.available_cpu_count(), 7) def test_cgroup_v2_quota_without_affinity_limit(self): self.quota("job", 1100000) - self.assertEqual(_cpu_count.available_cpu_count(), 11) + self.assertEqual(cpu_count.available_cpu_count(), 11) def test_tightest_visible_parent_quota(self): self.proc_cgroup = "0::/parent/job\n" self.quota("parent/job", "max") self.quota("parent", 1100000) self.quota("", 800000) - self.assertEqual(_cpu_count.available_cpu_count(), 8) + self.assertEqual(cpu_count.available_cpu_count(), 8) def test_mount_root_maps_host_cgroup_path(self): self.proc_cgroup = "0::/containers/container/job\n" @@ -79,7 +79,7 @@ def test_mount_root_maps_host_cgroup_path(self): f"1 0 0:1 /containers/container {self.mount} rw - cgroup2 cgroup rw\n" ) self.quota("job", 1100000) - self.assertEqual(_cpu_count.available_cpu_count(), 11) + self.assertEqual(cpu_count.available_cpu_count(), 11) def test_cgroup_namespace_root_maps_mountpoint(self): self.proc_cgroup = "0::/\n" @@ -87,7 +87,7 @@ def test_cgroup_namespace_root_maps_mountpoint(self): f"1 0 0:1 /containers/container {self.mount} rw - cgroup2 cgroup rw\n" ) self.quota("", 1100000) - self.assertEqual(_cpu_count.available_cpu_count(), 11) + self.assertEqual(cpu_count.available_cpu_count(), 11) def test_cgroup_v1_combined_controllers(self): self.proc_cgroup = "4:cpu,cpuacct:/job\n2:memory:/other\n" @@ -95,29 +95,29 @@ def test_cgroup_v1_combined_controllers(self): f"1 0 0:1 / {self.mount} rw - cgroup cgroup rw,cpu,cpuacct\n" ) self.quota("job", 1100000, unified=False) - self.assertEqual(_cpu_count.available_cpu_count(), 11) + self.assertEqual(cpu_count.available_cpu_count(), 11) def test_cgroup_v1_unlimited_parent_can_restrict(self): self.proc_cgroup = "4:cpu:/parent/job\n" self.proc_mountinfo = f"1 0 0:1 / {self.mount} rw - cgroup cgroup rw,cpu\n" self.quota("parent/job", -1, unified=False) self.quota("parent", 1100000, unified=False) - self.assertEqual(_cpu_count.available_cpu_count(), 11) + self.assertEqual(cpu_count.available_cpu_count(), 11) def test_fractional_quota_rounds_up_to_a_worker(self): for quota, expected in ((50000, 1), (150000, 2)): with self.subTest(quota=quota): self.quota("job", quota) - self.assertEqual(_cpu_count.available_cpu_count(), expected) + self.assertEqual(cpu_count.available_cpu_count(), expected) def test_affinity_tighter_than_quota(self): self.quota("job", 1100000) - with mock.patch.object(_cpu_count.os, "sched_getaffinity", return_value={0, 1}): - self.assertEqual(_cpu_count.available_cpu_count(), 2) + with mock.patch.object(cpu_count.os, "sched_getaffinity", return_value={0, 1}): + self.assertEqual(cpu_count.available_cpu_count(), 2) def test_missing_proc_files_falls_back(self): with mock.patch.object(Path, "read_text", side_effect=PermissionError): - self.assertEqual(_cpu_count.available_cpu_count(), 44) + self.assertEqual(cpu_count.available_cpu_count(), 44) def test_invalid_or_missing_quota_does_not_hide_parent(self): self.proc_cgroup = "0::/parent/job\n" @@ -127,20 +127,20 @@ def test_invalid_or_missing_quota_does_not_hide_parent(self): if value is not None: self.quota("parent/job", 1) (self.mount / "parent/job/cpu.max").write_text(value) - self.assertEqual(_cpu_count.available_cpu_count(), 11) + self.assertEqual(cpu_count.available_cpu_count(), 11) def test_non_linux_without_affinity(self): with ( - mock.patch.object(_cpu_count.sys, "platform", "win32"), + mock.patch.object(cpu_count.sys, "platform", "win32"), mock.patch.object( - _cpu_count.os, "sched_getaffinity", side_effect=AttributeError + cpu_count.os, "sched_getaffinity", side_effect=AttributeError ), ): - self.assertEqual(_cpu_count.available_cpu_count(), 44) + self.assertEqual(cpu_count.available_cpu_count(), 44) def test_missing_cpu_count_still_returns_one(self): - with mock.patch.object(_cpu_count.os, "cpu_count", return_value=None): - self.assertEqual(_cpu_count.available_cpu_count(), 1) + with mock.patch.object(cpu_count.os, "cpu_count", return_value=None): + self.assertEqual(cpu_count.available_cpu_count(), 1) if __name__ == "__main__": diff --git a/backends/cuda/triton/kernels/quantized_gemm_family.py b/backends/cuda/triton/kernels/quantized_gemm_family.py index cb2e46b551c..3a30a031404 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -41,8 +41,8 @@ import torch import triton -from executorch.backends.cuda._cpu_count import available_cpu_count from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import splitk_reduce +from executorch.backends.cuda.utils.cpu_count import available_cpu_count from torch.library import triton_op, wrap_triton try: diff --git a/backends/cuda/utils/__init__.py b/backends/cuda/utils/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/cuda/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/backends/cuda/utils/cpu_count.py b/backends/cuda/utils/cpu_count.py new file mode 100644 index 00000000000..b03f6599f0e --- /dev/null +++ b/backends/cuda/utils/cpu_count.py @@ -0,0 +1,108 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import os +import re +import sys +from pathlib import Path, PurePosixPath + + +def _cgroup_memberships(memberships: list[str]) -> dict[str, PurePosixPath]: + groups = {} + for membership in memberships: + fields = membership.split(":", 2) + if len(fields) == 3: + for controller in fields[1].split(","): + groups[controller] = PurePosixPath(fields[2]) + return groups + + +def _unescape_mount_field(field: str) -> Path: + return Path(re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), field)) + + +def _cpu_mount( + mount: str, groups: dict[str, PurePosixPath] +) -> tuple[Path, Path, bool] | None: + """Return (cgroup directory, mountpoint, unified) for a CPU cgroup mount.""" + before, separator, after = mount.partition(" - ") + fields, filesystem = before.split(), after.split() + if not separator or len(fields) < 5 or len(filesystem) < 3: + return None + unified = filesystem[0] == "cgroup2" + if not unified and not ( + filesystem[0] == "cgroup" and "cpu" in filesystem[2].split(",") + ): + return None + group = groups.get("" if unified else "cpu") + if group is None: + return None + root, mountpoint = (_unescape_mount_field(field) for field in fields[3:5]) + try: + relative = ( + group.relative_to(root) + if group != PurePosixPath("/") + else PurePosixPath(".") + ) + except ValueError: + return None + if ".." in relative.parts: + return None + return mountpoint / relative, mountpoint, unified + + +def _quota_cpu_count(directory: Path, unified: bool) -> int | None: + try: + if unified: + quota_text, period_text = (directory / "cpu.max").read_text().split() + quota, period = int(quota_text), int(period_text) + else: + quota = int((directory / "cpu.cfs_quota_us").read_text()) + period = int((directory / "cpu.cfs_period_us").read_text()) + except (OSError, ValueError): + return None + if quota > 0 and period > 0: + return (quota + period - 1) // period + return None + + +def _cgroup_cpu_count() -> int | None: + if sys.platform != "linux": + return None + try: + memberships = Path("/proc/self/cgroup").read_text().splitlines() + mounts = Path("/proc/self/mountinfo").read_text().splitlines() + except OSError: + return None + + groups = _cgroup_memberships(memberships) + limits = [] + for mount in mounts: + cpu_mount = _cpu_mount(mount, groups) + if cpu_mount is None: + continue + directory, mountpoint, unified = cpu_mount + while True: + limit = _quota_cpu_count(directory, unified) + if limit is not None: + limits.append(limit) + if directory == mountpoint: + break + directory = directory.parent + return min(limits) if limits else None + + +def available_cpu_count() -> int: + """Respect affinity and CPU time quotas, including visible parent cgroups.""" + count = os.cpu_count() or 1 + try: + count = min(count, len(os.sched_getaffinity(0))) + except (AttributeError, OSError): + pass + quota_count = _cgroup_cpu_count() + if quota_count is not None: + count = min(count, quota_count) + return max(1, count) From 245be7bc88b6deda3b9ef34967131b57986aaa1a Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 11 Nov 2026 02:07:14 -0800 Subject: [PATCH 12/17] Saturate only usable CPUs in the CUDA autotune load tests SaturatedCpu started one spinner per host CPU. Under a cgroup quota that oversubscribes the quota now that the quantized GEMM compile pool is sized to the usable CPUs, which starved the pool and pushed the int4 under-load test past its job timeout. Spawn one spinner per usable CPU instead. --- backends/cuda/tests/autotune_test_utils.py | 5 +++-- backends/cuda/tests/targets.bzl | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/backends/cuda/tests/autotune_test_utils.py b/backends/cuda/tests/autotune_test_utils.py index 3e49bff3f94..43b1e677bd2 100644 --- a/backends/cuda/tests/autotune_test_utils.py +++ b/backends/cuda/tests/autotune_test_utils.py @@ -17,6 +17,7 @@ import torch import triton from executorch.backends.cuda.autotune import cuda_graph_timing as cgt +from executorch.backends.cuda.utils.cpu_count import available_cpu_count from torch._inductor.runtime.triton_heuristics import CachingAutotuner from torch._inductor.select_algorithm import AlgorithmSelectorCache @@ -58,13 +59,13 @@ def _spin(stop_at: float) -> None: class SaturatedCpu: - """One busy-loop process per core; can be paused to measure a reference.""" + """One busy-loop process per usable CPU; can be paused to measure a reference.""" def __init__(self) -> None: stop_at = time.time() + 3600 self.workers = [ multiprocessing.Process(target=_spin, args=(stop_at,), daemon=True) - for _ in range(os.cpu_count() or 1) + for _ in range(available_cpu_count()) ] def __enter__(self) -> "SaturatedCpu": diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index 339f01ecda6..09b3700cbf5 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -52,6 +52,7 @@ def define_common_targets(is_fbcode = False): deps = [ "//caffe2:torch", "//executorch/backends/cuda:autotune", + "//executorch/backends/cuda:cpu_count", ], ) From 838de723343a390ee71b8a4eeacb3a8b326a8c2b Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 11 Nov 2026 04:24:21 -0800 Subject: [PATCH 13/17] Give the kquant-17gb Muse Glimmer CUDA e2e job 480 minutes Every recent kquant-17gb dflash-image-batching run, on main and on pull requests, was cancelled at the 240 minute limit while export_dflash was still running (more than 3.5 hours into the export). Raise that matrix entry to 480 minutes and keep the two kquant-dynamic entries at 240. --- .github/workflows/cuda.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 56bf2f71262..d5736575fbe 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1221,14 +1221,18 @@ jobs: fail-fast: false matrix: include: + # The kquant-17gb dflash export alone runs past 3.5 hours on this runner. - variant: kquant-17gb mode: dflash-image-batching + timeout: 480 - variant: kquant-dynamic mode: solo-text + timeout: 240 - variant: kquant-dynamic mode: dflash-image + timeout: 240 with: - timeout: 240 + timeout: ${{ matrix.timeout }} runner: mt-l-x86iavx512-11-125-a100 gpu-arch-type: cuda gpu-arch-version: "13.0" From 23279df77d91b15b7b7a160ac96d453a9f093317 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 11 Nov 2026 10:06:28 -0800 Subject: [PATCH 14/17] Run the full Muse Glimmer CUDA e2e only on pushes to main test-muse-glimmer-cuda-e2e runs three A100 jobs, and the kquant-17gb dflash-image-batching export alone takes about five hours, so running it on every pull request push costs many A100 hours per iteration. Gate the job on pushes to main; pull requests, ciflow tags and manual dispatch no longer start it. Pull requests keep the unittest-cuda-* jobs, including unittest-cuda-muse-glimmer-batching. A CI script test pins the gate. --- .../tests/test_muse_glimmer_batching_ci.py | 20 +++++++++++++++++++ .github/workflows/cuda.yml | 6 +++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/.ci/scripts/tests/test_muse_glimmer_batching_ci.py b/.ci/scripts/tests/test_muse_glimmer_batching_ci.py index 5c8fcead860..b04eb9eab23 100644 --- a/.ci/scripts/tests/test_muse_glimmer_batching_ci.py +++ b/.ci/scripts/tests/test_muse_glimmer_batching_ci.py @@ -129,6 +129,26 @@ def test_smoke_cmake_is_opt_in_and_reuses_runtime(self): self.assertNotIn("BENCHMARK", smoke) self.assertNotIn("batching_bench", smoke) + def test_full_e2e_runs_only_on_main_pushes(self): + workflow = (ROOT / ".github/workflows/cuda.yml").read_text() + e2e_if = ( + workflow.split(" test-muse-glimmer-cuda-e2e:", 1)[1] + .split(" if: |\n", 1)[1] + .split(" uses:", 1)[0] + .strip() + ) + self.assertTrue( + e2e_if.startswith( + "github.event_name == 'push' && github.ref == 'refs/heads/main' && (" + ) + ) + self.assertTrue(e2e_if.endswith(")")) + self.assertEqual(e2e_if.count("("), e2e_if.count(")")) + unit_if = workflow.split(" unittest-cuda-muse-glimmer-batching:", 1)[1].split( + " uses:", 1 + )[0] + self.assertNotIn("github.event_name", unit_if) + def test_owned_shell_and_unit_workflow_syntax(self): for script in ("test_model_e2e.sh", "build_muse_glimmer_batching_smoke.sh"): result = subprocess.run(["bash", "-n", str(SCRIPTS / script)]) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index d5736575fbe..83e60609cee 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1197,7 +1197,11 @@ jobs: test-muse-glimmer-cuda-e2e: name: test-muse-glimmer-cuda-e2e-${{ matrix.variant }}-${{ matrix.mode }} needs: [changed-files, run-decision] + # Hours of A100 time per run, so only after merge to main (not on pull + # requests, ciflow tags or manual dispatch); pull requests get the + # unittest-cuda-* jobs above. if: | + github.event_name == 'push' && github.ref == 'refs/heads/main' && ( contains(needs.changed-files.outputs.changed-files, '.ci/scripts/build_muse_glimmer_batching_smoke.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_smoke.cpp') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_workload.h') || @@ -1212,7 +1216,7 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') || contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer') || - needs.run-decision.outputs.is-full-run == 'true' + needs.run-decision.outputs.is-full-run == 'true') uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write From b9f2213380858a97d4b06066785eb3a2018f71db Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 12 Nov 2026 22:52:26 -0800 Subject: [PATCH 15/17] Export four static forward sizes in the kquant-17gb Muse Glimmer batching e2e Each forward_{N} is its own AOTI compile and autotune. On the 11-vCPU A100 runner the full 1..64 set does not fit 240 minutes (#25018 still timed out with this alone, so the 480-minute timeout stays); 1,4,16,64 exercises the same CUDA executor padding path with three fewer compiles. --- .ci/scripts/export_model_artifact.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 9a214964ce1..13971dfdf68 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -716,7 +716,9 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then # executor drives, plus the vision encoder and forward_from_embeddings. OFFGRAPH_FLAG="" if [ "$MODE" = "dflash-image-batching" ]; then - OFFGRAPH_FLAG="--use-offgraph-kv-cache" + # Each static forward size is a separate AOTI compile; the default + # seven sizes add about an hour on the 11-vCPU CI runner. + OFFGRAPH_FLAG="--use-offgraph-kv-cache --forward-sizes 1,4,16,64" fi EXPORT_START_SECONDS=$SECONDS TMPDIR="$INDUCTOR_TMPDIR" \ From b8c8528523d667031188edb0aa0bb3571b9c394c Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 13 Nov 2026 00:02:02 -0800 Subject: [PATCH 16/17] Fit the kquant-17gb Muse Glimmer batching e2e in 240 minutes Export forward_{1,4} and a single draft owner in CI, and drop the job's 480-minute timeout back to the shared 240. Even with the CPU-quota-aware compile pool and forward_{1,4,16,64}, the export was still compiling the draft methods at the 240-minute mark (#25018); forward_16/64 and draft_forward_{2,3,4} account for about two of its hours. The run still covers static and dynamic target steps, vision, and DFlash drafting on the batching runtime. --- .ci/scripts/export_model_artifact.sh | 8 +++++--- .github/workflows/cuda.yml | 6 +----- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 13971dfdf68..9894a305254 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -716,9 +716,11 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then # executor drives, plus the vision encoder and forward_from_embeddings. OFFGRAPH_FLAG="" if [ "$MODE" = "dflash-image-batching" ]; then - # Each static forward size is a separate AOTI compile; the default - # seven sizes add about an hour on the 11-vCPU CI runner. - OFFGRAPH_FLAG="--use-offgraph-kv-cache --forward-sizes 1,4,16,64" + # Every static forward_{N} and draft_forward_{N} is its own AOTI + # compile and autotune, which the 11-vCPU runner cannot fit in the + # job timeout at the default sizes; these still cover static and + # dynamic target steps and DFlash drafting. + OFFGRAPH_FLAG="--use-offgraph-kv-cache --forward-sizes 1,4 --max-draft-owners 1" fi EXPORT_START_SECONDS=$SECONDS TMPDIR="$INDUCTOR_TMPDIR" \ diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 83e60609cee..4fd15c5ac31 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1225,18 +1225,14 @@ jobs: fail-fast: false matrix: include: - # The kquant-17gb dflash export alone runs past 3.5 hours on this runner. - variant: kquant-17gb mode: dflash-image-batching - timeout: 480 - variant: kquant-dynamic mode: solo-text - timeout: 240 - variant: kquant-dynamic mode: dflash-image - timeout: 240 with: - timeout: ${{ matrix.timeout }} + timeout: 240 runner: mt-l-x86iavx512-11-125-a100 gpu-arch-type: cuda gpu-arch-version: "13.0" From b8f1ad5c241654e0297a84d9e8ae227f13a2d378 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Sat, 14 Nov 2026 12:04:41 -0800 Subject: [PATCH 17/17] [DO NOT LAND] Run the Muse Glimmer CUDA e2e on this pull request to observe the batching job --- .github/workflows/cuda.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 4fd15c5ac31..cfc9006f557 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1201,7 +1201,7 @@ jobs: # requests, ciflow tags or manual dispatch); pull requests get the # unittest-cuda-* jobs above. if: | - github.event_name == 'push' && github.ref == 'refs/heads/main' && ( + ( contains(needs.changed-files.outputs.changed-files, '.ci/scripts/build_muse_glimmer_batching_smoke.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_smoke.cpp') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_workload.h') ||