From 6b88a806c2b91336222ac82529930f36c36c2aff Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 12 Aug 2026 16:29:39 -0700 Subject: [PATCH] Avoid rereading CUDA weights blob for SHA-256 --- backends/cuda/cuda_backend.py | 29 ++++++++++++++------ backends/cuda/tests/test_cuda_partitioner.py | 13 +++++++++ exir/_serialize/_cord.py | 13 ++++++--- 3 files changed, 43 insertions(+), 12 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 9232dd57324..5b6ae5427d2 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -10,6 +10,7 @@ import ctypes import functools import gc +import hashlib import logging import os import shutil @@ -306,12 +307,17 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool: return value == "ON" -def _write_aoti_weights_blob(weights, blob_path: str) -> None: - """Stream AOTI tensor storages without creating a model-sized bytes object.""" +def _write_aoti_weights_blob(weights, blob_path: str) -> bytes: + """Stream AOTI tensor storages and return their SHA-256 digest.""" _trim_host_memory() tensors = [tensor for tensor, _ in weights.values()] all_cuda = all(tensor.is_cuda for tensor in tensors) chunk_size = 8 * 1024 * 1024 + digest = hashlib.sha256() + + def write_chunk(output, chunk) -> None: + digest.update(chunk) + output.write(chunk) with open(blob_path, "wb") as output: for tensor in tensors: @@ -325,18 +331,19 @@ def _write_aoti_weights_blob(weights, blob_path: str) -> None: ).set_(storage, 0, (nbytes,), (1,)) for offset in range(0, nbytes, chunk_size): cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu() - output.write(memoryview(cpu_chunk.numpy())) + write_chunk(output, memoryview(cpu_chunk.numpy())) del byte_tensor, cpu_chunk elif nbytes: raw_array = (ctypes.c_ubyte * nbytes).from_address(storage.data_ptr()) raw_view = memoryview(raw_array).cast("B") for offset in range(0, nbytes, chunk_size): - output.write(raw_view[offset : offset + chunk_size]) + write_chunk(output, raw_view[offset : offset + chunk_size]) del raw_view, raw_array if not all_cuda and (padding := (-nbytes) % 64): - output.write(bytes(padding)) + write_chunk(output, bytes(padding)) del storage _trim_host_memory() + return digest.digest() @final @@ -350,6 +357,8 @@ class CudaBackend(AotiBackend, BackendDetails): using the Executorch runtime. """ + _materialized_blob_hashes: Dict[str, bytes] = {} + @classmethod def get_device_name(cls) -> str: return "cuda" @@ -465,8 +474,10 @@ def load_weights_blob( """ if not cls._is_low_memory_mode(compile_specs): return super().load_weights_blob(blob_path, compile_specs) - blob_data = FileBackedData.move_from(blob_path) - return blob_data, blob_data.sha256().hex() + known_hash = cls._materialized_blob_hashes.pop(blob_path, None) + blob_data = FileBackedData.move_from(blob_path, sha256=known_hash) + weights_blob_hash = known_hash or blob_data.sha256() + return blob_data, weights_blob_hash.hex() @classmethod def materialize_weights_blob( @@ -491,7 +502,9 @@ def materialize_weights_blob( if isinstance(path, str) and path.endswith(".wrapper.so") ) blob_path = os.path.splitext(so_path)[0] + "_weights.blob" - _write_aoti_weights_blob(weights[0], blob_path) + cls._materialized_blob_hashes[blob_path] = _write_aoti_weights_blob( + weights[0], blob_path + ) # Forcing the external-weights ABI makes Inductor emit an empty blob # path alongside the Weights object. Replace that file in place and do diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 38a65da1c6f..b5be7a0df5c 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -79,6 +79,19 @@ def test_low_memory_weights_are_streamed_in_binary_blob_format(self) -> None: ) self.assertEqual(expected, data) + # The streaming write computes the digest in the same pass. Loading + # the file-backed blob must not reread it solely for hashing. + with patch.object( + FileBackedData, + "sha256", + side_effect=AssertionError("unexpected blob reread"), + ): + blob, digest = CudaBackend.load_weights_blob( + blob_path, [CompileSpec("low_memory_mode", b"ON")] + ) + self.assertEqual(hashlib.sha256(expected).hexdigest(), digest) + self.assertEqual(expected, blob.to_bytes()) + def test_low_memory_blob_stays_file_backed(self) -> None: data = b"cuda weights" with tempfile.TemporaryDirectory() as directory: diff --git a/exir/_serialize/_cord.py b/exir/_serialize/_cord.py index 4177f9088cc..c4bd4415ae9 100644 --- a/exir/_serialize/_cord.py +++ b/exir/_serialize/_cord.py @@ -18,10 +18,15 @@ class FileBackedData: _COPY_CHUNK_SIZE = 8 * 1024 * 1024 - def __init__(self, path: str, cleanup: bool = False) -> None: + def __init__( + self, + path: str, + cleanup: bool = False, + sha256: Optional[bytes] = None, + ) -> None: self._path = path self._size = os.path.getsize(path) - self._sha256: Optional[bytes] = None + self._sha256 = sha256 self._finalizer = ( weakref.finalize(self, self._remove, path) if cleanup else None ) @@ -34,7 +39,7 @@ def _remove(path: str) -> None: pass @classmethod - def move_from(cls, path: str) -> "FileBackedData": + def move_from(cls, path: str, sha256: Optional[bytes] = None) -> "FileBackedData": """Take ownership of ``path`` without loading its contents.""" directory = os.path.dirname(path) or "." fd, owned_path = tempfile.mkstemp( @@ -46,7 +51,7 @@ def move_from(cls, path: str) -> "FileBackedData": except Exception: os.remove(owned_path) raise - return cls(owned_path, cleanup=True) + return cls(owned_path, cleanup=True, sha256=sha256) def __len__(self) -> int: return self._size