From 1d6744e6845100aba8609952aeb8f35dcd7a0d66 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 13 Aug 2026 11:30:04 -0700 Subject: [PATCH 01/12] cuda: share AOTI weights by FQN across methods --- backends/cuda/CMakeLists.txt | 6 + backends/cuda/cuda_backend.py | 401 ++++++++++++--- backends/cuda/runtime/cuda_backend.cpp | 474 +++++++++++++++++- backends/cuda/runtime/cuda_delegate_handle.h | 49 ++ backends/cuda/runtime/cuda_weight_manifest.h | 217 ++++++++ backends/cuda/runtime/targets.bzl | 20 + .../test/test_cuda_weight_manifest.cpp | 110 ++++ backends/cuda/tests/test_cuda_partitioner.py | 130 +++-- 8 files changed, 1294 insertions(+), 113 deletions(-) create mode 100644 backends/cuda/runtime/cuda_weight_manifest.h create mode 100644 backends/cuda/runtime/test/test_cuda_weight_manifest.cpp diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 21f8f5e2914..0cae1bc77f3 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -448,4 +448,10 @@ if(BUILD_TESTING) EXTRA_LIBS aoti_cuda_backend ) target_compile_definitions(test_cuda_mutable_state PRIVATE CUDA_AVAILABLE=1) + + et_cxx_test( + test_cuda_weight_manifest SOURCES + runtime/test/test_cuda_weight_manifest.cpp EXTRA_LIBS aoti_cuda_backend + ) + target_compile_definitions(test_cuda_weight_manifest PRIVATE CUDA_AVAILABLE=1) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 1d6e0b11370..5b2a0ef737c 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -14,10 +14,13 @@ import logging import os import shutil +import struct +import tempfile import threading import typing +from dataclasses import dataclass from importlib import resources -from typing import Any, Dict, final, List, Optional +from typing import Any, Dict, final, List, Optional, Tuple import torch from executorch.backends.aoti.aoti_backend import AotiBackend @@ -31,9 +34,11 @@ ReplaceEdgeOpWithTritonOpPass, ) from executorch.exir._serialize._cord import FileBackedData +from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir._warnings import experimental -from executorch.exir.backend.backend_details import BackendDetails +from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.tensor import scalar_type_enum from torch._inductor.decomposition import conv1d_to_conv2d from torch.nn.attention import SDPBackend @@ -61,6 +66,34 @@ _CPU_CLONE_GUARD = threading.local() +_FQN_WEIGHTS_MAGIC = b"ETCUDAFQN1" +_FQN_WEIGHTS_CAPTURE = threading.local() + + +@dataclass +class _FqnWeightEntry: + fqn: str + storage_key: str + storage_group: int + storage_nbytes: int + dtype: int + storage_offset: int + sizes: Tuple[int, ...] + strides: Tuple[int, ...] + shareable: bool + + +@dataclass +class _FqnWeightArtifact: + entries: List[_FqnWeightEntry] + storages: Dict[str, FileBackedData] + + +@dataclass +class _FqnWeightCapture: + mutated_fqns: set[str] + artifact: Optional[_FqnWeightArtifact] = None + def _is_cpu_clone_active() -> bool: return getattr(_CPU_CLONE_GUARD, "active", False) @@ -195,10 +228,11 @@ def _compile_time_cpu_clones(target_device: torch.device): # noqa: C901 orig_tensor_properties = _codecache.TensorProperties orig_determine_aoti_mmap_flags = _codecache.determine_aoti_mmap_flags - def _force_external_weights_for_streaming(consts_size): - # ``pickle_weights`` normally tells AOTI that no external binary blob - # exists. We materialize that pickle output as a streamed blob below, - # so the generated wrapper must use the matching external-weights ABI. + def _force_external_weights_for_fqn_binding(consts_size): + # Structured weights are serialized as independently named storages, + # but the generated AOTI wrapper must still use the external-weights + # ABI. That mode preserves the original constant view metadata when + # the runtime replaces the dense blob with user-managed FQN tensors. if _is_cpu_clone_active(): return True, False return orig_determine_aoti_mmap_flags(consts_size) @@ -279,7 +313,7 @@ def _codegen_device_target_aware(self, device): _codecache.TensorProperties = functools.partial( _tensor_properties_for_low_memory, original=orig_tensor_properties ) - _codecache.determine_aoti_mmap_flags = _force_external_weights_for_streaming + _codecache.determine_aoti_mmap_flags = _force_external_weights_for_fqn_binding _Autotuner.run = _autotuner_run_with_rehydrated_emptied_args prev_active = getattr(_CPU_CLONE_GUARD, "active", False) _CPU_CLONE_GUARD.active = True @@ -402,11 +436,8 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool: return value == "ON" -def _write_aoti_weights_blob(weights, blob_path: str) -> bytes: - """Stream AOTI tensor storages to disk 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) +def _write_tensor_storage(tensor: torch.Tensor, path: str) -> bytes: + """Stream one AOTI storage to ``path`` and return its SHA-256 digest.""" chunk_size = 8 * 1024 * 1024 digest = hashlib.sha256() @@ -414,35 +445,210 @@ def write_chunk(output, chunk) -> None: digest.update(chunk) output.write(chunk) - with open(blob_path, "wb") as output: - for tensor in tensors: - if tensor.is_mkldnn: - raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI") - storage = tensor.untyped_storage() - nbytes = storage.nbytes() - if nbytes and tensor.is_cuda: - byte_tensor = torch.empty( - 0, dtype=torch.uint8, device=tensor.device - ).set_(storage, 0, (nbytes,), (1,)) - for offset in range(0, nbytes, chunk_size): - cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu() - 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): - write_chunk(output, raw_view[offset : offset + chunk_size]) - del raw_view, raw_array - # Match AOTInductor's binary_blob layout: CUDA-only constants are - # packed, while CPU/mixed constants are aligned to 64 bytes. - if not all_cuda and (padding := (-nbytes) % 64): - write_chunk(output, bytes(padding)) - del storage - _trim_host_memory() + if tensor.is_mkldnn: + raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI") + storage = tensor.untyped_storage() + nbytes = storage.nbytes() + with open(path, "wb") as output: + if nbytes and tensor.is_cuda: + byte_tensor = torch.empty(0, dtype=torch.uint8, device=tensor.device).set_( + storage, 0, (nbytes,), (1,) + ) + for offset in range(0, nbytes, chunk_size): + cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu() + 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): + write_chunk(output, raw_view[offset : offset + chunk_size]) + del raw_view, raw_array + del storage return digest.digest() +def _materialize_fqn_weights( # noqa: C901 + weights: Any, + directory: str, + mutated_fqns: set[str], +) -> _FqnWeightArtifact: + """Turn AOTI ``Weights`` into content-addressed storage files + views.""" + # The graph can contain hundreds of independent storages. Trimming around + # every storage is both ineffective (``records`` below still owns all of + # the tensors) and extremely expensive for large exported graphs. Trim + # once at artifact boundaries; streamed CPU chunks are released by normal + # reference counting as they are replaced. + _trim_host_memory() + entries: List[_FqnWeightEntry] = [] + storages: Dict[str, FileBackedData] = {} + records: List[Tuple[str, torch.Tensor, Any, Tuple[Any, ...], int]] = [] + record_indices_by_identity: Dict[Tuple[Any, ...], List[int]] = {} + + for index, (fqn, (tensor, properties)) in enumerate(weights.items()): + storage = tensor.untyped_storage() + storage_nbytes = storage.nbytes() + storage_ptr = storage.data_ptr() + property_storage_ptr = getattr(properties, "storage_ptr", None) + if property_storage_ptr not in (None, 0): + # TensorProperties describes the graph constant's real storage. + # The value tensor can be a clone (including a CPU clone in CUDA + # low-memory mode), so its data_ptr is not a stable alias key. + identity = ( + "aoti", + int(property_storage_ptr), + str(tensor.dtype), + ) + else: + identity = ( + tensor.device.type, + tensor.device.index if tensor.device.index is not None else -1, + storage_ptr if storage_ptr != 0 else -(index + 1), + storage_nbytes, + ) + del storage + records.append((fqn, tensor, properties, identity, storage_nbytes)) + record_indices_by_identity.setdefault(identity, []).append(index) + + storage_info_by_identity: Dict[Tuple[Any, ...], Tuple[str, int, int]] = {} + for storage_group, (identity, record_indices) in enumerate( + record_indices_by_identity.items() + ): + # AOTI's value can be a clone of a view. Pick the largest available + # backing storage in the alias group so every declared view can be + # reconstructed from the one serialized allocation. + candidate_index = max(record_indices, key=lambda item: records[item][4]) + candidate_tensor = records[candidate_index][1] + storage_nbytes = records[candidate_index][4] + expected_storage_nbytes = max( + ( + int(storage_size) + for item in record_indices + if (storage_size := getattr(records[item][2], "storage_size", None)) + is not None + ), + default=0, + ) + if storage_nbytes < expected_storage_nbytes: + raise RuntimeError( + "AOTI cloned storage is smaller than its TensorProperties " + f"({storage_nbytes} < {expected_storage_nbytes} bytes)" + ) + + fd, storage_path = tempfile.mkstemp( + prefix=".cuda_weight_", suffix=".storage", dir=directory + ) + os.close(fd) + try: + digest = _write_tensor_storage(candidate_tensor, storage_path) + storage_key = digest.hex() + "_cuda_weight_storage" + data = FileBackedData.move_from(storage_path, sha256=digest) + except Exception: + try: + os.remove(storage_path) + except OSError: + pass + raise + + existing = storages.get(storage_key) + if existing is None: + storages[storage_key] = data + else: + data.close() + storage_info_by_identity[identity] = ( + storage_key, + storage_nbytes, + storage_group, + ) + + for fqn, tensor, properties, identity, _storage_nbytes in records: + storage_key, serialized_nbytes, storage_group = storage_info_by_identity[ + identity + ] + sizes = getattr(properties, "shape", tensor.shape) + strides = getattr(properties, "stride", tensor.stride()) + storage_offset = getattr(properties, "offset", tensor.storage_offset()) + sizes = tuple(int(size) for size in sizes) + strides = tuple(int(stride) for stride in strides) + storage_offset = int(storage_offset) + if ( + len(sizes) != len(strides) + or storage_offset < 0 + or any(size < 0 for size in sizes) + or any(stride < 0 for stride in strides) + ): + raise RuntimeError(f"AOTI view {fqn!r} has invalid tensor metadata") + required_nbytes = 0 + if all(size != 0 for size in sizes): + last_element = storage_offset + sum( + stride * (size - 1) for size, stride in zip(sizes, strides) + ) + required_nbytes = (last_element + 1) * tensor.element_size() + if required_nbytes > serialized_nbytes: + raise RuntimeError( + f"AOTI view {fqn!r} requires {required_nbytes} bytes from a " + f"{serialized_nbytes}-byte cloned storage" + ) + entries.append( + _FqnWeightEntry( + fqn=fqn, + storage_key=storage_key, + storage_group=storage_group, + storage_nbytes=serialized_nbytes, + dtype=int(scalar_type_enum(tensor.dtype)), + storage_offset=storage_offset, + sizes=sizes, + strides=strides, + shareable=fqn not in mutated_fqns, + ) + ) + + # A mutable view makes its complete physical storage stateful. The runtime + # shares such storages by FQN (not by content hash), including aliases that + # share the same backing buffer. + local_storage_groups = { + entry.storage_group for entry in entries if not entry.shareable + } + for entry in entries: + if entry.storage_group in local_storage_groups: + entry.shareable = False + + _trim_host_memory() + return _FqnWeightArtifact(entries=entries, storages=storages) + + +def _encode_fqn_weight_manifest( + so_blob_key: str, entries: List[_FqnWeightEntry] +) -> bytes: + """Encode the CUDA per-storage manifest consumed by the runtime.""" + output = bytearray(_FQN_WEIGHTS_MAGIC) + + def write_string(value: str) -> None: + encoded = value.encode("utf-8") + output.extend(struct.pack(" bool: @classmethod def save_data_externally(cls) -> bool: """ - CUDA backend saves SO blob and weights blob to an external .ptd file. + CUDA backend saves weight storages (and, when configured, SO blobs) in + external named data such as a .ptd file. This file must be provided at runtime via --data_path argument. """ return True + @classmethod + def preprocess( # noqa: C901 + cls, edge_program: Any, compile_specs: List[CompileSpec] + ) -> PreprocessResult: + """Compile CUDA weights as independently addressable AOTI storages.""" + mutated_fqns = set( + getattr(edge_program.graph_signature, "buffers_to_mutate", {}).values() + ) + previous_capture = getattr(_FQN_WEIGHTS_CAPTURE, "current", None) + capture = _FqnWeightCapture(mutated_fqns=mutated_fqns) + _FQN_WEIGHTS_CAPTURE.current = capture + try: + result = super().preprocess(edge_program, compile_specs) + finally: + _FQN_WEIGHTS_CAPTURE.current = previous_capture + + artifact = capture.artifact + if artifact is None: + raise RuntimeError("CUDA AOTI did not return a structured Weights output") + if result.data_store_output is None: + raise RuntimeError("CUDA AOTI preprocess returned no named data") + + try: + parent_keys = result.processed_bytes.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise RuntimeError("Malformed CUDA AOTI named-data payload") from error + if not parent_keys or not parent_keys[0]: + raise RuntimeError("CUDA AOTI payload is missing its shared-object key") + so_blob_key = parent_keys[0] + compatibility_blob_key = parent_keys[1] if len(parent_keys) > 1 else None + + # Rebuild AotiBackend's store without the empty compatibility blob, + # then add each physical weight storage as separately named external + # data. This leaves a new PTD containing only real storages while the + # legacy runtime path remains able to consume old dense blobs. + parent_store = result.data_store_output + named_data_store = NamedDataStore() + for key, entry in parent_store.pte_data.items(): + if key != compatibility_blob_key: + named_data_store.add_named_data( + key, + parent_store.buffers[entry.buffer_index], + alignment=entry.alignment, + tensor_layout=entry.tensor_layout, + ) + for tag, entries in parent_store.external_data.items(): + for key, entry in entries.items(): + if key != compatibility_blob_key: + named_data_store.add_named_data( + key, + parent_store.buffers[entry.buffer_index], + alignment=entry.alignment, + external_tag=tag, + tensor_layout=entry.tensor_layout, + ) + + external_tag = f"aoti_{cls.get_device_name()}_blob" + for storage_key, data in artifact.storages.items(): + named_data_store.add_named_data( + storage_key, data, alignment=1, external_tag=external_tag + ) + + result.processed_bytes = _encode_fqn_weight_manifest( + so_blob_key, artifact.entries + ) + result.data_store_output = named_data_store.get_named_data_store_output() + return result + @classmethod def load_weights_blob( cls, blob_path: str, compile_specs: List[CompileSpec] ) -> tuple[Any, str]: - """Keep low-memory CUDA weights file-backed during PTE serialization. + """Keep low-memory CUDA named data file-backed during serialization. - The streamed file has the same layout as AOTInductor's ``binary_blob``. - Keeping it file-backed avoids reading another model-sized copy into - host memory without changing its bytes. + New FQN artifacts use this path only for AotiBackend's empty + compatibility placeholder. Legacy binary-blob handling remains + unchanged for callers that still provide a real blob. """ + known_hash = cls._materialized_blob_hashes.pop(blob_path, None) if not cls._is_low_memory_mode(compile_specs): return super().load_weights_blob(blob_path, compile_specs) - 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() @@ -583,7 +858,7 @@ def load_weights_blob( def materialize_weights_blob( cls, paths: Any, compile_specs: List[CompileSpec] ) -> Any: - if not cls._is_low_memory_mode(compile_specs) or not isinstance(paths, list): + if not isinstance(paths, list): return paths from torch.export.pt2_archive._package_weights import Weights @@ -607,13 +882,24 @@ def materialize_weights_blob( if so_path is None: raise RuntimeError(f"Expected a CUDA AOTI .wrapper.so output, got {paths}") blob_path = os.path.splitext(so_path)[0] + "_weights.blob" - cls._materialized_blob_hashes[blob_path] = _write_aoti_weights_blob( - weights[0], blob_path + capture = getattr(_FQN_WEIGHTS_CAPTURE, "current", None) + if capture is None: + raise RuntimeError( + "CUDA structured weights must be materialized inside preprocess" + ) + capture.artifact = _materialize_fqn_weights( + weights[0], os.path.dirname(blob_path), capture.mutated_fqns ) - # Forcing the external-weights ABI makes Inductor emit an empty blob - # path alongside the Weights object. Replace that file in place and do - # not add a duplicate path to the returned package outputs. + # Keep AotiBackend's existing path contract intact. The compatibility + # blob is empty and ignored by the versioned CUDA runtime path; old + # artifacts continue to carry and load their original dense blob. + with open(blob_path, "wb"): + pass + cls._materialized_blob_hashes[blob_path] = hashlib.sha256(b"").digest() + + # Replace the structured Weights output with the compatibility path + # expected by AotiBackend's existing named-data packaging contract. materialized = [path for path in paths if not isinstance(path, Weights)] if blob_path not in materialized: materialized.append(blob_path) @@ -754,10 +1040,8 @@ def get_aoti_compile_options( # Separate weight constants from the .so file "aot_inductor.package": True, "aot_inductor.package_constants_in_so": False, - # Store weight constants on disk in a binary blob. Low-memory mode - # asks AOTI for a Weights object and streams the equivalent blob in - # materialize_weights_blob; its context also forces the generated - # wrapper to use the required external-weights ABI. + # Ask AOTI for structured constants. CUDABackend converts these to + # independently named physical storages plus an FQN view manifest. "aot_inductor.package_constants_on_disk_format": cls._weights_format( compile_specs ), @@ -892,11 +1176,10 @@ def _is_low_memory_mode(compile_specs: List[CompileSpec]) -> bool: @classmethod def _weights_format(cls, compile_specs: List[CompileSpec]) -> str: - return ( - "pickle_weights" - if cls._is_low_memory_mode(compile_specs) - else "binary_blob" - ) + # CUDA consumes the structured AOTI output directly and emits a + # versioned per-storage manifest. This is backend-wide rather than a + # model/export-script option. + return "pickle_weights" @classmethod def move_program_to_device( diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 349082ad690..d355cf67f1e 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -17,10 +17,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -46,6 +48,7 @@ #include #include #include +#include #include #include #include @@ -165,8 +168,8 @@ class ET_EXPERIMENTAL CudaBackend final return shared_cuda_stream_ != nullptr; } - // Enable cross-method per-FQN weight caching. Set via the - // kWeightSharingAcrossMethods runtime backend option. + // Enable the legacy dense-blob per-FQN cache. New manifest artifacts carry + // enough ownership metadata to share immutable storages automatically. void set_weight_sharing_across_methods(bool enabled) { weight_sharing_across_methods_.store(enabled, std::memory_order_relaxed); } @@ -177,7 +180,7 @@ class ET_EXPERIMENTAL CudaBackend final Error load_function_pointers_into_handle( void* so_handle, - AOTIDelegateHandle* handle) const { + cuda::CudaDelegateHandle* handle) const { #define LOAD_SYMBOL(member, name) \ do { \ auto symbol_res = get_function(so_handle, #name); \ @@ -226,6 +229,10 @@ class ET_EXPERIMENTAL CudaBackend final LOAD_OPTIONAL_SYMBOL( get_constant_original_fqn, AOTInductorModelContainerGetConstantOriginalFQN); + LOAD_OPTIONAL_SYMBOL( + get_constant_dtype, AOTInductorModelContainerGetConstantDtype); + LOAD_OPTIONAL_SYMBOL( + get_constant_data_size, AOTInductorModelContainerGetConstantDataSize); LOAD_OPTIONAL_SYMBOL( extract_constants_map, AOTInductorModelContainerExtractConstantsMap); LOAD_OPTIONAL_SYMBOL( @@ -327,10 +334,21 @@ class ET_EXPERIMENTAL CudaBackend final std::string so_blob_key; std::string weights_blob_key; - ET_CHECK_OK_OR_RETURN_ERROR( - executorch::backends::aoti::resolve_blob_keys( - processed, method_name, so_blob_key, weights_blob_key), - "Malformed named-data key payload"); + CudaFqnWeightManifest fqn_weight_manifest; + const bool has_fqn_weights = + is_cuda_fqn_weight_manifest(processed->data(), processed->size()); + if (has_fqn_weights) { + ET_CHECK_OK_OR_RETURN_ERROR( + parse_cuda_fqn_weight_manifest( + processed->data(), processed->size(), fqn_weight_manifest), + "Malformed CUDA FQN weight manifest"); + so_blob_key = fqn_weight_manifest.so_blob_key; + } else { + ET_CHECK_OK_OR_RETURN_ERROR( + executorch::backends::aoti::resolve_blob_keys( + processed, method_name, so_blob_key, weights_blob_key), + "Malformed named-data key payload"); + } const NamedDataMap* named_data_map = context.get_named_data_map(); auto aoti_dso_buffer = named_data_map->get_data(so_blob_key.c_str()); @@ -404,6 +422,14 @@ class ET_EXPERIMENTAL CudaBackend final handle->container_handle = container_handle; + // Versioned FQN artifacts carry complete storage/view and mutability + // metadata, so immutable storages are safely shared without a load-order + // contract. Legacy artifacts keep their historical dense-blob behavior + // and runtime option semantics. + if (has_fqn_weights) { + ET_CHECK_OK_OR_RETURN_ERROR(load_constants_from_fqn_manifest( + handle, named_data_map, fqn_weight_manifest)); + } // Load constants. When weight_sharing_across_methods is enabled (opt-in // via the kWeightSharingAcrossMethods runtime backend option set by the // runner), use the per-weight FQN cache so methods that share weights @@ -411,7 +437,7 @@ class ET_EXPERIMENTAL CudaBackend final // back to the legacy per-method blob load — required for models whose // methods are independent sub-graphs that may have FQN collisions // (e.g. parakeet). - if (is_weight_sharing_across_methods_enabled()) { + else if (is_weight_sharing_across_methods_enabled()) { ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache( handle, named_data_map, method_name, weights_blob_key)); } else { @@ -899,9 +925,9 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::mutex cuda_stream_mutex_; std::shared_ptr shared_cuda_stream_ = nullptr; - // Whether to enable cross-method per-FQN weight caching at init time. + // Whether to enable cross-method caching for legacy dense-blob artifacts. // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default - // OFF — see set_weight_sharing_across_methods() for safety constraints. + // OFF; versioned manifest artifacts do not consult this option. std::atomic weight_sharing_across_methods_{false}; // --------------------------------------------------------------- @@ -1013,6 +1039,427 @@ class ET_EXPERIMENTAL CudaBackend final return Error::Ok; } + static Error validate_fqn_weight_view(const CudaFqnWeightEntry& entry) { + uint64_t item_size = 0; + switch (static_cast(entry.dtype)) { + case slim::c10::ScalarType::Byte: + case slim::c10::ScalarType::Char: + case slim::c10::ScalarType::Bool: + item_size = 1; + break; + case slim::c10::ScalarType::Short: + case slim::c10::ScalarType::Half: + case slim::c10::ScalarType::BFloat16: + item_size = 2; + break; + case slim::c10::ScalarType::Int: + case slim::c10::ScalarType::Float: + item_size = 4; + break; + case slim::c10::ScalarType::Long: + item_size = 8; + break; + default: + return Error::InvalidProgram; + } + + ET_CHECK_OR_RETURN_ERROR( + entry.storage_nbytes <= std::numeric_limits::max(), + InvalidProgram, + "CUDA FQN storage '%s' is too large for this platform", + entry.storage_key.c_str()); + + bool empty = false; + uint64_t last_element = static_cast(entry.storage_offset); + for (size_t dim = 0; dim < entry.sizes.size(); ++dim) { + const uint64_t size = static_cast(entry.sizes[dim]); + const uint64_t stride = static_cast(entry.strides[dim]); + if (size == 0) { + empty = true; + break; + } + const uint64_t extent = size - 1; + ET_CHECK_OR_RETURN_ERROR( + extent == 0 || + stride <= std::numeric_limits::max() / extent, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing shape/stride metadata", + entry.fqn.c_str()); + const uint64_t span = stride * extent; + ET_CHECK_OR_RETURN_ERROR( + last_element <= std::numeric_limits::max() - span, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing storage metadata", + entry.fqn.c_str()); + last_element += span; + } + + uint64_t required_nbytes = 0; + if (!empty) { + ET_CHECK_OR_RETURN_ERROR( + last_element < std::numeric_limits::max() && + last_element + 1 <= + std::numeric_limits::max() / item_size, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing storage size", + entry.fqn.c_str()); + required_nbytes = (last_element + 1) * item_size; + } + ET_CHECK_OR_RETURN_ERROR( + required_nbytes <= entry.storage_nbytes, + InvalidProgram, + "CUDA FQN weight '%s' requires %llu bytes from a %llu-byte storage", + entry.fqn.c_str(), + static_cast(required_nbytes), + static_cast(entry.storage_nbytes)); + return Error::Ok; + } + + Error acquire_fqn_weight_storage( + const NamedDataMap* named_data_map, + const CudaFqnWeightEntry& entry, + const std::vector* mutable_group, + int device_index, + std::shared_ptr& storage, + bool& reused) const { + reused = false; + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "CUDA FQN weights require a named data map"); + + uintptr_t mutable_scope = 0; + if (!entry.shareable) { + // Method::init wraps the same external PTD map in a distinct + // MergedDataMap for every method, so the wrapper address is not a model + // instance identity. get_key(), however, forwards the pointer owned by + // the underlying PTD map. That pointer is stable across methods, unique + // to a live PTD instance, and valid for the map's lifetime. + auto num_keys = named_data_map->get_num_keys(); + ET_CHECK_OR_RETURN_ERROR( + num_keys.ok(), + InvalidProgram, + "Failed to enumerate CUDA named data while loading mutable FQN storage"); + for (uint32_t index = 0; index < num_keys.get(); ++index) { + auto key = named_data_map->get_key(index); + ET_CHECK_OR_RETURN_ERROR( + key.ok(), + InvalidProgram, + "Failed to read CUDA named data key %u", + index); + if (entry.storage_key == key.get()) { + mutable_scope = reinterpret_cast(key.get()); + break; + } + } + ET_CHECK_OR_RETURN_ERROR( + mutable_scope != 0, + NotFound, + "CUDA mutable FQN storage '%s' is missing from named data", + entry.storage_key.c_str()); + } + + const auto cache_key = [&](const CudaFqnWeightEntry& item) { + if (item.shareable) { + // Immutable bytes are safe to reuse across methods and model + // instances solely by content identity. + return std::string("immutable:") + item.storage_key + + "@cuda:" + std::to_string(device_index); + } + // Stateful buffers with identical initial bytes are not interchangeable. + // Scope their logical FQN identity to the underlying PTD instance so + // methods in one model share state without leaking it to another live + // model instance. + return std::string("mutable:") + std::to_string(mutable_scope) + ":" + + item.fqn + "@cuda:" + std::to_string(device_index); + }; + + std::vector cache_keys; + if (entry.shareable) { + cache_keys.push_back(cache_key(entry)); + } else { + ET_CHECK_OR_RETURN_ERROR( + mutable_group != nullptr && !mutable_group->empty(), + InvalidProgram, + "CUDA mutable FQN storage group %u is empty", + entry.storage_group); + cache_keys.reserve(mutable_group->size()); + for (const CudaFqnWeightEntry* alias : *mutable_group) { + ET_CHECK_OR_RETURN_ERROR( + alias != nullptr && !alias->shareable && + alias->storage_nbytes == entry.storage_nbytes, + InvalidProgram, + "CUDA mutable FQN storage group %u is inconsistent", + entry.storage_group); + cache_keys.push_back(cache_key(*alias)); + } + } + + std::unique_lock cache_lock(fqn_weight_storage_mutex_); + for (const std::string& key : cache_keys) { + auto cached = shared_fqn_weight_storages_.find(key); + if (cached != shared_fqn_weight_storages_.end()) { + std::shared_ptr candidate = cached->second.lock(); + if (candidate != nullptr) { + ET_CHECK_OR_RETURN_ERROR( + candidate->nbytes == entry.storage_nbytes, + InvalidProgram, + "CUDA FQN storage '%s' has inconsistent sizes (%zu vs %llu)", + entry.storage_key.c_str(), + candidate->nbytes, + static_cast(entry.storage_nbytes)); + ET_CHECK_OR_RETURN_ERROR( + storage == nullptr || storage.get() == candidate.get(), + InvalidProgram, + "CUDA mutable FQN storage group %u resolves to multiple allocations", + entry.storage_group); + storage = std::move(candidate); + } + } + } + if (storage != nullptr) { + for (const std::string& key : cache_keys) { + shared_fqn_weight_storages_[key] = storage; + } + reused = true; + return Error::Ok; + } + + void* device_data = nullptr; + const size_t allocation_size = + std::max(1, static_cast(entry.storage_nbytes)); + const cudaError_t allocation_error = + cudaMalloc(&device_data, allocation_size); + if (allocation_error != cudaSuccess) { + ET_LOG( + Error, + "cudaMalloc failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(allocation_error)); + return Error::MemoryAllocationFailed; + } + + if (entry.storage_nbytes > 0) { + auto host_data = named_data_map->get_data(entry.storage_key.c_str()); + if (!host_data.ok()) { + cudaFree(device_data); + ET_LOG( + Error, + "CUDA FQN storage '%s' is missing", + entry.storage_key.c_str()); + return Error::NotFound; + } + if (host_data->size() != entry.storage_nbytes) { + cudaFree(device_data); + ET_LOG( + Error, + "CUDA FQN storage '%s' has size %zu, expected %llu", + entry.storage_key.c_str(), + host_data->size(), + static_cast(entry.storage_nbytes)); + return Error::InvalidProgram; + } + const cudaError_t copy_error = cudaMemcpy( + device_data, + host_data->data(), + static_cast(entry.storage_nbytes), + cudaMemcpyHostToDevice); + host_data->Free(); + if (copy_error != cudaSuccess) { + cudaFree(device_data); + ET_LOG( + Error, + "cudaMemcpy failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(copy_error)); + return Error::Internal; + } + } + + storage = std::make_shared( + device_data, static_cast(entry.storage_nbytes), device_index); + for (const std::string& key : cache_keys) { + shared_fqn_weight_storages_[key] = storage; + } + return Error::Ok; + } + + Error load_constants_from_fqn_manifest( + cuda::CudaDelegateHandle* handle, + const NamedDataMap* named_data_map, + const CudaFqnWeightManifest& manifest) const { + ET_CHECK_OR_RETURN_ERROR( + handle->get_num_constants && handle->get_constant_name && + handle->get_constant_original_fqn && handle->get_constant_dtype && + handle->get_constant_data_size && + handle->update_user_managed_constant_buffer_pairs, + NotSupported, + "AOTI library does not expose the APIs required by CUDA FQN weights"); + + size_t num_constants = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_num_constants(handle->container_handle, &num_constants), + "Failed to enumerate CUDA AOTI constants"); + std::unordered_map> + fqn_to_internal_names; + std::unordered_map> + fqn_to_aoti_metadata; + for (size_t index = 0; index < num_constants; ++index) { + const char* internal_name = nullptr; + const char* fqn = nullptr; + int32_t dtype = 0; + size_t data_size = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_name( + handle->container_handle, index, &internal_name), + "Failed to read CUDA AOTI constant name at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_original_fqn( + handle->container_handle, index, &fqn), + "Failed to read CUDA AOTI constant FQN at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_dtype(handle->container_handle, index, &dtype), + "Failed to read CUDA AOTI constant dtype at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_data_size( + handle->container_handle, index, &data_size), + "Failed to read CUDA AOTI constant size at index %zu", + index); + if (internal_name != nullptr && fqn != nullptr && fqn[0] != '\0') { + fqn_to_internal_names[fqn].emplace_back(internal_name); + auto [metadata, inserted] = + fqn_to_aoti_metadata.emplace(fqn, std::make_pair(dtype, data_size)); + ET_CHECK_OR_RETURN_ERROR( + inserted || + (metadata->second.first == dtype && + metadata->second.second == data_size), + InvalidProgram, + "CUDA AOTI constant FQN '%s' has inconsistent metadata", + fqn); + } + } + + int device_index = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index)); + struct LocalStorage { + std::string storage_key; + uint64_t storage_nbytes; + std::shared_ptr storage; + }; + std::unordered_map local_storages; + std::unordered_map> + mutable_storage_groups; + for (const CudaFqnWeightEntry& entry : manifest.entries) { + if (!entry.shareable) { + mutable_storage_groups[entry.storage_group].push_back(&entry); + } + } + std::vector pairs; + pairs.reserve(manifest.entries.size()); + std::unordered_set bound_fqns; + size_t reused_storages = 0; + handle->fqn_weight_tensors.reserve(manifest.entries.size()); + + for (const CudaFqnWeightEntry& entry : manifest.entries) { + ET_CHECK_OK_OR_RETURN_ERROR( + validate_fqn_weight_view(entry), + "Invalid CUDA FQN view '%s'", + entry.fqn.c_str()); + auto internal_names = fqn_to_internal_names.find(entry.fqn); + ET_CHECK_OR_RETURN_ERROR( + internal_names != fqn_to_internal_names.end(), + InvalidProgram, + "CUDA FQN weight '%s' is not present in its AOTI library", + entry.fqn.c_str()); + const auto aoti_metadata = fqn_to_aoti_metadata.find(entry.fqn); + ET_CHECK_OR_RETURN_ERROR( + aoti_metadata != fqn_to_aoti_metadata.end() && + aoti_metadata->second.first == entry.dtype && + aoti_metadata->second.second == entry.storage_nbytes, + InvalidProgram, + "CUDA FQN weight '%s' metadata does not match its AOTI library", + entry.fqn.c_str()); + ET_CHECK_OR_RETURN_ERROR( + bound_fqns.emplace(entry.fqn).second, + InvalidProgram, + "CUDA FQN weight '%s' appears more than once in its manifest", + entry.fqn.c_str()); + + const std::string local_key = entry.shareable + ? "shared:" + entry.storage_key + : "local:" + std::to_string(entry.storage_group); + auto local_storage = local_storages.find(local_key); + std::shared_ptr storage; + if (local_storage == local_storages.end()) { + bool reused = false; + const auto mutable_entries = + mutable_storage_groups.find(entry.storage_group); + const std::vector* mutable_group = + entry.shareable ? nullptr + : (mutable_entries == mutable_storage_groups.end() + ? nullptr + : &mutable_entries->second); + ET_CHECK_OK_OR_RETURN_ERROR( + acquire_fqn_weight_storage( + named_data_map, + entry, + mutable_group, + device_index, + storage, + reused), + "Failed to load CUDA FQN storage '%s'", + entry.storage_key.c_str()); + reused_storages += reused ? 1 : 0; + local_storages.emplace( + local_key, + LocalStorage{entry.storage_key, entry.storage_nbytes, storage}); + handle->fqn_weight_storages.push_back(storage); + } else { + ET_CHECK_OR_RETURN_ERROR( + local_storage->second.storage_key == entry.storage_key && + local_storage->second.storage_nbytes == entry.storage_nbytes, + InvalidProgram, + "CUDA FQN storage group %u has inconsistent backing storage", + entry.storage_group); + storage = local_storage->second.storage; + } + + auto tensor = std::make_unique(slim::from_blob( + storage->data, + slim::makeArrayRef(entry.sizes), + slim::makeArrayRef(entry.strides), + static_cast(entry.dtype), + Device(slim::c10::DeviceType::CUDA, device_index), + entry.storage_offset)); + AtenTensorHandle tensor_handle = + reinterpret_cast(tensor.get()); + handle->fqn_weight_tensors.push_back(std::move(tensor)); + for (const std::string& internal_name : internal_names->second) { + pairs.push_back({internal_name.c_str(), tensor_handle}); + } + } + + ET_CHECK_OK_OR_RETURN_ERROR( + handle->update_user_managed_constant_buffer_pairs( + handle->container_handle, + pairs.data(), + pairs.size(), + /*use_inactive=*/false, + /*validate_full_update=*/true), + "Failed to bind CUDA FQN weights"); + ET_LOG( + Info, + "Loaded %zu CUDA FQN views from %zu physical storages (%zu reused " + "across methods)", + manifest.entries.size(), + local_storages.size(), + reused_storages); + return Error::Ok; + } + // Load constants for a method using per-weight caching. // Returns Error::Ok on success. // @@ -1266,6 +1713,13 @@ class ET_EXPERIMENTAL CudaBackend final // explicitly deleted — see destroy() comment). mutable std::unordered_map shared_constant_tensors_; + + // New-format artifacts share immutable physical storages by their + // content-addressed named-data key. Weak ownership lets the allocation be + // reclaimed after the last delegate using it is destroyed. + mutable std::mutex fqn_weight_storage_mutex_; + mutable std::unordered_map> + shared_fqn_weight_storages_; }; } // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 83d88b65c5a..ffbfafa097d 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include #include @@ -17,6 +18,44 @@ namespace executorch { namespace backends { namespace cuda { +using AOTInductorModelContainerGetConstantDtypeFunc = + aoti::AOTIRuntimeError (*)( + aoti::AOTInductorModelContainerHandle container_handle, + size_t idx, + int32_t* dtype); +using AOTInductorModelContainerGetConstantDataSizeFunc = + aoti::AOTIRuntimeError (*)( + aoti::AOTInductorModelContainerHandle container_handle, + size_t idx, + size_t* data_size); + +struct CudaWeightStorage { + void* data{nullptr}; + size_t nbytes{0}; + int device_index{0}; + + CudaWeightStorage(void* data_, size_t nbytes_, int device_index_) + : data(data_), nbytes(nbytes_), device_index(device_index_) {} + + ~CudaWeightStorage() { + if (data == nullptr) { + return; + } + int previous_device = 0; + const cudaError_t get_device_error = cudaGetDevice(&previous_device); + if (get_device_error == cudaSuccess && previous_device != device_index) { + (void)cudaSetDevice(device_index); + } + (void)cudaFree(data); + if (get_device_error == cudaSuccess && previous_device != device_index) { + (void)cudaSetDevice(previous_device); + } + } + + CudaWeightStorage(const CudaWeightStorage&) = delete; + CudaWeightStorage& operator=(const CudaWeightStorage&) = delete; +}; + // Shared CUDA stream wrapper with proper RAII cleanup. // This ensures the stream is destroyed when all handles using it are destroyed. struct CudaStreamDeleter { @@ -148,6 +187,11 @@ struct CudaGraphState { // CUDA-specific delegate handle that extends AOTIDelegateHandle. // This consolidates CUDA stream management into a single location. struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { + // Extra AOTI metadata used to validate per-FQN manifests before binding. + AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr}; + AOTInductorModelContainerGetConstantDataSizeFunc get_constant_data_size{ + nullptr}; + // CUDA stream for this handle, support both shared mode and single mode. // In shared mode, all cuda delegate handles share the same stream (e.g., for // skip-copy optimization), they will all hold a reference to the same @@ -168,6 +212,11 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { // CUDA graph state (warmup, capture, replay, static buffers) CudaGraphState cuda_graph_state; + + // Per-storage weight artifacts keep the CUDA allocations and the original + // SlimTensor handles alive for as long as AOTI may reference their views. + std::vector> fqn_weight_storages; + std::vector> fqn_weight_tensors; }; } // namespace cuda diff --git a/backends/cuda/runtime/cuda_weight_manifest.h b/backends/cuda/runtime/cuda_weight_manifest.h new file mode 100644 index 00000000000..325a8b25fa7 --- /dev/null +++ b/backends/cuda/runtime/cuda_weight_manifest.h @@ -0,0 +1,217 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include + +namespace executorch::backends::cuda { + +constexpr char kCudaFqnWeightsMagic[] = "ETCUDAFQN1"; +constexpr size_t kCudaFqnWeightsMagicSize = sizeof(kCudaFqnWeightsMagic) - 1; + +struct CudaFqnWeightEntry { + std::string fqn; + std::string storage_key; + uint32_t storage_group{0}; + uint64_t storage_nbytes{0}; + int32_t dtype{0}; + int64_t storage_offset{0}; + std::vector sizes; + std::vector strides; + bool shareable{false}; +}; + +struct CudaFqnWeightManifest { + std::string so_blob_key; + std::vector entries; +}; + +inline bool is_supported_cuda_fqn_dtype(int32_t dtype) { + // Values match c10::ScalarType and the slim AOTI runtime. + switch (dtype) { + case 0: // Byte + case 1: // Char + case 2: // Short + case 3: // Int + case 4: // Long + case 5: // Half + case 6: // Float + case 11: // Bool + case 15: // BFloat16 + return true; + default: + return false; + } +} + +inline bool is_cuda_fqn_weight_manifest(const void* data, size_t size) { + return data != nullptr && size >= kCudaFqnWeightsMagicSize && + std::memcmp(data, kCudaFqnWeightsMagic, kCudaFqnWeightsMagicSize) == 0; +} + +namespace detail { + +class CudaWeightManifestReader final { + public: + CudaWeightManifestReader(const void* data, size_t size) + : cursor_(static_cast(data)), end_(cursor_ + size) {} + + bool skip(size_t size) { + if (remaining() < size) { + return false; + } + cursor_ += size; + return true; + } + + bool read_u8(uint8_t& value) { + if (remaining() < 1) { + return false; + } + value = *cursor_++; + return true; + } + + bool read_u32(uint32_t& value) { + uint64_t wide = 0; + if (!read_unsigned(wide, 4)) { + return false; + } + value = static_cast(wide); + return true; + } + + bool read_i32(int32_t& value) { + uint32_t raw = 0; + if (!read_u32(raw)) { + return false; + } + std::memcpy(&value, &raw, sizeof(value)); + return true; + } + + bool read_u64(uint64_t& value) { + return read_unsigned(value, 8); + } + + bool read_i64(int64_t& value) { + uint64_t raw = 0; + if (!read_u64(raw)) { + return false; + } + std::memcpy(&value, &raw, sizeof(value)); + return true; + } + + bool read_string(std::string& value) { + uint32_t size = 0; + if (!read_u32(size) || remaining() < size) { + return false; + } + value.assign(reinterpret_cast(cursor_), size); + cursor_ += size; + return true; + } + + bool empty() const { + return cursor_ == end_; + } + + private: + size_t remaining() const { + return static_cast(end_ - cursor_); + } + + bool read_unsigned(uint64_t& value, size_t width) { + if (remaining() < width) { + return false; + } + value = 0; + for (size_t index = 0; index < width; ++index) { + value |= static_cast(cursor_[index]) << (index * 8); + } + cursor_ += width; + return true; + } + + const uint8_t* cursor_; + const uint8_t* end_; +}; + +} // namespace detail + +inline executorch::runtime::Error parse_cuda_fqn_weight_manifest( + const void* data, + size_t size, + CudaFqnWeightManifest& manifest) { + using executorch::runtime::Error; + if (!is_cuda_fqn_weight_manifest(data, size)) { + return Error::InvalidProgram; + } + + detail::CudaWeightManifestReader reader(data, size); + if (!reader.skip(kCudaFqnWeightsMagicSize) || + !reader.read_string(manifest.so_blob_key) || + manifest.so_blob_key.empty()) { + return Error::InvalidProgram; + } + + uint32_t num_entries = 0; + constexpr uint32_t kMaxManifestEntries = 1U << 20; + if (!reader.read_u32(num_entries) || num_entries > kMaxManifestEntries) { + return Error::InvalidProgram; + } + manifest.entries.clear(); + manifest.entries.reserve(num_entries); + + constexpr uint32_t kMaxTensorDimensions = 64; + for (uint32_t index = 0; index < num_entries; ++index) { + CudaFqnWeightEntry entry; + uint32_t ndim = 0; + uint8_t shareable = 0; + if (!reader.read_string(entry.fqn) || entry.fqn.empty() || + !reader.read_string(entry.storage_key) || entry.storage_key.empty() || + !reader.read_u32(entry.storage_group) || + !reader.read_u64(entry.storage_nbytes) || + !reader.read_i32(entry.dtype) || + !is_supported_cuda_fqn_dtype(entry.dtype) || + !reader.read_i64(entry.storage_offset) || !reader.read_u32(ndim) || + ndim > kMaxTensorDimensions) { + return Error::InvalidProgram; + } + + entry.sizes.resize(ndim); + entry.strides.resize(ndim); + for (uint32_t dim = 0; dim < ndim; ++dim) { + if (!reader.read_i64(entry.sizes[dim]) || entry.sizes[dim] < 0) { + return Error::InvalidProgram; + } + } + for (uint32_t dim = 0; dim < ndim; ++dim) { + if (!reader.read_i64(entry.strides[dim]) || entry.strides[dim] < 0) { + return Error::InvalidProgram; + } + } + if (!reader.read_u8(shareable) || shareable > 1 || + entry.storage_offset < 0) { + return Error::InvalidProgram; + } + entry.shareable = shareable != 0; + manifest.entries.push_back(std::move(entry)); + } + + return reader.empty() ? Error::Ok : Error::InvalidProgram; +} + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index cef4988536f..0da63d55513 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -128,6 +128,7 @@ def define_common_targets(is_fbcode = False): headers = [ "cuda_delegate_handle.h", "cuda_mutable_state.h", + "cuda_weight_manifest.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) link_whole = True, @@ -180,6 +181,25 @@ def define_common_targets(is_fbcode = False): ), ) + cpp_unittest( + name = "test_cuda_weight_manifest", + srcs = [ + "test/test_cuda_weight_manifest.cpp", + ], + deps = [ + ":cuda_backend", + "//executorch/runtime/core:core", + ], + external_deps = [ + ("cuda", None, "cuda-lazy"), + ], + preprocessor_flags = ["-DCUDA_AVAILABLE=1"], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + ), + ) + cpp_unittest( name = "test_cuda_allocator", srcs = ["test/test_cuda_allocator.cpp"], diff --git a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp b/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp new file mode 100644 index 00000000000..db656482ffc --- /dev/null +++ b/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp @@ -0,0 +1,110 @@ +/* + * 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. + */ + +#include + +#include + +#include +#include +#include + +namespace cuda = ::executorch::backends::cuda; +using ::executorch::runtime::Error; + +namespace { + +void append_u32(std::vector& output, uint32_t value) { + for (size_t index = 0; index < 4; ++index) { + output.push_back(static_cast(value >> (index * 8))); + } +} + +void append_u64(std::vector& output, uint64_t value) { + for (size_t index = 0; index < 8; ++index) { + output.push_back(static_cast(value >> (index * 8))); + } +} + +void append_string(std::vector& output, const std::string& value) { + append_u32(output, static_cast(value.size())); + output.insert(output.end(), value.begin(), value.end()); +} + +std::vector valid_manifest(uint32_t dtype = 6) { + std::vector output( + cuda::kCudaFqnWeightsMagic, + cuda::kCudaFqnWeightsMagic + cuda::kCudaFqnWeightsMagicSize); + append_string(output, "so-key"); + append_u32(output, 1); // entries + append_string(output, "model.weight"); + append_string(output, "storage-key"); + append_u32(output, 7); // method-local storage group + append_u64(output, 24); // storage bytes + append_u32(output, dtype); // dtype + append_u64(output, 0); // storage offset + append_u32(output, 2); // ndim + append_u64(output, 2); + append_u64(output, 3); + append_u64(output, 3); + append_u64(output, 1); + output.push_back(1); // shareable + return output; +} + +} // namespace + +TEST(CudaWeightManifestTest, LegacyPayloadIsNotMisdetected) { + const std::string legacy = "so-key\nweights-key"; + EXPECT_FALSE(cuda::is_cuda_fqn_weight_manifest(legacy.data(), legacy.size())); +} + +TEST(CudaWeightManifestTest, ParsesVersionedManifest) { + const std::vector bytes = valid_manifest(); + cuda::CudaFqnWeightManifest manifest; + ASSERT_EQ( + cuda::parse_cuda_fqn_weight_manifest( + bytes.data(), bytes.size(), manifest), + Error::Ok); + ASSERT_EQ(manifest.so_blob_key, "so-key"); + ASSERT_EQ(manifest.entries.size(), 1u); + const auto& entry = manifest.entries[0]; + EXPECT_EQ(entry.fqn, "model.weight"); + EXPECT_EQ(entry.storage_key, "storage-key"); + EXPECT_EQ(entry.storage_group, 7u); + EXPECT_EQ(entry.storage_nbytes, 24u); + EXPECT_EQ(entry.dtype, 6); + EXPECT_EQ(entry.sizes, (std::vector{2, 3})); + EXPECT_EQ(entry.strides, (std::vector{3, 1})); + EXPECT_TRUE(entry.shareable); +} + +TEST(CudaWeightManifestTest, RejectsTruncationAndTrailingData) { + std::vector bytes = valid_manifest(); + cuda::CudaFqnWeightManifest manifest; + ASSERT_GT(bytes.size(), 1u); + EXPECT_EQ( + cuda::parse_cuda_fqn_weight_manifest( + bytes.data(), bytes.size() - 1, manifest), + Error::InvalidProgram); + bytes.push_back(0); + EXPECT_EQ( + cuda::parse_cuda_fqn_weight_manifest( + bytes.data(), bytes.size(), manifest), + Error::InvalidProgram); +} + +TEST(CudaWeightManifestTest, RejectsUnsupportedDtype) { + const std::vector bytes = + valid_manifest(7); // Double is unsupported. + cuda::CudaFqnWeightManifest manifest; + EXPECT_EQ( + cuda::parse_cuda_fqn_weight_manifest( + bytes.data(), bytes.size(), manifest), + Error::InvalidProgram); +} diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index a3e8b6bbee1..33305c85438 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -13,7 +13,12 @@ from unittest.mock import patch import torch -from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_backend import ( + _encode_fqn_weight_manifest, + _FQN_WEIGHTS_MAGIC, + _materialize_fqn_weights, + CudaBackend, +) from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.exir._serialize._cord import FileBackedData from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -27,9 +32,7 @@ class TestCudaLowMemoryExport(unittest.TestCase): @patch.object(CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True) - def test_low_memory_streaming_keeps_external_weights_abi(self, _) -> None: - from torch._inductor import codecache - + def test_all_cuda_exports_request_structured_weights(self, _) -> None: options = CudaBackend.get_aoti_compile_options( [CompileSpec("low_memory_mode", b"ON")] ) @@ -37,15 +40,14 @@ def test_low_memory_streaming_keeps_external_weights_abi(self, _) -> None: options["aot_inductor.package_constants_on_disk_format"], "pickle_weights", ) + self.assertEqual( + CudaBackend.get_aoti_compile_options([])[ + "aot_inductor.package_constants_on_disk_format" + ], + "pickle_weights", + ) - original = codecache.determine_aoti_mmap_flags - with CudaBackend.get_extra_aoti_compile_context_manager( - [CompileSpec("low_memory_mode", b"ON")] - ): - self.assertEqual(codecache.determine_aoti_mmap_flags(0), (True, False)) - self.assertIs(codecache.determine_aoti_mmap_flags, original) - - def test_low_memory_weights_are_streamed_in_binary_blob_format(self) -> None: + def test_weights_are_materialized_as_independent_storages(self) -> None: first = torch.tensor([1, 2, 3], dtype=torch.int16) second = torch.tensor([4, 5], dtype=torch.int32) weights = Weights( @@ -56,41 +58,81 @@ def test_low_memory_weights_are_streamed_in_binary_blob_format(self) -> None: ) with tempfile.TemporaryDirectory() as directory: - so_path = os.path.join(directory, "model.wrapper.so") - blob_path = os.path.join(directory, "model.wrapper_weights.blob") - # AOTI emits this empty placeholder when the wrapper is compiled - # with the external-weights ABI and the tensor values are pickled. - with open(blob_path, "wb"): - pass - - paths = CudaBackend.materialize_weights_blob( - [so_path, blob_path, weights], - [CompileSpec("low_memory_mode", b"ON")], + artifact = _materialize_fqn_weights( + weights, directory, mutated_fqns={"second"} + ) + self.assertEqual(2, len(artifact.entries)) + self.assertEqual(2, len(artifact.storages)) + self.assertTrue(artifact.entries[0].shareable) + self.assertFalse(artifact.entries[1].shareable) + self.assertEqual( + bytes(first.untyped_storage()), + artifact.storages[artifact.entries[0].storage_key].to_bytes(), + ) + self.assertEqual( + bytes(second.untyped_storage()), + artifact.storages[artifact.entries[1].storage_key].to_bytes(), + ) + + manifest = _encode_fqn_weight_manifest("so-key", artifact.entries) + self.assertTrue(manifest.startswith(_FQN_WEIGHTS_MAGIC)) + self.assertIn(b"first", manifest) + self.assertIn(b"second", manifest) + for storage in artifact.storages.values(): + storage.close() + + def test_views_share_one_physical_storage(self) -> None: + base = torch.arange(12, dtype=torch.float32).reshape(3, 4) + view = base[:, 1:] + weights = Weights( + { + "base": (base, TensorProperties(base)), + # AOTI may return a cloned value tensor; TensorProperties is + # the source of truth for reconstructing the original view. + "view": (base, TensorProperties(view)), + } + ) + + with tempfile.TemporaryDirectory() as directory: + artifact = _materialize_fqn_weights(weights, directory, set()) + self.assertEqual(1, len(artifact.storages)) + self.assertEqual( + artifact.entries[0].storage_group, + artifact.entries[1].storage_group, ) + self.assertEqual(1, artifact.entries[1].storage_offset) + self.assertEqual((3, 3), artifact.entries[1].sizes) + self.assertEqual((4, 1), artifact.entries[1].strides) + for storage in artifact.storages.values(): + storage.close() + + def test_identical_mutable_storages_remain_distinct_groups(self) -> None: + first = torch.zeros(4) + second = torch.zeros(4) + weights = Weights( + { + "first": (first, TensorProperties(first)), + "second": (second, TensorProperties(second)), + } + ) - self.assertEqual([so_path, blob_path], paths) - with open(blob_path, "rb") as blob: - data = blob.read() - expected = ( - bytes(first.untyped_storage()) - + bytes(58) - + bytes(second.untyped_storage()) - + bytes(56) + with tempfile.TemporaryDirectory() as directory: + artifact = _materialize_fqn_weights( + weights, directory, mutated_fqns={"first", "second"} ) - 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()) + self.assertEqual(1, len(artifact.storages)) + self.assertEqual( + artifact.entries[0].storage_key, + artifact.entries[1].storage_key, + ) + self.assertNotEqual( + artifact.entries[0].storage_group, + artifact.entries[1].storage_group, + ) + self.assertFalse(artifact.entries[0].shareable) + self.assertFalse(artifact.entries[1].shareable) + for storage in artifact.storages.values(): + storage.close() def test_low_memory_weights_require_wrapper_so(self) -> None: tensor = torch.tensor([1], dtype=torch.int16) From 6cf890d26cb0457d08c66f650b7faf7235acff3f Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 18 Aug 2026 23:36:28 -0700 Subject: [PATCH 02/12] Address FQN weight sharing review feedback Authored with Codex. --- backends/cuda/cuda_backend.py | 3 ++ backends/cuda/runtime/cuda_backend.cpp | 70 +++++++++++++------------- runtime/core/named_data_map.h | 5 +- 3 files changed, 42 insertions(+), 36 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 5b2a0ef737c..76d97726d6d 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -779,6 +779,9 @@ def preprocess( # noqa: C901 ) previous_capture = getattr(_FQN_WEIGHTS_CAPTURE, "current", None) capture = _FqnWeightCapture(mutated_fqns=mutated_fqns) + # AotiBackend packages weights synchronously on this thread. TLS keeps + # nested or concurrent preprocess calls isolated while that callback + # passes the structured artifact back to this invocation. _FQN_WEIGHTS_CAPTURE.current = capture try: result = super().preprocess(edge_program, compile_specs) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index d355cf67f1e..7c2fdbc9167 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -424,20 +424,14 @@ class ET_EXPERIMENTAL CudaBackend final // Versioned FQN artifacts carry complete storage/view and mutability // metadata, so immutable storages are safely shared without a load-order - // contract. Legacy artifacts keep their historical dense-blob behavior - // and runtime option semantics. + // contract. Legacy artifacts keep their historical dense-blob behavior: + // the runtime option selects the per-FQN cache, otherwise each method + // loads its own blob (required for independent methods with FQN + // collisions). if (has_fqn_weights) { ET_CHECK_OK_OR_RETURN_ERROR(load_constants_from_fqn_manifest( handle, named_data_map, fqn_weight_manifest)); - } - // Load constants. When weight_sharing_across_methods is enabled (opt-in - // via the kWeightSharingAcrossMethods runtime backend option set by the - // runner), use the per-weight FQN cache so methods that share weights - // (e.g. prefill/decode) avoid duplicate GPU allocations. Otherwise fall - // back to the legacy per-method blob load — required for models whose - // methods are independent sub-graphs that may have FQN collisions - // (e.g. parakeet). - else if (is_weight_sharing_across_methods_enabled()) { + } else if (is_weight_sharing_across_methods_enabled()) { ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache( handle, named_data_map, method_name, weights_blob_key)); } else { @@ -1119,6 +1113,7 @@ class ET_EXPERIMENTAL CudaBackend final const NamedDataMap* named_data_map, const CudaFqnWeightEntry& entry, const std::vector* mutable_group, + const std::unordered_map& storage_scope_by_key, int device_index, std::shared_ptr& storage, bool& reused) const { @@ -1130,33 +1125,13 @@ class ET_EXPERIMENTAL CudaBackend final uintptr_t mutable_scope = 0; if (!entry.shareable) { - // Method::init wraps the same external PTD map in a distinct - // MergedDataMap for every method, so the wrapper address is not a model - // instance identity. get_key(), however, forwards the pointer owned by - // the underlying PTD map. That pointer is stable across methods, unique - // to a live PTD instance, and valid for the map's lifetime. - auto num_keys = named_data_map->get_num_keys(); - ET_CHECK_OR_RETURN_ERROR( - num_keys.ok(), - InvalidProgram, - "Failed to enumerate CUDA named data while loading mutable FQN storage"); - for (uint32_t index = 0; index < num_keys.get(); ++index) { - auto key = named_data_map->get_key(index); - ET_CHECK_OR_RETURN_ERROR( - key.ok(), - InvalidProgram, - "Failed to read CUDA named data key %u", - index); - if (entry.storage_key == key.get()) { - mutable_scope = reinterpret_cast(key.get()); - break; - } - } + const auto scope = storage_scope_by_key.find(entry.storage_key); ET_CHECK_OR_RETURN_ERROR( - mutable_scope != 0, + scope != storage_scope_by_key.end(), NotFound, "CUDA mutable FQN storage '%s' is missing from named data", entry.storage_key.c_str()); + mutable_scope = scope->second; } const auto cache_key = [&](const CudaFqnWeightEntry& item) { @@ -1288,6 +1263,10 @@ class ET_EXPERIMENTAL CudaBackend final cuda::CudaDelegateHandle* handle, const NamedDataMap* named_data_map, const CudaFqnWeightManifest& manifest) const { + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "CUDA FQN weights require a named data map"); ET_CHECK_OR_RETURN_ERROR( handle->get_num_constants && handle->get_constant_name && handle->get_constant_original_fqn && handle->get_constant_dtype && @@ -1357,6 +1336,28 @@ class ET_EXPERIMENTAL CudaBackend final mutable_storage_groups[entry.storage_group].push_back(&entry); } } + std::unordered_map storage_scope_by_key; + if (!mutable_storage_groups.empty()) { + // Method::init creates a distinct MergedDataMap wrapper per method, but + // get_key() forwards the stable pointer owned by the underlying PTD map. + // Use that pointer to scope mutable state to one live model instance. + auto num_keys = named_data_map->get_num_keys(); + ET_CHECK_OR_RETURN_ERROR( + num_keys.ok(), + InvalidProgram, + "Failed to enumerate CUDA named data while loading mutable FQN storage"); + storage_scope_by_key.reserve(num_keys.get()); + for (uint32_t index = 0; index < num_keys.get(); ++index) { + auto key = named_data_map->get_key(index); + ET_CHECK_OR_RETURN_ERROR( + key.ok() && key.get() != nullptr, + InvalidProgram, + "Failed to read CUDA named data key %u", + index); + storage_scope_by_key.emplace( + key.get(), reinterpret_cast(key.get())); + } + } std::vector pairs; pairs.reserve(manifest.entries.size()); std::unordered_set bound_fqns; @@ -1407,6 +1408,7 @@ class ET_EXPERIMENTAL CudaBackend final named_data_map, entry, mutable_group, + storage_scope_by_key, device_index, storage, reused), diff --git a/runtime/core/named_data_map.h b/runtime/core/named_data_map.h index dbd5b21a66f..e6a53b7d46e 100644 --- a/runtime/core/named_data_map.h +++ b/runtime/core/named_data_map.h @@ -65,8 +65,9 @@ class NamedDataMap { * Get the key at the given index. * * @param index The index of the key to retrieve. - * @return Result containing the key at the given index. Note: the returned - * pointer is only valid for the lifetime of the DataMap. + * @return Result containing the key at the given index. The returned pointer + * is owned by the DataMap, remains stable across calls, and is only valid for + * the lifetime of the DataMap. */ ET_NODISCARD virtual Result get_key(uint32_t index) const = 0; }; From 02304be492e65b195e75169c19bbe22eed8ea326 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 21 Aug 2026 16:16:34 -0700 Subject: [PATCH 03/12] cuda: keep leaked AOTI libraries loaded Generated with Codex. --- backends/cuda/runtime/cuda_backend.cpp | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 7c2fdbc9167..8c68357b5c0 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -879,20 +879,14 @@ class ET_EXPERIMENTAL CudaBackend final // NOTE: AOTInductorModelContainerDelete does not work correctly with // multiple .so files. Deleting one container frees shared resources, // which causes segmentation faults when attempting to delete other - // containers. As a workaround, we skip explicit container deletion - // and defer cleanup to the OS. + // containers. As a workaround, we skip explicit container deletion and + // defer cleanup to the OS. The corresponding shared library must remain + // loaded as well: the leaked container still owns objects whose code and + // process-wide state live in that library, so dlclose/FreeLibrary can + // invalidate them and crash during multi-method teardown. // TODO(gasoonjia): Find a proper solution for safe container deletion. // AOTInductorModelContainerDelete(handle->container_handle); - // Now close the shared library - if (handle->so_handle != nullptr) { - Error err = close_library(handle->so_handle); - ET_CHECK_OR_LOG_ERROR( - err == Error::Ok, - "Failed to close shared library for %s", - handle->so_path.c_str()); - } - // Remove the temporary shared library file if (!handle->so_path.empty()) { std::error_code remove_error; From fb5f24b3a07408ed75909a148660543b1a9192dd Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 24 Aug 2026 16:11:10 -0700 Subject: [PATCH 04/12] cuda: preserve FQN storage semantics Generated with Codex. --- backends/cuda/cuda_backend.py | 59 ++++++-- backends/cuda/runtime/cuda_backend.cpp | 142 +++++++++++------- backends/cuda/runtime/cuda_delegate_handle.h | 25 +-- backends/cuda/runtime/cuda_weight_manifest.h | 10 +- .../test/test_cuda_weight_manifest.cpp | 15 +- backends/cuda/tests/test_cuda_partitioner.py | 28 ++++ 6 files changed, 202 insertions(+), 77 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 76d97726d6d..a5042cb077e 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -66,9 +66,12 @@ _CPU_CLONE_GUARD = threading.local() -_FQN_WEIGHTS_MAGIC = b"ETCUDAFQN1" +_FQN_WEIGHTS_MAGIC = b"ETCUDAFQN2" _FQN_WEIGHTS_CAPTURE = threading.local() +_AOTI_DEVICE_TYPE_CPU = 0 +_AOTI_DEVICE_TYPE_CUDA = 1 + @dataclass class _FqnWeightEntry: @@ -77,6 +80,7 @@ class _FqnWeightEntry: storage_group: int storage_nbytes: int dtype: int + device_type: int storage_offset: int sizes: Tuple[int, ...] strides: Tuple[int, ...] @@ -107,6 +111,29 @@ def _trim_host_memory() -> None: pass +def _aoti_device_type_for_weight(tensor: torch.Tensor) -> int: + # Low-memory compilation clones lifted CUDA buffers onto CPU, while the + # patched wrapper records them as CUDA constants. Mirror that target-device + # substitution in the manifest. Outside that scoped mode the serialized + # tensor's actual device is the AOTI constant's device. + if _is_cpu_clone_active() or tensor.device.type == "cuda": + return _AOTI_DEVICE_TYPE_CUDA + if tensor.device.type == "cpu": + return _AOTI_DEVICE_TYPE_CPU + raise RuntimeError( + f"Unsupported AOTI constant device for CUDA export: {tensor.device}" + ) + + +def _stateful_buffer_fqns(graph_signature: Any) -> set[str]: + # Some generated-code mutations (notably conditional cache updates) are + # not surfaced through buffers_to_mutate. Treat every registered buffer as + # model-instance-local, then include any explicitly reported mutations. + fqns = set(getattr(graph_signature, "buffers", ())) + fqns.update(getattr(graph_signature, "buffers_to_mutate", {}).values()) + return fqns + + @contextlib.contextmanager def _keep_triton_reduction_loads_loop_scoped(): """Conservatively treating loads as reduction-masked keeps each definition in its @@ -482,13 +509,14 @@ def _materialize_fqn_weights( # noqa: C901 _trim_host_memory() entries: List[_FqnWeightEntry] = [] storages: Dict[str, FileBackedData] = {} - records: List[Tuple[str, torch.Tensor, Any, Tuple[Any, ...], int]] = [] + records: List[Tuple[str, torch.Tensor, Any, Tuple[Any, ...], int, int]] = [] record_indices_by_identity: Dict[Tuple[Any, ...], List[int]] = {} for index, (fqn, (tensor, properties)) in enumerate(weights.items()): storage = tensor.untyped_storage() storage_nbytes = storage.nbytes() storage_ptr = storage.data_ptr() + device_type = _aoti_device_type_for_weight(tensor) property_storage_ptr = getattr(properties, "storage_ptr", None) if property_storage_ptr not in (None, 0): # TensorProperties describes the graph constant's real storage. @@ -498,6 +526,7 @@ def _materialize_fqn_weights( # noqa: C901 "aoti", int(property_storage_ptr), str(tensor.dtype), + device_type, ) else: identity = ( @@ -505,9 +534,10 @@ def _materialize_fqn_weights( # noqa: C901 tensor.device.index if tensor.device.index is not None else -1, storage_ptr if storage_ptr != 0 else -(index + 1), storage_nbytes, + device_type, ) del storage - records.append((fqn, tensor, properties, identity, storage_nbytes)) + records.append((fqn, tensor, properties, identity, storage_nbytes, device_type)) record_indices_by_identity.setdefault(identity, []).append(index) storage_info_by_identity: Dict[Tuple[Any, ...], Tuple[str, int, int]] = {} @@ -561,7 +591,7 @@ def _materialize_fqn_weights( # noqa: C901 storage_group, ) - for fqn, tensor, properties, identity, _storage_nbytes in records: + for fqn, tensor, properties, identity, _storage_nbytes, device_type in records: storage_key, serialized_nbytes, storage_group = storage_info_by_identity[ identity ] @@ -596,6 +626,7 @@ def _materialize_fqn_weights( # noqa: C901 storage_group=storage_group, storage_nbytes=serialized_nbytes, dtype=int(scalar_type_enum(tensor.dtype)), + device_type=device_type, storage_offset=storage_offset, sizes=sizes, strides=strides, @@ -635,10 +666,11 @@ def write_string(value: str) -> None: write_string(entry.storage_key) output.extend( struct.pack( - " PreprocessResult: """Compile CUDA weights as independently addressable AOTI storages.""" - mutated_fqns = set( - getattr(edge_program.graph_signature, "buffers_to_mutate", {}).values() - ) + # Keep every buffer model-instance-local while still sharing it by FQN + # across methods in that model instance. + mutated_fqns = _stateful_buffer_fqns(edge_program.graph_signature) previous_capture = getattr(_FQN_WEIGHTS_CAPTURE, "current", None) capture = _FqnWeightCapture(mutated_fqns=mutated_fqns) # AotiBackend packages weights synchronously on this thread. TLS keeps @@ -809,8 +841,9 @@ def preprocess( # noqa: C901 # legacy runtime path remains able to consume old dense blobs. parent_store = result.data_store_output named_data_store = NamedDataStore() + keep_compatibility_blob = not artifact.storages for key, entry in parent_store.pte_data.items(): - if key != compatibility_blob_key: + if key != compatibility_blob_key or keep_compatibility_blob: named_data_store.add_named_data( key, parent_store.buffers[entry.buffer_index], @@ -819,7 +852,7 @@ def preprocess( # noqa: C901 ) for tag, entries in parent_store.external_data.items(): for key, entry in entries.items(): - if key != compatibility_blob_key: + if key != compatibility_blob_key or keep_compatibility_blob: named_data_store.add_named_data( key, parent_store.buffers[entry.buffer_index], @@ -1110,9 +1143,9 @@ def get_aoti_compile_options( else: # Linux platform - assert ( - shim_library_path is None - ), "shim_library_path should not be set for Linux" + assert shim_library_path is None, ( + "shim_library_path should not be set for Linux" + ) return options @classmethod diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 8c68357b5c0..c46431c1a0b 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -231,8 +231,6 @@ class ET_EXPERIMENTAL CudaBackend final AOTInductorModelContainerGetConstantOriginalFQN); LOAD_OPTIONAL_SYMBOL( get_constant_dtype, AOTInductorModelContainerGetConstantDtype); - LOAD_OPTIONAL_SYMBOL( - get_constant_data_size, AOTInductorModelContainerGetConstantDataSize); LOAD_OPTIONAL_SYMBOL( extract_constants_map, AOTInductorModelContainerExtractConstantsMap); LOAD_OPTIONAL_SYMBOL( @@ -1117,6 +1115,11 @@ class ET_EXPERIMENTAL CudaBackend final InvalidArgument, "CUDA FQN weights require a named data map"); + const auto device_type = + static_cast(entry.device_type); + const bool is_cuda_storage = device_type == slim::c10::DeviceType::CUDA; + const int storage_device_index = is_cuda_storage ? device_index : 0; + uintptr_t mutable_scope = 0; if (!entry.shareable) { const auto scope = storage_scope_by_key.find(entry.storage_key); @@ -1133,14 +1136,16 @@ class ET_EXPERIMENTAL CudaBackend final // Immutable bytes are safe to reuse across methods and model // instances solely by content identity. return std::string("immutable:") + item.storage_key + - "@cuda:" + std::to_string(device_index); + (is_cuda_storage ? "@cuda:" + std::to_string(device_index) + : "@cpu"); } // Stateful buffers with identical initial bytes are not interchangeable. // Scope their logical FQN identity to the underlying PTD instance so // methods in one model share state without leaking it to another live // model instance. return std::string("mutable:") + std::to_string(mutable_scope) + ":" + - item.fqn + "@cuda:" + std::to_string(device_index); + item.fqn + + (is_cuda_storage ? "@cuda:" + std::to_string(device_index) : "@cpu"); }; std::vector cache_keys; @@ -1156,7 +1161,8 @@ class ET_EXPERIMENTAL CudaBackend final for (const CudaFqnWeightEntry* alias : *mutable_group) { ET_CHECK_OR_RETURN_ERROR( alias != nullptr && !alias->shareable && - alias->storage_nbytes == entry.storage_nbytes, + alias->storage_nbytes == entry.storage_nbytes && + alias->device_type == entry.device_type, InvalidProgram, "CUDA mutable FQN storage group %u is inconsistent", entry.storage_group); @@ -1171,12 +1177,12 @@ class ET_EXPERIMENTAL CudaBackend final std::shared_ptr candidate = cached->second.lock(); if (candidate != nullptr) { ET_CHECK_OR_RETURN_ERROR( - candidate->nbytes == entry.storage_nbytes, + candidate->nbytes == entry.storage_nbytes && + candidate->device_type == device_type && + candidate->device_index == storage_device_index, InvalidProgram, - "CUDA FQN storage '%s' has inconsistent sizes (%zu vs %llu)", - entry.storage_key.c_str(), - candidate->nbytes, - static_cast(entry.storage_nbytes)); + "CUDA FQN storage '%s' has inconsistent allocation metadata", + entry.storage_key.c_str()); ET_CHECK_OR_RETURN_ERROR( storage == nullptr || storage.get() == candidate.get(), InvalidProgram, @@ -1194,24 +1200,43 @@ class ET_EXPERIMENTAL CudaBackend final return Error::Ok; } - void* device_data = nullptr; + void* storage_data = nullptr; const size_t allocation_size = std::max(1, static_cast(entry.storage_nbytes)); - const cudaError_t allocation_error = - cudaMalloc(&device_data, allocation_size); - if (allocation_error != cudaSuccess) { - ET_LOG( - Error, - "cudaMalloc failed for FQN storage '%s': %s", - entry.storage_key.c_str(), - cudaGetErrorString(allocation_error)); - return Error::MemoryAllocationFailed; + if (is_cuda_storage) { + const cudaError_t allocation_error = + cudaMalloc(&storage_data, allocation_size); + if (allocation_error != cudaSuccess) { + ET_LOG( + Error, + "cudaMalloc failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(allocation_error)); + return Error::MemoryAllocationFailed; + } + } else { + storage_data = std::malloc(allocation_size); + if (storage_data == nullptr) { + ET_LOG( + Error, + "malloc failed for CPU FQN storage '%s'", + entry.storage_key.c_str()); + return Error::MemoryAllocationFailed; + } } + const auto free_storage_data = [&]() { + if (is_cuda_storage) { + (void)cudaFree(storage_data); + } else { + std::free(storage_data); + } + }; + if (entry.storage_nbytes > 0) { auto host_data = named_data_map->get_data(entry.storage_key.c_str()); if (!host_data.ok()) { - cudaFree(device_data); + free_storage_data(); ET_LOG( Error, "CUDA FQN storage '%s' is missing", @@ -1219,7 +1244,7 @@ class ET_EXPERIMENTAL CudaBackend final return Error::NotFound; } if (host_data->size() != entry.storage_nbytes) { - cudaFree(device_data); + free_storage_data(); ET_LOG( Error, "CUDA FQN storage '%s' has size %zu, expected %llu", @@ -1228,14 +1253,22 @@ class ET_EXPERIMENTAL CudaBackend final static_cast(entry.storage_nbytes)); return Error::InvalidProgram; } - const cudaError_t copy_error = cudaMemcpy( - device_data, - host_data->data(), - static_cast(entry.storage_nbytes), - cudaMemcpyHostToDevice); + cudaError_t copy_error = cudaSuccess; + if (is_cuda_storage) { + copy_error = cudaMemcpy( + storage_data, + host_data->data(), + static_cast(entry.storage_nbytes), + cudaMemcpyHostToDevice); + } else { + std::memcpy( + storage_data, + host_data->data(), + static_cast(entry.storage_nbytes)); + } host_data->Free(); if (copy_error != cudaSuccess) { - cudaFree(device_data); + free_storage_data(); ET_LOG( Error, "cudaMemcpy failed for FQN storage '%s': %s", @@ -1246,7 +1279,10 @@ class ET_EXPERIMENTAL CudaBackend final } storage = std::make_shared( - device_data, static_cast(entry.storage_nbytes), device_index); + storage_data, + static_cast(entry.storage_nbytes), + device_type, + storage_device_index); for (const std::string& key : cache_keys) { shared_fqn_weight_storages_[key] = storage; } @@ -1264,7 +1300,6 @@ class ET_EXPERIMENTAL CudaBackend final ET_CHECK_OR_RETURN_ERROR( handle->get_num_constants && handle->get_constant_name && handle->get_constant_original_fqn && handle->get_constant_dtype && - handle->get_constant_data_size && handle->update_user_managed_constant_buffer_pairs, NotSupported, "AOTI library does not expose the APIs required by CUDA FQN weights"); @@ -1275,13 +1310,11 @@ class ET_EXPERIMENTAL CudaBackend final "Failed to enumerate CUDA AOTI constants"); std::unordered_map> fqn_to_internal_names; - std::unordered_map> - fqn_to_aoti_metadata; + std::unordered_map fqn_to_aoti_dtype; for (size_t index = 0; index < num_constants; ++index) { const char* internal_name = nullptr; const char* fqn = nullptr; int32_t dtype = 0; - size_t data_size = 0; ET_CHECK_OK_OR_RETURN_ERROR( handle->get_constant_name( handle->container_handle, index, &internal_name), @@ -1296,19 +1329,11 @@ class ET_EXPERIMENTAL CudaBackend final handle->get_constant_dtype(handle->container_handle, index, &dtype), "Failed to read CUDA AOTI constant dtype at index %zu", index); - ET_CHECK_OK_OR_RETURN_ERROR( - handle->get_constant_data_size( - handle->container_handle, index, &data_size), - "Failed to read CUDA AOTI constant size at index %zu", - index); if (internal_name != nullptr && fqn != nullptr && fqn[0] != '\0') { fqn_to_internal_names[fqn].emplace_back(internal_name); - auto [metadata, inserted] = - fqn_to_aoti_metadata.emplace(fqn, std::make_pair(dtype, data_size)); + auto [metadata, inserted] = fqn_to_aoti_dtype.emplace(fqn, dtype); ET_CHECK_OR_RETURN_ERROR( - inserted || - (metadata->second.first == dtype && - metadata->second.second == data_size), + inserted || metadata->second == dtype, InvalidProgram, "CUDA AOTI constant FQN '%s' has inconsistent metadata", fqn); @@ -1320,6 +1345,7 @@ class ET_EXPERIMENTAL CudaBackend final struct LocalStorage { std::string storage_key; uint64_t storage_nbytes; + int32_t device_type; std::shared_ptr storage; }; std::unordered_map local_storages; @@ -1369,14 +1395,16 @@ class ET_EXPERIMENTAL CudaBackend final InvalidProgram, "CUDA FQN weight '%s' is not present in its AOTI library", entry.fqn.c_str()); - const auto aoti_metadata = fqn_to_aoti_metadata.find(entry.fqn); + const auto aoti_dtype = fqn_to_aoti_dtype.find(entry.fqn); ET_CHECK_OR_RETURN_ERROR( - aoti_metadata != fqn_to_aoti_metadata.end() && - aoti_metadata->second.first == entry.dtype && - aoti_metadata->second.second == entry.storage_nbytes, + aoti_dtype != fqn_to_aoti_dtype.end() && + aoti_dtype->second == entry.dtype, InvalidProgram, - "CUDA FQN weight '%s' metadata does not match its AOTI library", - entry.fqn.c_str()); + "CUDA FQN weight '%s' dtype does not match its AOTI library " + "(manifest=%d, AOTI=%d)", + entry.fqn.c_str(), + entry.dtype, + aoti_dtype == fqn_to_aoti_dtype.end() ? -1 : aoti_dtype->second); ET_CHECK_OR_RETURN_ERROR( bound_fqns.emplace(entry.fqn).second, InvalidProgram, @@ -1411,12 +1439,17 @@ class ET_EXPERIMENTAL CudaBackend final reused_storages += reused ? 1 : 0; local_storages.emplace( local_key, - LocalStorage{entry.storage_key, entry.storage_nbytes, storage}); + LocalStorage{ + entry.storage_key, + entry.storage_nbytes, + entry.device_type, + storage}); handle->fqn_weight_storages.push_back(storage); } else { ET_CHECK_OR_RETURN_ERROR( local_storage->second.storage_key == entry.storage_key && - local_storage->second.storage_nbytes == entry.storage_nbytes, + local_storage->second.storage_nbytes == entry.storage_nbytes && + local_storage->second.device_type == entry.device_type, InvalidProgram, "CUDA FQN storage group %u has inconsistent backing storage", entry.storage_group); @@ -1428,7 +1461,12 @@ class ET_EXPERIMENTAL CudaBackend final slim::makeArrayRef(entry.sizes), slim::makeArrayRef(entry.strides), static_cast(entry.dtype), - Device(slim::c10::DeviceType::CUDA, device_index), + Device( + static_cast(entry.device_type), + entry.device_type == + static_cast(slim::c10::DeviceType::CUDA) + ? device_index + : 0), entry.storage_offset)); AtenTensorHandle tensor_handle = reinterpret_cast(tensor.get()); diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index ffbfafa097d..585e2eca7db 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -23,24 +24,30 @@ using AOTInductorModelContainerGetConstantDtypeFunc = aoti::AOTInductorModelContainerHandle container_handle, size_t idx, int32_t* dtype); -using AOTInductorModelContainerGetConstantDataSizeFunc = - aoti::AOTIRuntimeError (*)( - aoti::AOTInductorModelContainerHandle container_handle, - size_t idx, - size_t* data_size); - struct CudaWeightStorage { void* data{nullptr}; size_t nbytes{0}; + aoti::slim::c10::DeviceType device_type{aoti::slim::c10::DeviceType::CUDA}; int device_index{0}; - CudaWeightStorage(void* data_, size_t nbytes_, int device_index_) - : data(data_), nbytes(nbytes_), device_index(device_index_) {} + CudaWeightStorage( + void* data_, + size_t nbytes_, + aoti::slim::c10::DeviceType device_type_, + int device_index_) + : data(data_), + nbytes(nbytes_), + device_type(device_type_), + device_index(device_index_) {} ~CudaWeightStorage() { if (data == nullptr) { return; } + if (device_type == aoti::slim::c10::DeviceType::CPU) { + std::free(data); + return; + } int previous_device = 0; const cudaError_t get_device_error = cudaGetDevice(&previous_device); if (get_device_error == cudaSuccess && previous_device != device_index) { @@ -189,8 +196,6 @@ struct CudaGraphState { struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { // Extra AOTI metadata used to validate per-FQN manifests before binding. AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr}; - AOTInductorModelContainerGetConstantDataSizeFunc get_constant_data_size{ - nullptr}; // CUDA stream for this handle, support both shared mode and single mode. // In shared mode, all cuda delegate handles share the same stream (e.g., for diff --git a/backends/cuda/runtime/cuda_weight_manifest.h b/backends/cuda/runtime/cuda_weight_manifest.h index 325a8b25fa7..a451f61a667 100644 --- a/backends/cuda/runtime/cuda_weight_manifest.h +++ b/backends/cuda/runtime/cuda_weight_manifest.h @@ -17,7 +17,7 @@ namespace executorch::backends::cuda { -constexpr char kCudaFqnWeightsMagic[] = "ETCUDAFQN1"; +constexpr char kCudaFqnWeightsMagic[] = "ETCUDAFQN2"; constexpr size_t kCudaFqnWeightsMagicSize = sizeof(kCudaFqnWeightsMagic) - 1; struct CudaFqnWeightEntry { @@ -26,6 +26,7 @@ struct CudaFqnWeightEntry { uint32_t storage_group{0}; uint64_t storage_nbytes{0}; int32_t dtype{0}; + int32_t device_type{0}; int64_t storage_offset{0}; std::vector sizes; std::vector strides; @@ -55,6 +56,11 @@ inline bool is_supported_cuda_fqn_dtype(int32_t dtype) { } } +inline bool is_supported_cuda_fqn_device_type(int32_t device_type) { + // Values match c10::DeviceType and the slim AOTI runtime. + return device_type == 0 || device_type == 1; // CPU or CUDA +} + inline bool is_cuda_fqn_weight_manifest(const void* data, size_t size) { return data != nullptr && size >= kCudaFqnWeightsMagicSize && std::memcmp(data, kCudaFqnWeightsMagic, kCudaFqnWeightsMagicSize) == 0; @@ -186,6 +192,8 @@ inline executorch::runtime::Error parse_cuda_fqn_weight_manifest( !reader.read_u64(entry.storage_nbytes) || !reader.read_i32(entry.dtype) || !is_supported_cuda_fqn_dtype(entry.dtype) || + !reader.read_i32(entry.device_type) || + !is_supported_cuda_fqn_device_type(entry.device_type) || !reader.read_i64(entry.storage_offset) || !reader.read_u32(ndim) || ndim > kMaxTensorDimensions) { return Error::InvalidProgram; diff --git a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp b/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp index db656482ffc..8fa643fa555 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp @@ -36,7 +36,9 @@ void append_string(std::vector& output, const std::string& value) { output.insert(output.end(), value.begin(), value.end()); } -std::vector valid_manifest(uint32_t dtype = 6) { +std::vector valid_manifest( + uint32_t dtype = 6, + uint32_t device_type = 1) { std::vector output( cuda::kCudaFqnWeightsMagic, cuda::kCudaFqnWeightsMagic + cuda::kCudaFqnWeightsMagicSize); @@ -47,6 +49,7 @@ std::vector valid_manifest(uint32_t dtype = 6) { append_u32(output, 7); // method-local storage group append_u64(output, 24); // storage bytes append_u32(output, dtype); // dtype + append_u32(output, device_type); // device type (CUDA) append_u64(output, 0); // storage offset append_u32(output, 2); // ndim append_u64(output, 2); @@ -79,6 +82,7 @@ TEST(CudaWeightManifestTest, ParsesVersionedManifest) { EXPECT_EQ(entry.storage_group, 7u); EXPECT_EQ(entry.storage_nbytes, 24u); EXPECT_EQ(entry.dtype, 6); + EXPECT_EQ(entry.device_type, 1); EXPECT_EQ(entry.sizes, (std::vector{2, 3})); EXPECT_EQ(entry.strides, (std::vector{3, 1})); EXPECT_TRUE(entry.shareable); @@ -108,3 +112,12 @@ TEST(CudaWeightManifestTest, RejectsUnsupportedDtype) { bytes.data(), bytes.size(), manifest), Error::InvalidProgram); } + +TEST(CudaWeightManifestTest, RejectsUnsupportedDeviceType) { + const std::vector bytes = valid_manifest(6, 2); + cuda::CudaFqnWeightManifest manifest; + EXPECT_EQ( + cuda::parse_cuda_fqn_weight_manifest( + bytes.data(), bytes.size(), manifest), + Error::InvalidProgram); +} diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 33305c85438..b3c4f69cc39 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -9,6 +9,7 @@ import os import tempfile import unittest +from types import SimpleNamespace from typing import Tuple from unittest.mock import patch @@ -17,6 +18,7 @@ _encode_fqn_weight_manifest, _FQN_WEIGHTS_MAGIC, _materialize_fqn_weights, + _stateful_buffer_fqns, CudaBackend, ) from executorch.backends.cuda.cuda_partitioner import CudaPartitioner @@ -31,6 +33,16 @@ class TestCudaLowMemoryExport(unittest.TestCase): + def test_all_buffers_are_model_instance_local(self) -> None: + signature = SimpleNamespace( + buffers=("persistent", "conditional_cache"), + buffers_to_mutate={"copy_out": "explicitly_mutated"}, + ) + self.assertEqual( + _stateful_buffer_fqns(signature), + {"persistent", "conditional_cache", "explicitly_mutated"}, + ) + @patch.object(CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True) def test_all_cuda_exports_request_structured_weights(self, _) -> None: options = CudaBackend.get_aoti_compile_options( @@ -65,6 +77,8 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: self.assertEqual(2, len(artifact.storages)) self.assertTrue(artifact.entries[0].shareable) self.assertFalse(artifact.entries[1].shareable) + self.assertEqual(artifact.entries[0].device_type, 0) + self.assertEqual(artifact.entries[1].device_type, 0) self.assertEqual( bytes(first.untyped_storage()), artifact.storages[artifact.entries[0].storage_key].to_bytes(), @@ -81,6 +95,20 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: for storage in artifact.storages.values(): storage.close() + @patch( + "executorch.backends.cuda.cuda_backend._is_cpu_clone_active", + return_value=True, + ) + def test_low_memory_cpu_clones_keep_cuda_device_type(self, _) -> None: + tensor = torch.tensor([1, 2], dtype=torch.int16) + weights = Weights({"weight": (tensor, TensorProperties(tensor))}) + + with tempfile.TemporaryDirectory() as directory: + artifact = _materialize_fqn_weights(weights, directory, set()) + self.assertEqual(artifact.entries[0].device_type, 1) + for storage in artifact.storages.values(): + storage.close() + def test_views_share_one_physical_storage(self) -> None: base = torch.arange(12, dtype=torch.float32).reshape(3, 4) view = base[:, 1:] From 89de71f1b24e577360cafaa4c17d7495bf378a33 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Mon, 24 Aug 2026 16:21:16 -0700 Subject: [PATCH 05/12] cuda: apply Python formatting Generated with Codex. --- backends/cuda/cuda_backend.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index a5042cb077e..5da832aaa10 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1143,9 +1143,9 @@ def get_aoti_compile_options( else: # Linux platform - assert shim_library_path is None, ( - "shim_library_path should not be set for Linux" - ) + assert ( + shim_library_path is None + ), "shim_library_path should not be set for Linux" return options @classmethod From 38b0c993b89b33511c049120c935a3be0f60b1f8 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 11:17:07 -0700 Subject: [PATCH 06/12] cuda: centralize multimethod FQN weights Move structured CUDA weight collection into a dedicated export-side collector, key values by device and FQN, and reject conflicting values or metadata. Keep the runtime backend integration small by delegating manifest loading and cross-method allocation reuse to a focused CUDA weight cache.\n\nGenerated with Codex. --- backends/cuda/BUCK | 1 + backends/cuda/CMakeLists.txt | 1 + backends/cuda/cuda_backend.py | 424 +++------------ backends/cuda/cuda_weight_collector.py | 321 +++++++++++ backends/cuda/runtime/cuda_backend.cpp | 497 +----------------- backends/cuda/runtime/cuda_delegate_handle.h | 3 +- backends/cuda/runtime/cuda_weight_cache.cpp | 398 ++++++++++++++ backends/cuda/runtime/cuda_weight_cache.h | 47 ++ backends/cuda/runtime/cuda_weight_manifest.h | 18 +- backends/cuda/runtime/targets.bzl | 2 + .../test/test_cuda_weight_manifest.cpp | 4 - backends/cuda/tests/test_cuda_partitioner.py | 185 +++++-- 12 files changed, 990 insertions(+), 911 deletions(-) create mode 100644 backends/cuda/cuda_weight_collector.py create mode 100644 backends/cuda/runtime/cuda_weight_cache.cpp create mode 100644 backends/cuda/runtime/cuda_weight_cache.h diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 1fbbf61be3c..2d0da4af2af 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -93,6 +93,7 @@ fbcode_target( name = "cuda_backend", srcs = [ "cuda_backend.py", + "cuda_weight_collector.py", ], visibility = ["PUBLIC"], deps = [ diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 0cae1bc77f3..7bb51e4b0eb 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -339,6 +339,7 @@ install( # CUDA backend implementation set(_aoti_cuda_backend_sources runtime/cuda_backend.cpp runtime/cuda_mutable_state.cpp + runtime/cuda_weight_cache.cpp ) if(_cuda_is_msvc_toolchain) # MSVC links aoti_cuda_backend into portable_lib without relying on C++ diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 5da832aaa10..7b2157fc4e0 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -7,23 +7,23 @@ import contextlib import copy -import ctypes import functools -import gc -import hashlib import logging import os import shutil -import struct -import tempfile import threading import typing -from dataclasses import dataclass from importlib import resources -from typing import Any, Dict, final, List, Optional, Tuple +from typing import Any, Dict, final, List, Optional import torch from executorch.backends.aoti.aoti_backend import AotiBackend +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CPU, + AOTI_DEVICE_TYPE_CUDA, + CudaWeightCollector, + trim_host_memory, +) from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) @@ -34,11 +34,9 @@ ReplaceEdgeOpWithTritonOpPass, ) from executorch.exir._serialize._cord import FileBackedData -from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir._warnings import experimental from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec -from executorch.exir.tensor import scalar_type_enum from torch._inductor.decomposition import conv1d_to_conv2d from torch.nn.attention import SDPBackend @@ -66,74 +64,25 @@ _CPU_CLONE_GUARD = threading.local() -_FQN_WEIGHTS_MAGIC = b"ETCUDAFQN2" -_FQN_WEIGHTS_CAPTURE = threading.local() - -_AOTI_DEVICE_TYPE_CPU = 0 -_AOTI_DEVICE_TYPE_CUDA = 1 - - -@dataclass -class _FqnWeightEntry: - fqn: str - storage_key: str - storage_group: int - storage_nbytes: int - dtype: int - device_type: int - storage_offset: int - sizes: Tuple[int, ...] - strides: Tuple[int, ...] - shareable: bool - - -@dataclass -class _FqnWeightArtifact: - entries: List[_FqnWeightEntry] - storages: Dict[str, FileBackedData] - - -@dataclass -class _FqnWeightCapture: - mutated_fqns: set[str] - artifact: Optional[_FqnWeightArtifact] = None - def _is_cpu_clone_active() -> bool: return getattr(_CPU_CLONE_GUARD, "active", False) -def _trim_host_memory() -> None: - gc.collect() - try: - ctypes.CDLL(None).malloc_trim(0) - except AttributeError: - pass - - def _aoti_device_type_for_weight(tensor: torch.Tensor) -> int: # Low-memory compilation clones lifted CUDA buffers onto CPU, while the # patched wrapper records them as CUDA constants. Mirror that target-device # substitution in the manifest. Outside that scoped mode the serialized # tensor's actual device is the AOTI constant's device. if _is_cpu_clone_active() or tensor.device.type == "cuda": - return _AOTI_DEVICE_TYPE_CUDA + return AOTI_DEVICE_TYPE_CUDA if tensor.device.type == "cpu": - return _AOTI_DEVICE_TYPE_CPU + return AOTI_DEVICE_TYPE_CPU raise RuntimeError( f"Unsupported AOTI constant device for CUDA export: {tensor.device}" ) -def _stateful_buffer_fqns(graph_signature: Any) -> set[str]: - # Some generated-code mutations (notably conditional cache updates) are - # not surfaced through buffers_to_mutate. Treat every registered buffer as - # model-instance-local, then include any explicitly reported mutations. - fqns = set(getattr(graph_signature, "buffers", ())) - fqns.update(getattr(graph_signature, "buffers_to_mutate", {}).values()) - return fqns - - @contextlib.contextmanager def _keep_triton_reduction_loads_loop_scoped(): """Conservatively treating loads as reduction-masked keeps each definition in its @@ -463,224 +412,6 @@ def _on_off_compile_spec_value(spec: CompileSpec) -> bool: return value == "ON" -def _write_tensor_storage(tensor: torch.Tensor, path: str) -> bytes: - """Stream one AOTI storage to ``path`` and return its SHA-256 digest.""" - chunk_size = 8 * 1024 * 1024 - digest = hashlib.sha256() - - def write_chunk(output, chunk) -> None: - digest.update(chunk) - output.write(chunk) - - if tensor.is_mkldnn: - raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI") - storage = tensor.untyped_storage() - nbytes = storage.nbytes() - with open(path, "wb") as output: - if nbytes and tensor.is_cuda: - byte_tensor = torch.empty(0, dtype=torch.uint8, device=tensor.device).set_( - storage, 0, (nbytes,), (1,) - ) - for offset in range(0, nbytes, chunk_size): - cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu() - 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): - write_chunk(output, raw_view[offset : offset + chunk_size]) - del raw_view, raw_array - del storage - return digest.digest() - - -def _materialize_fqn_weights( # noqa: C901 - weights: Any, - directory: str, - mutated_fqns: set[str], -) -> _FqnWeightArtifact: - """Turn AOTI ``Weights`` into content-addressed storage files + views.""" - # The graph can contain hundreds of independent storages. Trimming around - # every storage is both ineffective (``records`` below still owns all of - # the tensors) and extremely expensive for large exported graphs. Trim - # once at artifact boundaries; streamed CPU chunks are released by normal - # reference counting as they are replaced. - _trim_host_memory() - entries: List[_FqnWeightEntry] = [] - storages: Dict[str, FileBackedData] = {} - records: List[Tuple[str, torch.Tensor, Any, Tuple[Any, ...], int, int]] = [] - record_indices_by_identity: Dict[Tuple[Any, ...], List[int]] = {} - - for index, (fqn, (tensor, properties)) in enumerate(weights.items()): - storage = tensor.untyped_storage() - storage_nbytes = storage.nbytes() - storage_ptr = storage.data_ptr() - device_type = _aoti_device_type_for_weight(tensor) - property_storage_ptr = getattr(properties, "storage_ptr", None) - if property_storage_ptr not in (None, 0): - # TensorProperties describes the graph constant's real storage. - # The value tensor can be a clone (including a CPU clone in CUDA - # low-memory mode), so its data_ptr is not a stable alias key. - identity = ( - "aoti", - int(property_storage_ptr), - str(tensor.dtype), - device_type, - ) - else: - identity = ( - tensor.device.type, - tensor.device.index if tensor.device.index is not None else -1, - storage_ptr if storage_ptr != 0 else -(index + 1), - storage_nbytes, - device_type, - ) - del storage - records.append((fqn, tensor, properties, identity, storage_nbytes, device_type)) - record_indices_by_identity.setdefault(identity, []).append(index) - - storage_info_by_identity: Dict[Tuple[Any, ...], Tuple[str, int, int]] = {} - for storage_group, (identity, record_indices) in enumerate( - record_indices_by_identity.items() - ): - # AOTI's value can be a clone of a view. Pick the largest available - # backing storage in the alias group so every declared view can be - # reconstructed from the one serialized allocation. - candidate_index = max(record_indices, key=lambda item: records[item][4]) - candidate_tensor = records[candidate_index][1] - storage_nbytes = records[candidate_index][4] - expected_storage_nbytes = max( - ( - int(storage_size) - for item in record_indices - if (storage_size := getattr(records[item][2], "storage_size", None)) - is not None - ), - default=0, - ) - if storage_nbytes < expected_storage_nbytes: - raise RuntimeError( - "AOTI cloned storage is smaller than its TensorProperties " - f"({storage_nbytes} < {expected_storage_nbytes} bytes)" - ) - - fd, storage_path = tempfile.mkstemp( - prefix=".cuda_weight_", suffix=".storage", dir=directory - ) - os.close(fd) - try: - digest = _write_tensor_storage(candidate_tensor, storage_path) - storage_key = digest.hex() + "_cuda_weight_storage" - data = FileBackedData.move_from(storage_path, sha256=digest) - except Exception: - try: - os.remove(storage_path) - except OSError: - pass - raise - - existing = storages.get(storage_key) - if existing is None: - storages[storage_key] = data - else: - data.close() - storage_info_by_identity[identity] = ( - storage_key, - storage_nbytes, - storage_group, - ) - - for fqn, tensor, properties, identity, _storage_nbytes, device_type in records: - storage_key, serialized_nbytes, storage_group = storage_info_by_identity[ - identity - ] - sizes = getattr(properties, "shape", tensor.shape) - strides = getattr(properties, "stride", tensor.stride()) - storage_offset = getattr(properties, "offset", tensor.storage_offset()) - sizes = tuple(int(size) for size in sizes) - strides = tuple(int(stride) for stride in strides) - storage_offset = int(storage_offset) - if ( - len(sizes) != len(strides) - or storage_offset < 0 - or any(size < 0 for size in sizes) - or any(stride < 0 for stride in strides) - ): - raise RuntimeError(f"AOTI view {fqn!r} has invalid tensor metadata") - required_nbytes = 0 - if all(size != 0 for size in sizes): - last_element = storage_offset + sum( - stride * (size - 1) for size, stride in zip(sizes, strides) - ) - required_nbytes = (last_element + 1) * tensor.element_size() - if required_nbytes > serialized_nbytes: - raise RuntimeError( - f"AOTI view {fqn!r} requires {required_nbytes} bytes from a " - f"{serialized_nbytes}-byte cloned storage" - ) - entries.append( - _FqnWeightEntry( - fqn=fqn, - storage_key=storage_key, - storage_group=storage_group, - storage_nbytes=serialized_nbytes, - dtype=int(scalar_type_enum(tensor.dtype)), - device_type=device_type, - storage_offset=storage_offset, - sizes=sizes, - strides=strides, - shareable=fqn not in mutated_fqns, - ) - ) - - # A mutable view makes its complete physical storage stateful. The runtime - # shares such storages by FQN (not by content hash), including aliases that - # share the same backing buffer. - local_storage_groups = { - entry.storage_group for entry in entries if not entry.shareable - } - for entry in entries: - if entry.storage_group in local_storage_groups: - entry.shareable = False - - _trim_host_memory() - return _FqnWeightArtifact(entries=entries, storages=storages) - - -def _encode_fqn_weight_manifest( - so_blob_key: str, entries: List[_FqnWeightEntry] -) -> bytes: - """Encode the CUDA per-storage manifest consumed by the runtime.""" - output = bytearray(_FQN_WEIGHTS_MAGIC) - - def write_string(value: str) -> None: - encoded = value.encode("utf-8") - output.extend(struct.pack(" str: return "cuda" @@ -802,77 +528,56 @@ def save_data_externally(cls) -> bool: return True @classmethod - def preprocess( # noqa: C901 - cls, edge_program: Any, compile_specs: List[CompileSpec] + def _preprocess_with_weight_collector( + cls, + edge_program: Any, + compile_specs: List[CompileSpec], + collector: CudaWeightCollector, ) -> PreprocessResult: - """Compile CUDA weights as independently addressable AOTI storages.""" - # Keep every buffer model-instance-local while still sharing it by FQN - # across methods in that model instance. - mutated_fqns = _stateful_buffer_fqns(edge_program.graph_signature) - previous_capture = getattr(_FQN_WEIGHTS_CAPTURE, "current", None) - capture = _FqnWeightCapture(mutated_fqns=mutated_fqns) - # AotiBackend packages weights synchronously on this thread. TLS keeps - # nested or concurrent preprocess calls isolated while that callback - # passes the structured artifact back to this invocation. - _FQN_WEIGHTS_CAPTURE.current = capture - try: + with collector.capture(_aoti_device_type_for_weight) as capture: result = super().preprocess(edge_program, compile_specs) - finally: - _FQN_WEIGHTS_CAPTURE.current = previous_capture - - artifact = capture.artifact - if artifact is None: + if capture.artifact is None: raise RuntimeError("CUDA AOTI did not return a structured Weights output") - if result.data_store_output is None: - raise RuntimeError("CUDA AOTI preprocess returned no named data") - - try: - parent_keys = result.processed_bytes.decode("utf-8").splitlines() - except UnicodeDecodeError as error: - raise RuntimeError("Malformed CUDA AOTI named-data payload") from error - if not parent_keys or not parent_keys[0]: - raise RuntimeError("CUDA AOTI payload is missing its shared-object key") - so_blob_key = parent_keys[0] - compatibility_blob_key = parent_keys[1] if len(parent_keys) > 1 else None - - # Rebuild AotiBackend's store without the empty compatibility blob, - # then add each physical weight storage as separately named external - # data. This leaves a new PTD containing only real storages while the - # legacy runtime path remains able to consume old dense blobs. - parent_store = result.data_store_output - named_data_store = NamedDataStore() - keep_compatibility_blob = not artifact.storages - for key, entry in parent_store.pte_data.items(): - if key != compatibility_blob_key or keep_compatibility_blob: - named_data_store.add_named_data( - key, - parent_store.buffers[entry.buffer_index], - alignment=entry.alignment, - tensor_layout=entry.tensor_layout, - ) - for tag, entries in parent_store.external_data.items(): - for key, entry in entries.items(): - if key != compatibility_blob_key or keep_compatibility_blob: - named_data_store.add_named_data( - key, - parent_store.buffers[entry.buffer_index], - alignment=entry.alignment, - external_tag=tag, - tensor_layout=entry.tensor_layout, - ) - - external_tag = f"aoti_{cls.get_device_name()}_blob" - for storage_key, data in artifact.storages.items(): - named_data_store.add_named_data( - storage_key, data, alignment=1, external_tag=external_tag - ) + collector.add_preprocess_result(result, capture.artifact, cls.get_device_name()) + return result - result.processed_bytes = _encode_fqn_weight_manifest( - so_blob_key, artifact.entries + @classmethod + def preprocess( + cls, edge_program: Any, compile_specs: List[CompileSpec] + ) -> PreprocessResult: + """Compile one CUDA method with its own weight collector.""" + collector = CudaWeightCollector() + result = cls._preprocess_with_weight_collector( + edge_program, compile_specs, collector ) - result.data_store_output = named_data_store.get_named_data_store_output() + collector.finish() return result + @classmethod + def preprocess_multimethod( + cls, + edge_programs: Dict[str, List[Any]], + compile_specs: Dict[str, List[List[CompileSpec]]], + ) -> Dict[str, List[PreprocessResult]]: + """Compile every method against one global ``(device, FQN)`` store.""" + collector = CudaWeightCollector() + results: Dict[str, List[PreprocessResult]] = {} + for method_name, programs in edge_programs.items(): + if method_name not in compile_specs: + raise ValueError(f"Missing CUDA compile specs for {method_name!r}") + method_specs = compile_specs[method_name] + if len(programs) != len(method_specs): + raise ValueError( + f"Method {method_name!r} has {len(programs)} partitions but " + f"{len(method_specs)} compile-spec lists" + ) + results[method_name] = [ + cls._preprocess_with_weight_collector(program, specs, collector) + for program, specs in zip(programs, method_specs) + ] + collector.finish() + return results + @classmethod def load_weights_blob( cls, blob_path: str, compile_specs: List[CompileSpec] @@ -883,12 +588,10 @@ def load_weights_blob( compatibility placeholder. Legacy binary-blob handling remains unchanged for callers that still provide a real blob. """ - known_hash = cls._materialized_blob_hashes.pop(blob_path, None) 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, sha256=known_hash) - weights_blob_hash = known_hash or blob_data.sha256() - return blob_data, weights_blob_hash.hex() + blob_data = FileBackedData.move_from(blob_path) + return blob_data, blob_data.sha256().hex() @classmethod def materialize_weights_blob( @@ -918,13 +621,9 @@ def materialize_weights_blob( if so_path is None: raise RuntimeError(f"Expected a CUDA AOTI .wrapper.so output, got {paths}") blob_path = os.path.splitext(so_path)[0] + "_weights.blob" - capture = getattr(_FQN_WEIGHTS_CAPTURE, "current", None) - if capture is None: - raise RuntimeError( - "CUDA structured weights must be materialized inside preprocess" - ) - capture.artifact = _materialize_fqn_weights( - weights[0], os.path.dirname(blob_path), capture.mutated_fqns + capture = CudaWeightCollector.current_capture() + capture.artifact = capture.collector.materialize( + weights[0], os.path.dirname(blob_path), capture.device_type_for_weight ) # Keep AotiBackend's existing path contract intact. The compatibility @@ -932,7 +631,6 @@ def materialize_weights_blob( # artifacts continue to carry and load their original dense blob. with open(blob_path, "wb"): pass - cls._materialized_blob_hashes[blob_path] = hashlib.sha256(b"").digest() # Replace the structured Weights output with the compatibility path # expected by AotiBackend's existing named-data packaging contract. @@ -1076,8 +774,8 @@ def get_aoti_compile_options( # Separate weight constants from the .so file "aot_inductor.package": True, "aot_inductor.package_constants_in_so": False, - # Ask AOTI for structured constants. CUDABackend converts these to - # independently named physical storages plus an FQN view manifest. + # Ask AOTI for structured constants. CUDABackend collects them by + # (device, FQN) and emits the tensor metadata needed at runtime. "aot_inductor.package_constants_on_disk_format": cls._weights_format( compile_specs ), @@ -1197,7 +895,7 @@ def _combined(): stack.enter_context( _compile_time_cpu_clones(torch.device(cls.get_device_name())) ) - _trim_host_memory() + trim_host_memory() yield return _combined() @@ -1213,7 +911,7 @@ def _is_low_memory_mode(compile_specs: List[CompileSpec]) -> bool: @classmethod def _weights_format(cls, compile_specs: List[CompileSpec]) -> str: # CUDA consumes the structured AOTI output directly and emits a - # versioned per-storage manifest. This is backend-wide rather than a + # versioned per-FQN manifest. This is backend-wide rather than a # model/export-script option. return "pickle_weights" diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py new file mode 100644 index 00000000000..bfd02d232a3 --- /dev/null +++ b/backends/cuda/cuda_weight_collector.py @@ -0,0 +1,321 @@ +# 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 contextlib +import ctypes +import gc +import hashlib +import os +import struct +import tempfile +import threading +from dataclasses import dataclass +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple + +import torch +from executorch.exir._serialize._cord import FileBackedData +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 + + +CUDA_FQN_WEIGHTS_MAGIC = b"ETCUDAFQN3" + +AOTI_DEVICE_TYPE_CPU = 0 +AOTI_DEVICE_TYPE_CUDA = 1 + + +@dataclass(frozen=True) +class CudaWeightEntry: + fqn: str + storage_key: str + storage_nbytes: int + dtype: int + device_type: int + storage_offset: int + sizes: Tuple[int, ...] + strides: Tuple[int, ...] + + +@dataclass +class CudaWeightArtifact: + entries: List[CudaWeightEntry] + storages: Dict[str, FileBackedData] + + +@dataclass +class _CudaWeightCapture: + collector: "CudaWeightCollector" + device_type_for_weight: Callable[[torch.Tensor], int] + artifact: Optional[CudaWeightArtifact] = None + + +def trim_host_memory() -> None: + gc.collect() + try: + ctypes.CDLL(None).malloc_trim(0) + except AttributeError: + pass + + +def _write_tensor_storage(tensor: torch.Tensor, path: str) -> bytes: + """Stream one AOTI storage to ``path`` and return its SHA-256 digest.""" + chunk_size = 8 * 1024 * 1024 + digest = hashlib.sha256() + + def write_chunk(output, chunk) -> None: + digest.update(chunk) + output.write(chunk) + + if tensor.is_mkldnn: + raise RuntimeError("MKLDNN constants are not supported by CUDA AOTI") + storage = tensor.untyped_storage() + nbytes = storage.nbytes() + with open(path, "wb") as output: + if nbytes and tensor.is_cuda: + byte_tensor = torch.empty(0, dtype=torch.uint8, device=tensor.device).set_( + storage, 0, (nbytes,), (1,) + ) + for offset in range(0, nbytes, chunk_size): + cpu_chunk = byte_tensor[offset : offset + chunk_size].cpu() + 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): + write_chunk(output, raw_view[offset : offset + chunk_size]) + del raw_view, raw_array + del storage + return digest.digest() + + +def _storage_key(fqn: str, device_type: int) -> str: + if device_type == AOTI_DEVICE_TYPE_CPU: + device = "cpu" + elif device_type == AOTI_DEVICE_TYPE_CUDA: + device = "cuda" + else: + raise RuntimeError(f"Unsupported AOTI device type: {device_type}") + return f"cuda_fqn_weight:{device}:{fqn}" + + +def encode_cuda_weight_manifest( + so_blob_key: str, entries: List[CudaWeightEntry] +) -> bytes: + """Encode the per-method FQN-to-tensor metadata consumed by CUDA runtime.""" + output = bytearray(CUDA_FQN_WEIGHTS_MAGIC) + + def write_string(value: str) -> None: + encoded = value.encode("utf-8") + output.extend(struct.pack(" value`` store for all methods.""" + + _active = threading.local() + + def __init__(self) -> None: + self._store = NamedDataStore() + self._entries: Dict[str, CudaWeightEntry] = {} + self._results: List[PreprocessResult] = [] + + @contextlib.contextmanager + def capture( + self, device_type_for_weight: Callable[[torch.Tensor], int] + ) -> Iterator[_CudaWeightCapture]: + previous = getattr(self._active, "current", None) + capture = _CudaWeightCapture(self, device_type_for_weight) + self._active.current = capture + try: + yield capture + finally: + self._active.current = previous + + @classmethod + def current_capture(cls) -> _CudaWeightCapture: + capture = getattr(cls._active, "current", None) + if capture is None: + raise RuntimeError( + "CUDA structured weights must be materialized inside preprocess" + ) + return capture + + def materialize( + self, + weights: Any, + directory: str, + device_type_for_weight: Callable[[torch.Tensor], int], + ) -> CudaWeightArtifact: + """Turn AOTI ``Weights`` into one independently named blob per FQN.""" + trim_host_memory() + entries: List[CudaWeightEntry] = [] + storages: Dict[str, FileBackedData] = {} + + for fqn, (tensor, properties) in weights.items(): + storage = tensor.untyped_storage() + storage_nbytes = storage.nbytes() + del storage + device_type = device_type_for_weight(tensor) + expected_storage_nbytes = int( + getattr(properties, "storage_size", None) or 0 + ) + if storage_nbytes < expected_storage_nbytes: + raise RuntimeError( + "AOTI cloned storage is smaller than its TensorProperties " + f"({storage_nbytes} < {expected_storage_nbytes} bytes)" + ) + + fd, storage_path = tempfile.mkstemp( + prefix=".cuda_weight_", suffix=".storage", dir=directory + ) + os.close(fd) + try: + digest = _write_tensor_storage(tensor, storage_path) + data = FileBackedData.move_from(storage_path, sha256=digest) + except Exception: + try: + os.remove(storage_path) + except OSError: + pass + raise + + storage_key = _storage_key(fqn, device_type) + if storage_key in storages: + data.close() + raise RuntimeError(f"Duplicate CUDA FQN weight key for {fqn!r}") + storages[storage_key] = data + + sizes = tuple( + int(size) for size in getattr(properties, "shape", tensor.shape) + ) + strides = tuple( + int(stride) for stride in getattr(properties, "stride", tensor.stride()) + ) + storage_offset = int(getattr(properties, "offset", tensor.storage_offset())) + if ( + len(sizes) != len(strides) + or storage_offset < 0 + or any(size < 0 for size in sizes) + or any(stride < 0 for stride in strides) + ): + raise RuntimeError(f"AOTI view {fqn!r} has invalid tensor metadata") + required_nbytes = 0 + if all(size != 0 for size in sizes): + last_element = storage_offset + sum( + stride * (size - 1) for size, stride in zip(sizes, strides) + ) + required_nbytes = (last_element + 1) * tensor.element_size() + if required_nbytes > storage_nbytes: + raise RuntimeError( + f"AOTI view {fqn!r} requires {required_nbytes} bytes from a " + f"{storage_nbytes}-byte cloned storage" + ) + entries.append( + CudaWeightEntry( + fqn=fqn, + storage_key=storage_key, + storage_nbytes=storage_nbytes, + dtype=int(scalar_type_enum(tensor.dtype)), + device_type=device_type, + storage_offset=storage_offset, + sizes=sizes, + strides=strides, + ) + ) + + trim_host_memory() + return CudaWeightArtifact(entries=entries, storages=storages) + + def add_preprocess_result( + self, + result: PreprocessResult, + artifact: CudaWeightArtifact, + device_name: str, + ) -> None: + if result.data_store_output is None: + raise RuntimeError("CUDA AOTI preprocess returned no named data") + try: + parent_keys = result.processed_bytes.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise RuntimeError("Malformed CUDA AOTI named-data payload") from error + if not parent_keys or not parent_keys[0]: + raise RuntimeError("CUDA AOTI payload is missing its shared-object key") + so_blob_key = parent_keys[0] + compatibility_blob_key = parent_keys[1] if len(parent_keys) > 1 else None + + parent_store = result.data_store_output + keep_compatibility_blob = not artifact.storages + for key, entry in parent_store.pte_data.items(): + if key != compatibility_blob_key or keep_compatibility_blob: + self._store.add_named_data( + key, + parent_store.buffers[entry.buffer_index], + alignment=entry.alignment, + tensor_layout=entry.tensor_layout, + ) + for tag, entries in parent_store.external_data.items(): + for key, entry in entries.items(): + if key != compatibility_blob_key or keep_compatibility_blob: + self._store.add_named_data( + key, + parent_store.buffers[entry.buffer_index], + alignment=entry.alignment, + external_tag=tag, + tensor_layout=entry.tensor_layout, + ) + + external_tag = f"aoti_{device_name}_blob" + for entry in artifact.entries: + data = artifact.storages[entry.storage_key] + previous = self._entries.get(entry.storage_key) + try: + if previous is not None and previous != entry: + raise ValueError( + f"Duplicate key {entry.storage_key} with different tensor " + "metadata." + ) + self._store.add_named_data( + entry.storage_key, + data, + alignment=1, + external_tag=external_tag, + ) + finally: + if previous is not None: + data.close() + self._entries.setdefault(entry.storage_key, entry) + + result.processed_bytes = encode_cuda_weight_manifest( + so_blob_key, artifact.entries + ) + self._results.append(result) + + def finish(self) -> None: + shared_output = self._store.get_named_data_store_output() + for result in self._results: + result.data_store_output = shared_output diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index c46431c1a0b..773992db068 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -17,12 +17,10 @@ #include #include -#include #include #include #include #include -#include #include #include #include @@ -48,7 +46,7 @@ #include #include #include -#include +#include #include #include #include @@ -168,8 +166,8 @@ class ET_EXPERIMENTAL CudaBackend final return shared_cuda_stream_ != nullptr; } - // Enable the legacy dense-blob per-FQN cache. New manifest artifacts carry - // enough ownership metadata to share immutable storages automatically. + // Enable the legacy dense-blob per-FQN cache. New manifest artifacts use + // their FQN-addressed data keys automatically. void set_weight_sharing_across_methods(bool enabled) { weight_sharing_across_methods_.store(enabled, std::memory_order_relaxed); } @@ -420,15 +418,12 @@ class ET_EXPERIMENTAL CudaBackend final handle->container_handle = container_handle; - // Versioned FQN artifacts carry complete storage/view and mutability - // metadata, so immutable storages are safely shared without a load-order - // contract. Legacy artifacts keep their historical dense-blob behavior: - // the runtime option selects the per-FQN cache, otherwise each method - // loads its own blob (required for independent methods with FQN - // collisions). + // Versioned artifacts load each (device, FQN) through the same process-wide + // cross-method cache model used by the legacy path. The manifest only adds + // the tensor metadata needed to reconstruct independently named PTD blobs. if (has_fqn_weights) { - ET_CHECK_OK_OR_RETURN_ERROR(load_constants_from_fqn_manifest( - handle, named_data_map, fqn_weight_manifest)); + ET_CHECK_OK_OR_RETURN_ERROR( + fqn_weight_cache_.load(handle, named_data_map, fqn_weight_manifest)); } else if (is_weight_sharing_across_methods_enabled()) { ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache( handle, named_data_map, method_name, weights_blob_key)); @@ -1025,475 +1020,6 @@ class ET_EXPERIMENTAL CudaBackend final return Error::Ok; } - static Error validate_fqn_weight_view(const CudaFqnWeightEntry& entry) { - uint64_t item_size = 0; - switch (static_cast(entry.dtype)) { - case slim::c10::ScalarType::Byte: - case slim::c10::ScalarType::Char: - case slim::c10::ScalarType::Bool: - item_size = 1; - break; - case slim::c10::ScalarType::Short: - case slim::c10::ScalarType::Half: - case slim::c10::ScalarType::BFloat16: - item_size = 2; - break; - case slim::c10::ScalarType::Int: - case slim::c10::ScalarType::Float: - item_size = 4; - break; - case slim::c10::ScalarType::Long: - item_size = 8; - break; - default: - return Error::InvalidProgram; - } - - ET_CHECK_OR_RETURN_ERROR( - entry.storage_nbytes <= std::numeric_limits::max(), - InvalidProgram, - "CUDA FQN storage '%s' is too large for this platform", - entry.storage_key.c_str()); - - bool empty = false; - uint64_t last_element = static_cast(entry.storage_offset); - for (size_t dim = 0; dim < entry.sizes.size(); ++dim) { - const uint64_t size = static_cast(entry.sizes[dim]); - const uint64_t stride = static_cast(entry.strides[dim]); - if (size == 0) { - empty = true; - break; - } - const uint64_t extent = size - 1; - ET_CHECK_OR_RETURN_ERROR( - extent == 0 || - stride <= std::numeric_limits::max() / extent, - InvalidProgram, - "CUDA FQN weight '%s' has overflowing shape/stride metadata", - entry.fqn.c_str()); - const uint64_t span = stride * extent; - ET_CHECK_OR_RETURN_ERROR( - last_element <= std::numeric_limits::max() - span, - InvalidProgram, - "CUDA FQN weight '%s' has overflowing storage metadata", - entry.fqn.c_str()); - last_element += span; - } - - uint64_t required_nbytes = 0; - if (!empty) { - ET_CHECK_OR_RETURN_ERROR( - last_element < std::numeric_limits::max() && - last_element + 1 <= - std::numeric_limits::max() / item_size, - InvalidProgram, - "CUDA FQN weight '%s' has overflowing storage size", - entry.fqn.c_str()); - required_nbytes = (last_element + 1) * item_size; - } - ET_CHECK_OR_RETURN_ERROR( - required_nbytes <= entry.storage_nbytes, - InvalidProgram, - "CUDA FQN weight '%s' requires %llu bytes from a %llu-byte storage", - entry.fqn.c_str(), - static_cast(required_nbytes), - static_cast(entry.storage_nbytes)); - return Error::Ok; - } - - Error acquire_fqn_weight_storage( - const NamedDataMap* named_data_map, - const CudaFqnWeightEntry& entry, - const std::vector* mutable_group, - const std::unordered_map& storage_scope_by_key, - int device_index, - std::shared_ptr& storage, - bool& reused) const { - reused = false; - ET_CHECK_OR_RETURN_ERROR( - named_data_map != nullptr, - InvalidArgument, - "CUDA FQN weights require a named data map"); - - const auto device_type = - static_cast(entry.device_type); - const bool is_cuda_storage = device_type == slim::c10::DeviceType::CUDA; - const int storage_device_index = is_cuda_storage ? device_index : 0; - - uintptr_t mutable_scope = 0; - if (!entry.shareable) { - const auto scope = storage_scope_by_key.find(entry.storage_key); - ET_CHECK_OR_RETURN_ERROR( - scope != storage_scope_by_key.end(), - NotFound, - "CUDA mutable FQN storage '%s' is missing from named data", - entry.storage_key.c_str()); - mutable_scope = scope->second; - } - - const auto cache_key = [&](const CudaFqnWeightEntry& item) { - if (item.shareable) { - // Immutable bytes are safe to reuse across methods and model - // instances solely by content identity. - return std::string("immutable:") + item.storage_key + - (is_cuda_storage ? "@cuda:" + std::to_string(device_index) - : "@cpu"); - } - // Stateful buffers with identical initial bytes are not interchangeable. - // Scope their logical FQN identity to the underlying PTD instance so - // methods in one model share state without leaking it to another live - // model instance. - return std::string("mutable:") + std::to_string(mutable_scope) + ":" + - item.fqn + - (is_cuda_storage ? "@cuda:" + std::to_string(device_index) : "@cpu"); - }; - - std::vector cache_keys; - if (entry.shareable) { - cache_keys.push_back(cache_key(entry)); - } else { - ET_CHECK_OR_RETURN_ERROR( - mutable_group != nullptr && !mutable_group->empty(), - InvalidProgram, - "CUDA mutable FQN storage group %u is empty", - entry.storage_group); - cache_keys.reserve(mutable_group->size()); - for (const CudaFqnWeightEntry* alias : *mutable_group) { - ET_CHECK_OR_RETURN_ERROR( - alias != nullptr && !alias->shareable && - alias->storage_nbytes == entry.storage_nbytes && - alias->device_type == entry.device_type, - InvalidProgram, - "CUDA mutable FQN storage group %u is inconsistent", - entry.storage_group); - cache_keys.push_back(cache_key(*alias)); - } - } - - std::unique_lock cache_lock(fqn_weight_storage_mutex_); - for (const std::string& key : cache_keys) { - auto cached = shared_fqn_weight_storages_.find(key); - if (cached != shared_fqn_weight_storages_.end()) { - std::shared_ptr candidate = cached->second.lock(); - if (candidate != nullptr) { - ET_CHECK_OR_RETURN_ERROR( - candidate->nbytes == entry.storage_nbytes && - candidate->device_type == device_type && - candidate->device_index == storage_device_index, - InvalidProgram, - "CUDA FQN storage '%s' has inconsistent allocation metadata", - entry.storage_key.c_str()); - ET_CHECK_OR_RETURN_ERROR( - storage == nullptr || storage.get() == candidate.get(), - InvalidProgram, - "CUDA mutable FQN storage group %u resolves to multiple allocations", - entry.storage_group); - storage = std::move(candidate); - } - } - } - if (storage != nullptr) { - for (const std::string& key : cache_keys) { - shared_fqn_weight_storages_[key] = storage; - } - reused = true; - return Error::Ok; - } - - void* storage_data = nullptr; - const size_t allocation_size = - std::max(1, static_cast(entry.storage_nbytes)); - if (is_cuda_storage) { - const cudaError_t allocation_error = - cudaMalloc(&storage_data, allocation_size); - if (allocation_error != cudaSuccess) { - ET_LOG( - Error, - "cudaMalloc failed for FQN storage '%s': %s", - entry.storage_key.c_str(), - cudaGetErrorString(allocation_error)); - return Error::MemoryAllocationFailed; - } - } else { - storage_data = std::malloc(allocation_size); - if (storage_data == nullptr) { - ET_LOG( - Error, - "malloc failed for CPU FQN storage '%s'", - entry.storage_key.c_str()); - return Error::MemoryAllocationFailed; - } - } - - const auto free_storage_data = [&]() { - if (is_cuda_storage) { - (void)cudaFree(storage_data); - } else { - std::free(storage_data); - } - }; - - if (entry.storage_nbytes > 0) { - auto host_data = named_data_map->get_data(entry.storage_key.c_str()); - if (!host_data.ok()) { - free_storage_data(); - ET_LOG( - Error, - "CUDA FQN storage '%s' is missing", - entry.storage_key.c_str()); - return Error::NotFound; - } - if (host_data->size() != entry.storage_nbytes) { - free_storage_data(); - ET_LOG( - Error, - "CUDA FQN storage '%s' has size %zu, expected %llu", - entry.storage_key.c_str(), - host_data->size(), - static_cast(entry.storage_nbytes)); - return Error::InvalidProgram; - } - cudaError_t copy_error = cudaSuccess; - if (is_cuda_storage) { - copy_error = cudaMemcpy( - storage_data, - host_data->data(), - static_cast(entry.storage_nbytes), - cudaMemcpyHostToDevice); - } else { - std::memcpy( - storage_data, - host_data->data(), - static_cast(entry.storage_nbytes)); - } - host_data->Free(); - if (copy_error != cudaSuccess) { - free_storage_data(); - ET_LOG( - Error, - "cudaMemcpy failed for FQN storage '%s': %s", - entry.storage_key.c_str(), - cudaGetErrorString(copy_error)); - return Error::Internal; - } - } - - storage = std::make_shared( - storage_data, - static_cast(entry.storage_nbytes), - device_type, - storage_device_index); - for (const std::string& key : cache_keys) { - shared_fqn_weight_storages_[key] = storage; - } - return Error::Ok; - } - - Error load_constants_from_fqn_manifest( - cuda::CudaDelegateHandle* handle, - const NamedDataMap* named_data_map, - const CudaFqnWeightManifest& manifest) const { - ET_CHECK_OR_RETURN_ERROR( - named_data_map != nullptr, - InvalidArgument, - "CUDA FQN weights require a named data map"); - ET_CHECK_OR_RETURN_ERROR( - handle->get_num_constants && handle->get_constant_name && - handle->get_constant_original_fqn && handle->get_constant_dtype && - handle->update_user_managed_constant_buffer_pairs, - NotSupported, - "AOTI library does not expose the APIs required by CUDA FQN weights"); - - size_t num_constants = 0; - ET_CHECK_OK_OR_RETURN_ERROR( - handle->get_num_constants(handle->container_handle, &num_constants), - "Failed to enumerate CUDA AOTI constants"); - std::unordered_map> - fqn_to_internal_names; - std::unordered_map fqn_to_aoti_dtype; - for (size_t index = 0; index < num_constants; ++index) { - const char* internal_name = nullptr; - const char* fqn = nullptr; - int32_t dtype = 0; - ET_CHECK_OK_OR_RETURN_ERROR( - handle->get_constant_name( - handle->container_handle, index, &internal_name), - "Failed to read CUDA AOTI constant name at index %zu", - index); - ET_CHECK_OK_OR_RETURN_ERROR( - handle->get_constant_original_fqn( - handle->container_handle, index, &fqn), - "Failed to read CUDA AOTI constant FQN at index %zu", - index); - ET_CHECK_OK_OR_RETURN_ERROR( - handle->get_constant_dtype(handle->container_handle, index, &dtype), - "Failed to read CUDA AOTI constant dtype at index %zu", - index); - if (internal_name != nullptr && fqn != nullptr && fqn[0] != '\0') { - fqn_to_internal_names[fqn].emplace_back(internal_name); - auto [metadata, inserted] = fqn_to_aoti_dtype.emplace(fqn, dtype); - ET_CHECK_OR_RETURN_ERROR( - inserted || metadata->second == dtype, - InvalidProgram, - "CUDA AOTI constant FQN '%s' has inconsistent metadata", - fqn); - } - } - - int device_index = 0; - ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index)); - struct LocalStorage { - std::string storage_key; - uint64_t storage_nbytes; - int32_t device_type; - std::shared_ptr storage; - }; - std::unordered_map local_storages; - std::unordered_map> - mutable_storage_groups; - for (const CudaFqnWeightEntry& entry : manifest.entries) { - if (!entry.shareable) { - mutable_storage_groups[entry.storage_group].push_back(&entry); - } - } - std::unordered_map storage_scope_by_key; - if (!mutable_storage_groups.empty()) { - // Method::init creates a distinct MergedDataMap wrapper per method, but - // get_key() forwards the stable pointer owned by the underlying PTD map. - // Use that pointer to scope mutable state to one live model instance. - auto num_keys = named_data_map->get_num_keys(); - ET_CHECK_OR_RETURN_ERROR( - num_keys.ok(), - InvalidProgram, - "Failed to enumerate CUDA named data while loading mutable FQN storage"); - storage_scope_by_key.reserve(num_keys.get()); - for (uint32_t index = 0; index < num_keys.get(); ++index) { - auto key = named_data_map->get_key(index); - ET_CHECK_OR_RETURN_ERROR( - key.ok() && key.get() != nullptr, - InvalidProgram, - "Failed to read CUDA named data key %u", - index); - storage_scope_by_key.emplace( - key.get(), reinterpret_cast(key.get())); - } - } - std::vector pairs; - pairs.reserve(manifest.entries.size()); - std::unordered_set bound_fqns; - size_t reused_storages = 0; - handle->fqn_weight_tensors.reserve(manifest.entries.size()); - - for (const CudaFqnWeightEntry& entry : manifest.entries) { - ET_CHECK_OK_OR_RETURN_ERROR( - validate_fqn_weight_view(entry), - "Invalid CUDA FQN view '%s'", - entry.fqn.c_str()); - auto internal_names = fqn_to_internal_names.find(entry.fqn); - ET_CHECK_OR_RETURN_ERROR( - internal_names != fqn_to_internal_names.end(), - InvalidProgram, - "CUDA FQN weight '%s' is not present in its AOTI library", - entry.fqn.c_str()); - const auto aoti_dtype = fqn_to_aoti_dtype.find(entry.fqn); - ET_CHECK_OR_RETURN_ERROR( - aoti_dtype != fqn_to_aoti_dtype.end() && - aoti_dtype->second == entry.dtype, - InvalidProgram, - "CUDA FQN weight '%s' dtype does not match its AOTI library " - "(manifest=%d, AOTI=%d)", - entry.fqn.c_str(), - entry.dtype, - aoti_dtype == fqn_to_aoti_dtype.end() ? -1 : aoti_dtype->second); - ET_CHECK_OR_RETURN_ERROR( - bound_fqns.emplace(entry.fqn).second, - InvalidProgram, - "CUDA FQN weight '%s' appears more than once in its manifest", - entry.fqn.c_str()); - - const std::string local_key = entry.shareable - ? "shared:" + entry.storage_key - : "local:" + std::to_string(entry.storage_group); - auto local_storage = local_storages.find(local_key); - std::shared_ptr storage; - if (local_storage == local_storages.end()) { - bool reused = false; - const auto mutable_entries = - mutable_storage_groups.find(entry.storage_group); - const std::vector* mutable_group = - entry.shareable ? nullptr - : (mutable_entries == mutable_storage_groups.end() - ? nullptr - : &mutable_entries->second); - ET_CHECK_OK_OR_RETURN_ERROR( - acquire_fqn_weight_storage( - named_data_map, - entry, - mutable_group, - storage_scope_by_key, - device_index, - storage, - reused), - "Failed to load CUDA FQN storage '%s'", - entry.storage_key.c_str()); - reused_storages += reused ? 1 : 0; - local_storages.emplace( - local_key, - LocalStorage{ - entry.storage_key, - entry.storage_nbytes, - entry.device_type, - storage}); - handle->fqn_weight_storages.push_back(storage); - } else { - ET_CHECK_OR_RETURN_ERROR( - local_storage->second.storage_key == entry.storage_key && - local_storage->second.storage_nbytes == entry.storage_nbytes && - local_storage->second.device_type == entry.device_type, - InvalidProgram, - "CUDA FQN storage group %u has inconsistent backing storage", - entry.storage_group); - storage = local_storage->second.storage; - } - - auto tensor = std::make_unique(slim::from_blob( - storage->data, - slim::makeArrayRef(entry.sizes), - slim::makeArrayRef(entry.strides), - static_cast(entry.dtype), - Device( - static_cast(entry.device_type), - entry.device_type == - static_cast(slim::c10::DeviceType::CUDA) - ? device_index - : 0), - entry.storage_offset)); - AtenTensorHandle tensor_handle = - reinterpret_cast(tensor.get()); - handle->fqn_weight_tensors.push_back(std::move(tensor)); - for (const std::string& internal_name : internal_names->second) { - pairs.push_back({internal_name.c_str(), tensor_handle}); - } - } - - ET_CHECK_OK_OR_RETURN_ERROR( - handle->update_user_managed_constant_buffer_pairs( - handle->container_handle, - pairs.data(), - pairs.size(), - /*use_inactive=*/false, - /*validate_full_update=*/true), - "Failed to bind CUDA FQN weights"); - ET_LOG( - Info, - "Loaded %zu CUDA FQN views from %zu physical storages (%zu reused " - "across methods)", - manifest.entries.size(), - local_storages.size(), - reused_storages); - return Error::Ok; - } - // Load constants for a method using per-weight caching. // Returns Error::Ok on success. // @@ -1748,12 +1274,7 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::unordered_map shared_constant_tensors_; - // New-format artifacts share immutable physical storages by their - // content-addressed named-data key. Weak ownership lets the allocation be - // reclaimed after the last delegate using it is destroyed. - mutable std::mutex fqn_weight_storage_mutex_; - mutable std::unordered_map> - shared_fqn_weight_storages_; + mutable CudaWeightCache fqn_weight_cache_; }; } // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 585e2eca7db..2734d68fc1f 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -218,7 +219,7 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { // CUDA graph state (warmup, capture, replay, static buffers) CudaGraphState cuda_graph_state; - // Per-storage weight artifacts keep the CUDA allocations and the original + // Per-FQN weight artifacts keep the allocations and their // SlimTensor handles alive for as long as AOTI may reference their views. std::vector> fqn_weight_storages; std::vector> fqn_weight_tensors; diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp new file mode 100644 index 00000000000..8bb00514828 --- /dev/null +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -0,0 +1,398 @@ +/* + * 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. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +using aoti::AOTInductorConstantMapEntry; +using aoti::AtenTensorHandle; +using aoti::slim::SlimTensor; +using aoti::slim::c10::Device; +using runtime::Error; +using runtime::NamedDataMap; + +Error CudaWeightCache::validate_view(const CudaFqnWeightEntry& entry) { + uint64_t item_size = 0; + switch (static_cast(entry.dtype)) { + case aoti::slim::c10::ScalarType::Byte: + case aoti::slim::c10::ScalarType::Char: + case aoti::slim::c10::ScalarType::Bool: + item_size = 1; + break; + case aoti::slim::c10::ScalarType::Short: + case aoti::slim::c10::ScalarType::Half: + case aoti::slim::c10::ScalarType::BFloat16: + item_size = 2; + break; + case aoti::slim::c10::ScalarType::Int: + case aoti::slim::c10::ScalarType::Float: + item_size = 4; + break; + case aoti::slim::c10::ScalarType::Long: + item_size = 8; + break; + default: + return Error::InvalidProgram; + } + + ET_CHECK_OR_RETURN_ERROR( + entry.storage_nbytes <= std::numeric_limits::max(), + InvalidProgram, + "CUDA FQN storage '%s' is too large for this platform", + entry.storage_key.c_str()); + + bool empty = false; + uint64_t last_element = static_cast(entry.storage_offset); + for (size_t dim = 0; dim < entry.sizes.size(); ++dim) { + const uint64_t size = static_cast(entry.sizes[dim]); + const uint64_t stride = static_cast(entry.strides[dim]); + if (size == 0) { + empty = true; + break; + } + const uint64_t extent = size - 1; + ET_CHECK_OR_RETURN_ERROR( + extent == 0 || stride <= std::numeric_limits::max() / extent, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing shape/stride metadata", + entry.fqn.c_str()); + const uint64_t span = stride * extent; + ET_CHECK_OR_RETURN_ERROR( + last_element <= std::numeric_limits::max() - span, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing storage metadata", + entry.fqn.c_str()); + last_element += span; + } + + uint64_t required_nbytes = 0; + if (!empty) { + ET_CHECK_OR_RETURN_ERROR( + last_element < std::numeric_limits::max() && + last_element + 1 <= + std::numeric_limits::max() / item_size, + InvalidProgram, + "CUDA FQN weight '%s' has overflowing storage size", + entry.fqn.c_str()); + required_nbytes = (last_element + 1) * item_size; + } + ET_CHECK_OR_RETURN_ERROR( + required_nbytes <= entry.storage_nbytes, + InvalidProgram, + "CUDA FQN weight '%s' requires %llu bytes from a %llu-byte storage", + entry.fqn.c_str(), + static_cast(required_nbytes), + static_cast(entry.storage_nbytes)); + return Error::Ok; +} + +Error CudaWeightCache::acquire_storage( + const NamedDataMap* named_data_map, + const CudaFqnWeightEntry& entry, + uintptr_t logical_scope, + int device_index, + std::shared_ptr& storage, + bool& reused) const { + reused = false; + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "CUDA FQN weights require a named data map"); + + const auto device_type = + static_cast(entry.device_type); + const bool is_cuda_storage = device_type == aoti::slim::c10::DeviceType::CUDA; + const int storage_device_index = is_cuda_storage ? device_index : 0; + const std::string cache_key = entry.storage_key + "@" + + std::to_string(logical_scope) + + (is_cuda_storage ? "@cuda:" + std::to_string(device_index) : "@cpu"); + + std::unique_lock lock(mutex_); + auto cached = storages_.find(cache_key); + if (cached != storages_.end()) { + storage = cached->second.lock(); + } + if (storage != nullptr) { + ET_CHECK_OR_RETURN_ERROR( + storage->nbytes == entry.storage_nbytes && + storage->device_type == device_type && + storage->device_index == storage_device_index, + InvalidProgram, + "CUDA FQN weight '%s' has inconsistent allocation metadata", + entry.fqn.c_str()); + reused = true; + return Error::Ok; + } + + auto host_data = named_data_map->get_data(entry.storage_key.c_str()); + ET_CHECK_OR_RETURN_ERROR( + host_data.ok(), + NotFound, + "CUDA FQN storage '%s' is missing from named data", + entry.storage_key.c_str()); + if (host_data->size() != entry.storage_nbytes) { + const size_t actual_size = host_data->size(); + host_data->Free(); + ET_LOG( + Error, + "CUDA FQN storage '%s' has size %zu, expected %llu", + entry.storage_key.c_str(), + actual_size, + static_cast(entry.storage_nbytes)); + return Error::InvalidProgram; + } + + void* storage_data = nullptr; + const size_t allocation_size = + std::max(1, static_cast(entry.storage_nbytes)); + if (is_cuda_storage) { + const cudaError_t allocation_error = + cudaMalloc(&storage_data, allocation_size); + if (allocation_error != cudaSuccess) { + host_data->Free(); + ET_LOG( + Error, + "cudaMalloc failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(allocation_error)); + return Error::MemoryAllocationFailed; + } + } else { + storage_data = std::malloc(allocation_size); + if (storage_data == nullptr) { + host_data->Free(); + ET_LOG( + Error, + "malloc failed for CPU FQN storage '%s'", + entry.storage_key.c_str()); + return Error::MemoryAllocationFailed; + } + } + + const auto free_storage_data = [&]() { + if (is_cuda_storage) { + (void)cudaFree(storage_data); + } else { + std::free(storage_data); + } + }; + + cudaError_t copy_error = cudaSuccess; + if (entry.storage_nbytes > 0) { + if (is_cuda_storage) { + copy_error = cudaMemcpy( + storage_data, + host_data->data(), + static_cast(entry.storage_nbytes), + cudaMemcpyHostToDevice); + } else { + std::memcpy( + storage_data, + host_data->data(), + static_cast(entry.storage_nbytes)); + } + } + host_data->Free(); + if (copy_error != cudaSuccess) { + free_storage_data(); + ET_LOG( + Error, + "cudaMemcpy failed for FQN storage '%s': %s", + entry.storage_key.c_str(), + cudaGetErrorString(copy_error)); + return Error::Internal; + } + + storage = std::make_shared( + storage_data, + static_cast(entry.storage_nbytes), + device_type, + storage_device_index); + storages_[cache_key] = storage; + return Error::Ok; +} + +Error CudaWeightCache::load( + CudaDelegateHandle* handle, + const NamedDataMap* named_data_map, + const CudaFqnWeightManifest& manifest) const { + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "CUDA FQN weights require a named data map"); + ET_CHECK_OR_RETURN_ERROR( + handle->get_num_constants && handle->get_constant_name && + handle->get_constant_original_fqn && handle->get_constant_dtype && + handle->update_user_managed_constant_buffer_pairs, + NotSupported, + "AOTI library does not expose the APIs required by CUDA FQN weights"); + + size_t num_constants = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_num_constants(handle->container_handle, &num_constants), + "Failed to enumerate CUDA AOTI constants"); + std::unordered_map> + fqn_to_internal_names; + std::unordered_map fqn_to_aoti_dtype; + for (size_t index = 0; index < num_constants; ++index) { + const char* internal_name = nullptr; + const char* fqn = nullptr; + int32_t dtype = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_name( + handle->container_handle, index, &internal_name), + "Failed to read CUDA AOTI constant name at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_original_fqn( + handle->container_handle, index, &fqn), + "Failed to read CUDA AOTI constant FQN at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_constant_dtype(handle->container_handle, index, &dtype), + "Failed to read CUDA AOTI constant dtype at index %zu", + index); + if (internal_name != nullptr && fqn != nullptr && fqn[0] != '\0') { + fqn_to_internal_names[fqn].emplace_back(internal_name); + auto [metadata, inserted] = fqn_to_aoti_dtype.emplace(fqn, dtype); + ET_CHECK_OR_RETURN_ERROR( + inserted || metadata->second == dtype, + InvalidProgram, + "CUDA AOTI constant FQN '%s' has inconsistent metadata", + fqn); + } + } + + int device_index = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index)); + // Each method receives a different MergedDataMap wrapper, but get_key() + // forwards the stable key owned by the shared PTD map. Use that identity to + // scope the process-wide FQN cache to one loaded model. + std::unordered_map key_scopes; + auto num_keys = named_data_map->get_num_keys(); + ET_CHECK_OR_RETURN_ERROR( + num_keys.ok(), + InvalidProgram, + "Failed to enumerate CUDA named data while loading FQN weights"); + key_scopes.reserve(num_keys.get()); + for (uint32_t index = 0; index < num_keys.get(); ++index) { + auto key = named_data_map->get_key(index); + ET_CHECK_OR_RETURN_ERROR( + key.ok() && key.get() != nullptr, + InvalidProgram, + "Failed to read CUDA named data key %u", + index); + key_scopes.emplace(key.get(), reinterpret_cast(key.get())); + } + std::vector pairs; + pairs.reserve(manifest.entries.size()); + std::unordered_set bound_fqns; + size_t reused_storages = 0; + handle->fqn_weight_tensors.reserve(manifest.entries.size()); + + for (const CudaFqnWeightEntry& entry : manifest.entries) { + ET_CHECK_OK_OR_RETURN_ERROR( + validate_view(entry), "Invalid CUDA FQN view '%s'", entry.fqn.c_str()); + auto internal_names = fqn_to_internal_names.find(entry.fqn); + ET_CHECK_OR_RETURN_ERROR( + internal_names != fqn_to_internal_names.end(), + InvalidProgram, + "CUDA FQN weight '%s' is not present in its AOTI library", + entry.fqn.c_str()); + const auto aoti_dtype = fqn_to_aoti_dtype.find(entry.fqn); + ET_CHECK_OR_RETURN_ERROR( + aoti_dtype != fqn_to_aoti_dtype.end() && + aoti_dtype->second == entry.dtype, + InvalidProgram, + "CUDA FQN weight '%s' dtype does not match its AOTI library " + "(manifest=%d, AOTI=%d)", + entry.fqn.c_str(), + entry.dtype, + aoti_dtype == fqn_to_aoti_dtype.end() ? -1 : aoti_dtype->second); + ET_CHECK_OR_RETURN_ERROR( + bound_fqns.emplace(entry.fqn).second, + InvalidProgram, + "CUDA FQN weight '%s' appears more than once in its manifest", + entry.fqn.c_str()); + + std::shared_ptr storage; + bool reused = false; + const auto logical_scope = key_scopes.find(entry.storage_key); + ET_CHECK_OR_RETURN_ERROR( + logical_scope != key_scopes.end(), + NotFound, + "CUDA FQN storage '%s' is missing from named data", + entry.storage_key.c_str()); + ET_CHECK_OK_OR_RETURN_ERROR( + acquire_storage( + named_data_map, + entry, + logical_scope->second, + device_index, + storage, + reused), + "Failed to load CUDA FQN storage '%s'", + entry.storage_key.c_str()); + reused_storages += reused ? 1 : 0; + handle->fqn_weight_storages.push_back(storage); + + auto tensor = std::make_unique(aoti::slim::from_blob( + storage->data, + aoti::slim::makeArrayRef(entry.sizes), + aoti::slim::makeArrayRef(entry.strides), + static_cast(entry.dtype), + Device( + static_cast(entry.device_type), + entry.device_type == + static_cast(aoti::slim::c10::DeviceType::CUDA) + ? device_index + : 0), + entry.storage_offset)); + AtenTensorHandle tensor_handle = + reinterpret_cast(tensor.get()); + handle->fqn_weight_tensors.push_back(std::move(tensor)); + for (const std::string& internal_name : internal_names->second) { + pairs.push_back({internal_name.c_str(), tensor_handle}); + } + } + + ET_CHECK_OK_OR_RETURN_ERROR( + handle->update_user_managed_constant_buffer_pairs( + handle->container_handle, + pairs.data(), + pairs.size(), + /*use_inactive=*/false, + /*validate_full_update=*/true), + "Failed to bind CUDA FQN weights"); + ET_LOG( + Info, + "Loaded %zu CUDA FQN weights (%zu reused across methods)", + manifest.entries.size(), + reused_storages); + return Error::Ok; +} + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_weight_cache.h b/backends/cuda/runtime/cuda_weight_cache.h new file mode 100644 index 00000000000..f34cf41e1d9 --- /dev/null +++ b/backends/cuda/runtime/cuda_weight_cache.h @@ -0,0 +1,47 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +class CudaWeightCache final { + public: + runtime::Error load( + CudaDelegateHandle* handle, + const runtime::NamedDataMap* named_data_map, + const CudaFqnWeightManifest& manifest) const; + + private: + static runtime::Error validate_view(const CudaFqnWeightEntry& entry); + + runtime::Error acquire_storage( + const runtime::NamedDataMap* named_data_map, + const CudaFqnWeightEntry& entry, + uintptr_t logical_scope, + int device_index, + std::shared_ptr& storage, + bool& reused) const; + + mutable std::mutex mutex_; + mutable std::unordered_map> + storages_; +}; + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_weight_manifest.h b/backends/cuda/runtime/cuda_weight_manifest.h index a451f61a667..10f5fcb5734 100644 --- a/backends/cuda/runtime/cuda_weight_manifest.h +++ b/backends/cuda/runtime/cuda_weight_manifest.h @@ -17,20 +17,18 @@ namespace executorch::backends::cuda { -constexpr char kCudaFqnWeightsMagic[] = "ETCUDAFQN2"; +constexpr char kCudaFqnWeightsMagic[] = "ETCUDAFQN3"; constexpr size_t kCudaFqnWeightsMagicSize = sizeof(kCudaFqnWeightsMagic) - 1; struct CudaFqnWeightEntry { std::string fqn; std::string storage_key; - uint32_t storage_group{0}; uint64_t storage_nbytes{0}; int32_t dtype{0}; int32_t device_type{0}; int64_t storage_offset{0}; std::vector sizes; std::vector strides; - bool shareable{false}; }; struct CudaFqnWeightManifest { @@ -81,14 +79,6 @@ class CudaWeightManifestReader final { return true; } - bool read_u8(uint8_t& value) { - if (remaining() < 1) { - return false; - } - value = *cursor_++; - return true; - } - bool read_u32(uint32_t& value) { uint64_t wide = 0; if (!read_unsigned(wide, 4)) { @@ -185,10 +175,8 @@ inline executorch::runtime::Error parse_cuda_fqn_weight_manifest( for (uint32_t index = 0; index < num_entries; ++index) { CudaFqnWeightEntry entry; uint32_t ndim = 0; - uint8_t shareable = 0; if (!reader.read_string(entry.fqn) || entry.fqn.empty() || !reader.read_string(entry.storage_key) || entry.storage_key.empty() || - !reader.read_u32(entry.storage_group) || !reader.read_u64(entry.storage_nbytes) || !reader.read_i32(entry.dtype) || !is_supported_cuda_fqn_dtype(entry.dtype) || @@ -211,11 +199,9 @@ inline executorch::runtime::Error parse_cuda_fqn_weight_manifest( return Error::InvalidProgram; } } - if (!reader.read_u8(shareable) || shareable > 1 || - entry.storage_offset < 0) { + if (entry.storage_offset < 0) { return Error::InvalidProgram; } - entry.shareable = shareable != 0; manifest.entries.push_back(std::move(entry)); } diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 0da63d55513..228a8470908 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -124,10 +124,12 @@ def define_common_targets(is_fbcode = False): srcs = [ "cuda_backend.cpp", "cuda_mutable_state.cpp", + "cuda_weight_cache.cpp", ], headers = [ "cuda_delegate_handle.h", "cuda_mutable_state.h", + "cuda_weight_cache.h", "cuda_weight_manifest.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) diff --git a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp b/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp index 8fa643fa555..0ed3d383abb 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp @@ -46,7 +46,6 @@ std::vector valid_manifest( append_u32(output, 1); // entries append_string(output, "model.weight"); append_string(output, "storage-key"); - append_u32(output, 7); // method-local storage group append_u64(output, 24); // storage bytes append_u32(output, dtype); // dtype append_u32(output, device_type); // device type (CUDA) @@ -56,7 +55,6 @@ std::vector valid_manifest( append_u64(output, 3); append_u64(output, 3); append_u64(output, 1); - output.push_back(1); // shareable return output; } @@ -79,13 +77,11 @@ TEST(CudaWeightManifestTest, ParsesVersionedManifest) { const auto& entry = manifest.entries[0]; EXPECT_EQ(entry.fqn, "model.weight"); EXPECT_EQ(entry.storage_key, "storage-key"); - EXPECT_EQ(entry.storage_group, 7u); EXPECT_EQ(entry.storage_nbytes, 24u); EXPECT_EQ(entry.dtype, 6); EXPECT_EQ(entry.device_type, 1); EXPECT_EQ(entry.sizes, (std::vector{2, 3})); EXPECT_EQ(entry.strides, (std::vector{3, 1})); - EXPECT_TRUE(entry.shareable); } TEST(CudaWeightManifestTest, RejectsTruncationAndTrailingData) { diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index b3c4f69cc39..51c7be8d381 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -9,20 +9,25 @@ import os import tempfile import unittest -from types import SimpleNamespace from typing import Tuple from unittest.mock import patch import torch from executorch.backends.cuda.cuda_backend import ( - _encode_fqn_weight_manifest, - _FQN_WEIGHTS_MAGIC, - _materialize_fqn_weights, - _stateful_buffer_fqns, + _aoti_device_type_for_weight, CudaBackend, ) from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CPU, + AOTI_DEVICE_TYPE_CUDA, + CUDA_FQN_WEIGHTS_MAGIC, + CudaWeightCollector, + encode_cuda_weight_manifest, +) from executorch.exir._serialize._cord import FileBackedData +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.backend.backend_details import PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import PartitionResult from executorch.exir.delegate import executorch_call_delegate @@ -33,14 +38,20 @@ class TestCudaLowMemoryExport(unittest.TestCase): - def test_all_buffers_are_model_instance_local(self) -> None: - signature = SimpleNamespace( - buffers=("persistent", "conditional_cache"), - buffers_to_mutate={"copy_out": "explicitly_mutated"}, + @staticmethod + def _materialize(weights, directory, device_type=AOTI_DEVICE_TYPE_CPU): + return CudaWeightCollector().materialize( + weights, directory, lambda _: device_type ) - self.assertEqual( - _stateful_buffer_fqns(signature), - {"persistent", "conditional_cache", "explicitly_mutated"}, + + @staticmethod + def _parent_result(so_key: str) -> PreprocessResult: + store = NamedDataStore() + store.add_named_data(so_key, so_key.encode()) + store.add_named_data("empty_weights", b"", external_tag="aoti_cuda_blob") + return PreprocessResult( + processed_bytes=f"{so_key}\nempty_weights".encode(), + data_store_output=store.get_named_data_store_output(), ) @patch.object(CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True) @@ -70,13 +81,9 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: ) with tempfile.TemporaryDirectory() as directory: - artifact = _materialize_fqn_weights( - weights, directory, mutated_fqns={"second"} - ) + artifact = self._materialize(weights, directory) self.assertEqual(2, len(artifact.entries)) self.assertEqual(2, len(artifact.storages)) - self.assertTrue(artifact.entries[0].shareable) - self.assertFalse(artifact.entries[1].shareable) self.assertEqual(artifact.entries[0].device_type, 0) self.assertEqual(artifact.entries[1].device_type, 0) self.assertEqual( @@ -88,8 +95,8 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: artifact.storages[artifact.entries[1].storage_key].to_bytes(), ) - manifest = _encode_fqn_weight_manifest("so-key", artifact.entries) - self.assertTrue(manifest.startswith(_FQN_WEIGHTS_MAGIC)) + manifest = encode_cuda_weight_manifest("so-key", artifact.entries) + self.assertTrue(manifest.startswith(CUDA_FQN_WEIGHTS_MAGIC)) self.assertIn(b"first", manifest) self.assertIn(b"second", manifest) for storage in artifact.storages.values(): @@ -104,12 +111,14 @@ def test_low_memory_cpu_clones_keep_cuda_device_type(self, _) -> None: weights = Weights({"weight": (tensor, TensorProperties(tensor))}) with tempfile.TemporaryDirectory() as directory: - artifact = _materialize_fqn_weights(weights, directory, set()) - self.assertEqual(artifact.entries[0].device_type, 1) + artifact = CudaWeightCollector().materialize( + weights, directory, _aoti_device_type_for_weight + ) + self.assertEqual(artifact.entries[0].device_type, AOTI_DEVICE_TYPE_CUDA) for storage in artifact.storages.values(): storage.close() - def test_views_share_one_physical_storage(self) -> None: + def test_different_fqn_views_have_distinct_logical_storage(self) -> None: base = torch.arange(12, dtype=torch.float32).reshape(3, 4) view = base[:, 1:] weights = Weights( @@ -122,11 +131,11 @@ def test_views_share_one_physical_storage(self) -> None: ) with tempfile.TemporaryDirectory() as directory: - artifact = _materialize_fqn_weights(weights, directory, set()) - self.assertEqual(1, len(artifact.storages)) - self.assertEqual( - artifact.entries[0].storage_group, - artifact.entries[1].storage_group, + artifact = self._materialize(weights, directory) + self.assertEqual(2, len(artifact.storages)) + self.assertNotEqual( + artifact.entries[0].storage_key, + artifact.entries[1].storage_key, ) self.assertEqual(1, artifact.entries[1].storage_offset) self.assertEqual((3, 3), artifact.entries[1].sizes) @@ -134,7 +143,7 @@ def test_views_share_one_physical_storage(self) -> None: for storage in artifact.storages.values(): storage.close() - def test_identical_mutable_storages_remain_distinct_groups(self) -> None: + def test_identical_values_keep_distinct_fqn_keys(self) -> None: first = torch.zeros(4) second = torch.zeros(4) weights = Weights( @@ -145,23 +154,121 @@ def test_identical_mutable_storages_remain_distinct_groups(self) -> None: ) with tempfile.TemporaryDirectory() as directory: - artifact = _materialize_fqn_weights( - weights, directory, mutated_fqns={"first", "second"} - ) - self.assertEqual(1, len(artifact.storages)) - self.assertEqual( + artifact = self._materialize(weights, directory) + self.assertEqual(2, len(artifact.storages)) + self.assertNotEqual(artifact.entries[0].fqn, artifact.entries[1].fqn) + self.assertNotEqual( artifact.entries[0].storage_key, artifact.entries[1].storage_key, ) - self.assertNotEqual( - artifact.entries[0].storage_group, - artifact.entries[1].storage_group, - ) - self.assertFalse(artifact.entries[0].shareable) - self.assertFalse(artifact.entries[1].shareable) + store = NamedDataStore() + for entry in artifact.entries: + store.add_named_data( + entry.storage_key, artifact.storages[entry.storage_key] + ) + self.assertEqual(1, len(store.buffers)) for storage in artifact.storages.values(): storage.close() + def test_same_fqn_with_different_data_is_rejected_when_merged(self) -> None: + first = torch.zeros(4) + second = torch.ones(4) + + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize( + Weights({"weight": (first, TensorProperties(first))}), first_dir + ) + second_artifact = self._materialize( + Weights({"weight": (second, TensorProperties(second))}), second_dir + ) + first_key = first_artifact.entries[0].storage_key + second_key = second_artifact.entries[0].storage_key + self.assertEqual(first_key, second_key) + + collector = CudaWeightCollector() + collector.add_preprocess_result( + self._parent_result("first_so"), first_artifact, "cuda" + ) + with self.assertRaises(ValueError): + collector.add_preprocess_result( + self._parent_result("second_so"), second_artifact, "cuda" + ) + + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_same_fqn_with_different_metadata_is_rejected(self) -> None: + tensor = torch.arange(4) + view = tensor.reshape(2, 2) + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize( + Weights({"weight": (tensor, TensorProperties(tensor))}), first_dir + ) + second_artifact = self._materialize( + Weights({"weight": (tensor, TensorProperties(view))}), second_dir + ) + collector = CudaWeightCollector() + collector.add_preprocess_result( + self._parent_result("first_so"), first_artifact, "cuda" + ) + with self.assertRaisesRegex(ValueError, "different tensor metadata"): + collector.add_preprocess_result( + self._parent_result("second_so"), second_artifact, "cuda" + ) + + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_methods_share_one_collected_named_data_store(self) -> None: + tensor = torch.arange(4) + weights = Weights({"weight": (tensor, TensorProperties(tensor))}) + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize(weights, first_dir) + second_artifact = self._materialize(weights, second_dir) + first_result = self._parent_result("first_so") + second_result = self._parent_result("second_so") + collector = CudaWeightCollector() + collector.add_preprocess_result(first_result, first_artifact, "cuda") + collector.add_preprocess_result(second_result, second_artifact, "cuda") + collector.finish() + + self.assertIs( + first_result.data_store_output, second_result.data_store_output + ) + weight_key = first_artifact.entries[0].storage_key + self.assertIn( + weight_key, + first_result.data_store_output.external_data["aoti_cuda_blob"], + ) + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + + def test_device_is_part_of_the_fqn_key(self) -> None: + tensor = torch.zeros(4) + weights = Weights({"weight": (tensor, TensorProperties(tensor))}) + with ( + tempfile.TemporaryDirectory() as cpu_dir, + tempfile.TemporaryDirectory() as cuda_dir, + ): + cpu = self._materialize(weights, cpu_dir, AOTI_DEVICE_TYPE_CPU) + cuda = self._materialize(weights, cuda_dir, AOTI_DEVICE_TYPE_CUDA) + self.assertNotEqual(cpu.entries[0].storage_key, cuda.entries[0].storage_key) + for artifact in (cpu, cuda): + for storage in artifact.storages.values(): + storage.close() + def test_low_memory_weights_require_wrapper_so(self) -> None: tensor = torch.tensor([1], dtype=torch.int16) weights = Weights({"weight": (tensor, TensorProperties(tensor))}) From 27ae1e87c77fb17007466d2c25d0ddd17576b1e5 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 11:45:40 -0700 Subject: [PATCH 07/12] cuda: fold weight metadata into cache Keep serialized FQN metadata parsing, validation, allocation, and AOTI binding behind CudaWeightCache, and remove the standalone manifest abstraction.\n\nGenerated with Codex. --- backends/cuda/CMakeLists.txt | 6 +- backends/cuda/cuda_backend.py | 4 +- backends/cuda/cuda_weight_collector.py | 8 +- backends/cuda/runtime/cuda_backend.cpp | 20 +- backends/cuda/runtime/cuda_delegate_handle.h | 2 +- backends/cuda/runtime/cuda_weight_cache.cpp | 186 ++++++++++++++- backends/cuda/runtime/cuda_weight_cache.h | 33 ++- backends/cuda/runtime/cuda_weight_manifest.h | 211 ------------------ backends/cuda/runtime/targets.bzl | 5 +- ...anifest.cpp => test_cuda_weight_cache.cpp} | 59 +++-- backends/cuda/tests/test_cuda_partitioner.py | 12 +- 11 files changed, 262 insertions(+), 284 deletions(-) delete mode 100644 backends/cuda/runtime/cuda_weight_manifest.h rename backends/cuda/runtime/test/{test_cuda_weight_manifest.cpp => test_cuda_weight_cache.cpp} (58%) diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 7bb51e4b0eb..4d3a26582e5 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -451,8 +451,8 @@ if(BUILD_TESTING) target_compile_definitions(test_cuda_mutable_state PRIVATE CUDA_AVAILABLE=1) et_cxx_test( - test_cuda_weight_manifest SOURCES - runtime/test/test_cuda_weight_manifest.cpp EXTRA_LIBS aoti_cuda_backend + test_cuda_weight_cache SOURCES runtime/test/test_cuda_weight_cache.cpp + EXTRA_LIBS aoti_cuda_backend ) - target_compile_definitions(test_cuda_weight_manifest PRIVATE CUDA_AVAILABLE=1) + target_compile_definitions(test_cuda_weight_cache PRIVATE CUDA_AVAILABLE=1) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 7b2157fc4e0..bf4f66f6d32 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -72,7 +72,7 @@ def _is_cpu_clone_active() -> bool: def _aoti_device_type_for_weight(tensor: torch.Tensor) -> int: # Low-memory compilation clones lifted CUDA buffers onto CPU, while the # patched wrapper records them as CUDA constants. Mirror that target-device - # substitution in the manifest. Outside that scoped mode the serialized + # substitution in the metadata. Outside that scoped mode the serialized # tensor's actual device is the AOTI constant's device. if _is_cpu_clone_active() or tensor.device.type == "cuda": return AOTI_DEVICE_TYPE_CUDA @@ -911,7 +911,7 @@ def _is_low_memory_mode(compile_specs: List[CompileSpec]) -> bool: @classmethod def _weights_format(cls, compile_specs: List[CompileSpec]) -> str: # CUDA consumes the structured AOTI output directly and emits a - # versioned per-FQN manifest. This is backend-wide rather than a + # versioned per-FQN metadata payload. This is backend-wide rather than a # model/export-script option. return "pickle_weights" diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index bfd02d232a3..4a6b315a5b5 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -22,7 +22,7 @@ from executorch.exir.tensor import scalar_type_enum -CUDA_FQN_WEIGHTS_MAGIC = b"ETCUDAFQN3" +CUDA_WEIGHT_CACHE_MAGIC = b"ETCUDAFQN3" AOTI_DEVICE_TYPE_CPU = 0 AOTI_DEVICE_TYPE_CUDA = 1 @@ -103,11 +103,11 @@ def _storage_key(fqn: str, device_type: int) -> str: return f"cuda_fqn_weight:{device}:{fqn}" -def encode_cuda_weight_manifest( +def encode_cuda_weight_metadata( so_blob_key: str, entries: List[CudaWeightEntry] ) -> bytes: """Encode the per-method FQN-to-tensor metadata consumed by CUDA runtime.""" - output = bytearray(CUDA_FQN_WEIGHTS_MAGIC) + output = bytearray(CUDA_WEIGHT_CACHE_MAGIC) def write_string(value: str) -> None: encoded = value.encode("utf-8") @@ -310,7 +310,7 @@ def add_preprocess_result( data.close() self._entries.setdefault(entry.storage_key, entry) - result.processed_bytes = encode_cuda_weight_manifest( + result.processed_bytes = encode_cuda_weight_metadata( so_blob_key, artifact.entries ) self._results.append(result) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 773992db068..29cba8b5ada 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -166,7 +166,7 @@ class ET_EXPERIMENTAL CudaBackend final return shared_cuda_stream_ != nullptr; } - // Enable the legacy dense-blob per-FQN cache. New manifest artifacts use + // Enable the legacy dense-blob per-FQN cache. New FQN artifacts use // their FQN-addressed data keys automatically. void set_weight_sharing_across_methods(bool enabled) { weight_sharing_across_methods_.store(enabled, std::memory_order_relaxed); @@ -330,15 +330,15 @@ class ET_EXPERIMENTAL CudaBackend final std::string so_blob_key; std::string weights_blob_key; - CudaFqnWeightManifest fqn_weight_manifest; + CudaWeightCache::Metadata fqn_weights; const bool has_fqn_weights = - is_cuda_fqn_weight_manifest(processed->data(), processed->size()); + CudaWeightCache::is_serialized(processed->data(), processed->size()); if (has_fqn_weights) { ET_CHECK_OK_OR_RETURN_ERROR( - parse_cuda_fqn_weight_manifest( - processed->data(), processed->size(), fqn_weight_manifest), - "Malformed CUDA FQN weight manifest"); - so_blob_key = fqn_weight_manifest.so_blob_key; + CudaWeightCache::parse( + processed->data(), processed->size(), fqn_weights), + "Malformed CUDA FQN weight metadata"); + so_blob_key = fqn_weights.so_blob_key; } else { ET_CHECK_OK_OR_RETURN_ERROR( executorch::backends::aoti::resolve_blob_keys( @@ -419,11 +419,11 @@ class ET_EXPERIMENTAL CudaBackend final handle->container_handle = container_handle; // Versioned artifacts load each (device, FQN) through the same process-wide - // cross-method cache model used by the legacy path. The manifest only adds + // cross-method cache model used by the legacy path. The payload only adds // the tensor metadata needed to reconstruct independently named PTD blobs. if (has_fqn_weights) { ET_CHECK_OK_OR_RETURN_ERROR( - fqn_weight_cache_.load(handle, named_data_map, fqn_weight_manifest)); + fqn_weight_cache_.load(handle, named_data_map, fqn_weights)); } else if (is_weight_sharing_across_methods_enabled()) { ET_CHECK_OK_OR_RETURN_ERROR(load_constants_with_cache( handle, named_data_map, method_name, weights_blob_key)); @@ -908,7 +908,7 @@ class ET_EXPERIMENTAL CudaBackend final // Whether to enable cross-method caching for legacy dense-blob artifacts. // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default - // OFF; versioned manifest artifacts do not consult this option. + // OFF; versioned FQN artifacts do not consult this option. std::atomic weight_sharing_across_methods_{false}; // --------------------------------------------------------------- diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 2734d68fc1f..32144ce139e 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -195,7 +195,7 @@ struct CudaGraphState { // CUDA-specific delegate handle that extends AOTIDelegateHandle. // This consolidates CUDA stream management into a single location. struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { - // Extra AOTI metadata used to validate per-FQN manifests before binding. + // Extra AOTI metadata used to validate per-FQN weights before binding. AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr}; // CUDA stream for this handle, support both shared mode and single mode. diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp index 8bb00514828..09b4bd76312 100644 --- a/backends/cuda/runtime/cuda_weight_cache.cpp +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -34,7 +34,175 @@ using aoti::slim::c10::Device; using runtime::Error; using runtime::NamedDataMap; -Error CudaWeightCache::validate_view(const CudaFqnWeightEntry& entry) { +namespace { + +class MetadataReader final { + public: + MetadataReader(const void* data, size_t size) + : cursor_(static_cast(data)), end_(cursor_ + size) {} + + bool skip(size_t size) { + if (remaining() < size) { + return false; + } + cursor_ += size; + return true; + } + + bool read_u32(uint32_t& value) { + uint64_t wide = 0; + if (!read_unsigned(wide, 4)) { + return false; + } + value = static_cast(wide); + return true; + } + + bool read_i32(int32_t& value) { + uint32_t raw = 0; + if (!read_u32(raw)) { + return false; + } + std::memcpy(&value, &raw, sizeof(value)); + return true; + } + + bool read_u64(uint64_t& value) { + return read_unsigned(value, 8); + } + + bool read_i64(int64_t& value) { + uint64_t raw = 0; + if (!read_u64(raw)) { + return false; + } + std::memcpy(&value, &raw, sizeof(value)); + return true; + } + + bool read_string(std::string& value) { + uint32_t size = 0; + if (!read_u32(size) || remaining() < size) { + return false; + } + value.assign(reinterpret_cast(cursor_), size); + cursor_ += size; + return true; + } + + bool empty() const { + return cursor_ == end_; + } + + private: + size_t remaining() const { + return static_cast(end_ - cursor_); + } + + bool read_unsigned(uint64_t& value, size_t width) { + if (remaining() < width) { + return false; + } + value = 0; + for (size_t index = 0; index < width; ++index) { + value |= static_cast(cursor_[index]) << (index * 8); + } + cursor_ += width; + return true; + } + + const uint8_t* cursor_; + const uint8_t* end_; +}; + +bool is_supported_dtype(int32_t dtype) { + switch (dtype) { + case 0: // Byte + case 1: // Char + case 2: // Short + case 3: // Int + case 4: // Long + case 5: // Half + case 6: // Float + case 11: // Bool + case 15: // BFloat16 + return true; + default: + return false; + } +} + +bool is_supported_device_type(int32_t device_type) { + return device_type == 0 || device_type == 1; // CPU or CUDA +} + +} // namespace + +bool CudaWeightCache::is_serialized(const void* data, size_t size) { + return data != nullptr && size >= kFormatMagicSize && + std::memcmp(data, kFormatMagic, kFormatMagicSize) == 0; +} + +Error CudaWeightCache::parse( + const void* data, + size_t size, + Metadata& metadata) { + if (!is_serialized(data, size)) { + return Error::InvalidProgram; + } + + MetadataReader reader(data, size); + if (!reader.skip(kFormatMagicSize) || + !reader.read_string(metadata.so_blob_key) || + metadata.so_blob_key.empty()) { + return Error::InvalidProgram; + } + + uint32_t num_entries = 0; + constexpr uint32_t kMaxEntries = 1U << 20; + if (!reader.read_u32(num_entries) || num_entries > kMaxEntries) { + return Error::InvalidProgram; + } + metadata.entries.clear(); + metadata.entries.reserve(num_entries); + + constexpr uint32_t kMaxTensorDimensions = 64; + for (uint32_t index = 0; index < num_entries; ++index) { + Entry entry; + uint32_t ndim = 0; + if (!reader.read_string(entry.fqn) || entry.fqn.empty() || + !reader.read_string(entry.storage_key) || entry.storage_key.empty() || + !reader.read_u64(entry.storage_nbytes) || + !reader.read_i32(entry.dtype) || !is_supported_dtype(entry.dtype) || + !reader.read_i32(entry.device_type) || + !is_supported_device_type(entry.device_type) || + !reader.read_i64(entry.storage_offset) || !reader.read_u32(ndim) || + ndim > kMaxTensorDimensions) { + return Error::InvalidProgram; + } + + entry.sizes.resize(ndim); + entry.strides.resize(ndim); + for (uint32_t dim = 0; dim < ndim; ++dim) { + if (!reader.read_i64(entry.sizes[dim]) || entry.sizes[dim] < 0) { + return Error::InvalidProgram; + } + } + for (uint32_t dim = 0; dim < ndim; ++dim) { + if (!reader.read_i64(entry.strides[dim]) || entry.strides[dim] < 0) { + return Error::InvalidProgram; + } + } + if (entry.storage_offset < 0) { + return Error::InvalidProgram; + } + metadata.entries.push_back(std::move(entry)); + } + + return reader.empty() ? Error::Ok : Error::InvalidProgram; +} + +Error CudaWeightCache::validate_view(const Entry& entry) { uint64_t item_size = 0; switch (static_cast(entry.dtype)) { case aoti::slim::c10::ScalarType::Byte: @@ -111,7 +279,7 @@ Error CudaWeightCache::validate_view(const CudaFqnWeightEntry& entry) { Error CudaWeightCache::acquire_storage( const NamedDataMap* named_data_map, - const CudaFqnWeightEntry& entry, + const Entry& entry, uintptr_t logical_scope, int device_index, std::shared_ptr& storage, @@ -238,7 +406,7 @@ Error CudaWeightCache::acquire_storage( Error CudaWeightCache::load( CudaDelegateHandle* handle, const NamedDataMap* named_data_map, - const CudaFqnWeightManifest& manifest) const { + const Metadata& metadata) const { ET_CHECK_OR_RETURN_ERROR( named_data_map != nullptr, InvalidArgument, @@ -308,12 +476,12 @@ Error CudaWeightCache::load( key_scopes.emplace(key.get(), reinterpret_cast(key.get())); } std::vector pairs; - pairs.reserve(manifest.entries.size()); + pairs.reserve(metadata.entries.size()); std::unordered_set bound_fqns; size_t reused_storages = 0; - handle->fqn_weight_tensors.reserve(manifest.entries.size()); + handle->fqn_weight_tensors.reserve(metadata.entries.size()); - for (const CudaFqnWeightEntry& entry : manifest.entries) { + for (const Entry& entry : metadata.entries) { ET_CHECK_OK_OR_RETURN_ERROR( validate_view(entry), "Invalid CUDA FQN view '%s'", entry.fqn.c_str()); auto internal_names = fqn_to_internal_names.find(entry.fqn); @@ -328,14 +496,14 @@ Error CudaWeightCache::load( aoti_dtype->second == entry.dtype, InvalidProgram, "CUDA FQN weight '%s' dtype does not match its AOTI library " - "(manifest=%d, AOTI=%d)", + "(serialized=%d, AOTI=%d)", entry.fqn.c_str(), entry.dtype, aoti_dtype == fqn_to_aoti_dtype.end() ? -1 : aoti_dtype->second); ET_CHECK_OR_RETURN_ERROR( bound_fqns.emplace(entry.fqn).second, InvalidProgram, - "CUDA FQN weight '%s' appears more than once in its manifest", + "CUDA FQN weight '%s' appears more than once in serialized metadata", entry.fqn.c_str()); std::shared_ptr storage; @@ -390,7 +558,7 @@ Error CudaWeightCache::load( ET_LOG( Info, "Loaded %zu CUDA FQN weights (%zu reused across methods)", - manifest.entries.size(), + metadata.entries.size(), reused_storages); return Error::Ok; } diff --git a/backends/cuda/runtime/cuda_weight_cache.h b/backends/cuda/runtime/cuda_weight_cache.h index f34cf41e1d9..eb58ebc8fb1 100644 --- a/backends/cuda/runtime/cuda_weight_cache.h +++ b/backends/cuda/runtime/cuda_weight_cache.h @@ -8,14 +8,15 @@ #pragma once +#include #include #include #include #include #include +#include #include -#include #include #include @@ -23,17 +24,41 @@ namespace executorch::backends::cuda { class CudaWeightCache final { public: + static constexpr char kFormatMagic[] = "ETCUDAFQN3"; + static constexpr size_t kFormatMagicSize = sizeof(kFormatMagic) - 1; + + struct Entry { + std::string fqn; + std::string storage_key; + uint64_t storage_nbytes{0}; + int32_t dtype{0}; + int32_t device_type{0}; + int64_t storage_offset{0}; + std::vector sizes; + std::vector strides; + }; + + struct Metadata { + std::string so_blob_key; + std::vector entries; + }; + + static bool is_serialized(const void* data, size_t size); + + static runtime::Error + parse(const void* data, size_t size, Metadata& metadata); + runtime::Error load( CudaDelegateHandle* handle, const runtime::NamedDataMap* named_data_map, - const CudaFqnWeightManifest& manifest) const; + const Metadata& metadata) const; private: - static runtime::Error validate_view(const CudaFqnWeightEntry& entry); + static runtime::Error validate_view(const Entry& entry); runtime::Error acquire_storage( const runtime::NamedDataMap* named_data_map, - const CudaFqnWeightEntry& entry, + const Entry& entry, uintptr_t logical_scope, int device_index, std::shared_ptr& storage, diff --git a/backends/cuda/runtime/cuda_weight_manifest.h b/backends/cuda/runtime/cuda_weight_manifest.h deleted file mode 100644 index 10f5fcb5734..00000000000 --- a/backends/cuda/runtime/cuda_weight_manifest.h +++ /dev/null @@ -1,211 +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. - */ - -#pragma once - -#include -#include -#include -#include - -#include - -namespace executorch::backends::cuda { - -constexpr char kCudaFqnWeightsMagic[] = "ETCUDAFQN3"; -constexpr size_t kCudaFqnWeightsMagicSize = sizeof(kCudaFqnWeightsMagic) - 1; - -struct CudaFqnWeightEntry { - std::string fqn; - std::string storage_key; - uint64_t storage_nbytes{0}; - int32_t dtype{0}; - int32_t device_type{0}; - int64_t storage_offset{0}; - std::vector sizes; - std::vector strides; -}; - -struct CudaFqnWeightManifest { - std::string so_blob_key; - std::vector entries; -}; - -inline bool is_supported_cuda_fqn_dtype(int32_t dtype) { - // Values match c10::ScalarType and the slim AOTI runtime. - switch (dtype) { - case 0: // Byte - case 1: // Char - case 2: // Short - case 3: // Int - case 4: // Long - case 5: // Half - case 6: // Float - case 11: // Bool - case 15: // BFloat16 - return true; - default: - return false; - } -} - -inline bool is_supported_cuda_fqn_device_type(int32_t device_type) { - // Values match c10::DeviceType and the slim AOTI runtime. - return device_type == 0 || device_type == 1; // CPU or CUDA -} - -inline bool is_cuda_fqn_weight_manifest(const void* data, size_t size) { - return data != nullptr && size >= kCudaFqnWeightsMagicSize && - std::memcmp(data, kCudaFqnWeightsMagic, kCudaFqnWeightsMagicSize) == 0; -} - -namespace detail { - -class CudaWeightManifestReader final { - public: - CudaWeightManifestReader(const void* data, size_t size) - : cursor_(static_cast(data)), end_(cursor_ + size) {} - - bool skip(size_t size) { - if (remaining() < size) { - return false; - } - cursor_ += size; - return true; - } - - bool read_u32(uint32_t& value) { - uint64_t wide = 0; - if (!read_unsigned(wide, 4)) { - return false; - } - value = static_cast(wide); - return true; - } - - bool read_i32(int32_t& value) { - uint32_t raw = 0; - if (!read_u32(raw)) { - return false; - } - std::memcpy(&value, &raw, sizeof(value)); - return true; - } - - bool read_u64(uint64_t& value) { - return read_unsigned(value, 8); - } - - bool read_i64(int64_t& value) { - uint64_t raw = 0; - if (!read_u64(raw)) { - return false; - } - std::memcpy(&value, &raw, sizeof(value)); - return true; - } - - bool read_string(std::string& value) { - uint32_t size = 0; - if (!read_u32(size) || remaining() < size) { - return false; - } - value.assign(reinterpret_cast(cursor_), size); - cursor_ += size; - return true; - } - - bool empty() const { - return cursor_ == end_; - } - - private: - size_t remaining() const { - return static_cast(end_ - cursor_); - } - - bool read_unsigned(uint64_t& value, size_t width) { - if (remaining() < width) { - return false; - } - value = 0; - for (size_t index = 0; index < width; ++index) { - value |= static_cast(cursor_[index]) << (index * 8); - } - cursor_ += width; - return true; - } - - const uint8_t* cursor_; - const uint8_t* end_; -}; - -} // namespace detail - -inline executorch::runtime::Error parse_cuda_fqn_weight_manifest( - const void* data, - size_t size, - CudaFqnWeightManifest& manifest) { - using executorch::runtime::Error; - if (!is_cuda_fqn_weight_manifest(data, size)) { - return Error::InvalidProgram; - } - - detail::CudaWeightManifestReader reader(data, size); - if (!reader.skip(kCudaFqnWeightsMagicSize) || - !reader.read_string(manifest.so_blob_key) || - manifest.so_blob_key.empty()) { - return Error::InvalidProgram; - } - - uint32_t num_entries = 0; - constexpr uint32_t kMaxManifestEntries = 1U << 20; - if (!reader.read_u32(num_entries) || num_entries > kMaxManifestEntries) { - return Error::InvalidProgram; - } - manifest.entries.clear(); - manifest.entries.reserve(num_entries); - - constexpr uint32_t kMaxTensorDimensions = 64; - for (uint32_t index = 0; index < num_entries; ++index) { - CudaFqnWeightEntry entry; - uint32_t ndim = 0; - if (!reader.read_string(entry.fqn) || entry.fqn.empty() || - !reader.read_string(entry.storage_key) || entry.storage_key.empty() || - !reader.read_u64(entry.storage_nbytes) || - !reader.read_i32(entry.dtype) || - !is_supported_cuda_fqn_dtype(entry.dtype) || - !reader.read_i32(entry.device_type) || - !is_supported_cuda_fqn_device_type(entry.device_type) || - !reader.read_i64(entry.storage_offset) || !reader.read_u32(ndim) || - ndim > kMaxTensorDimensions) { - return Error::InvalidProgram; - } - - entry.sizes.resize(ndim); - entry.strides.resize(ndim); - for (uint32_t dim = 0; dim < ndim; ++dim) { - if (!reader.read_i64(entry.sizes[dim]) || entry.sizes[dim] < 0) { - return Error::InvalidProgram; - } - } - for (uint32_t dim = 0; dim < ndim; ++dim) { - if (!reader.read_i64(entry.strides[dim]) || entry.strides[dim] < 0) { - return Error::InvalidProgram; - } - } - if (entry.storage_offset < 0) { - return Error::InvalidProgram; - } - manifest.entries.push_back(std::move(entry)); - } - - return reader.empty() ? Error::Ok : Error::InvalidProgram; -} - -} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 228a8470908..e8372a33462 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -130,7 +130,6 @@ def define_common_targets(is_fbcode = False): "cuda_delegate_handle.h", "cuda_mutable_state.h", "cuda_weight_cache.h", - "cuda_weight_manifest.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) link_whole = True, @@ -184,9 +183,9 @@ def define_common_targets(is_fbcode = False): ) cpp_unittest( - name = "test_cuda_weight_manifest", + name = "test_cuda_weight_cache", srcs = [ - "test/test_cuda_weight_manifest.cpp", + "test/test_cuda_weight_cache.cpp", ], deps = [ ":cuda_backend", diff --git a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp similarity index 58% rename from backends/cuda/runtime/test/test_cuda_weight_manifest.cpp rename to backends/cuda/runtime/test/test_cuda_weight_cache.cpp index 0ed3d383abb..d16e7143058 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_manifest.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include @@ -36,12 +36,13 @@ void append_string(std::vector& output, const std::string& value) { output.insert(output.end(), value.begin(), value.end()); } -std::vector valid_manifest( +std::vector serialized_metadata( uint32_t dtype = 6, uint32_t device_type = 1) { std::vector output( - cuda::kCudaFqnWeightsMagic, - cuda::kCudaFqnWeightsMagic + cuda::kCudaFqnWeightsMagicSize); + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); append_string(output, "so-key"); append_u32(output, 1); // entries append_string(output, "model.weight"); @@ -60,21 +61,21 @@ std::vector valid_manifest( } // namespace -TEST(CudaWeightManifestTest, LegacyPayloadIsNotMisdetected) { +TEST(CudaWeightCacheTest, LegacyPayloadIsNotMisdetected) { const std::string legacy = "so-key\nweights-key"; - EXPECT_FALSE(cuda::is_cuda_fqn_weight_manifest(legacy.data(), legacy.size())); + EXPECT_FALSE( + cuda::CudaWeightCache::is_serialized(legacy.data(), legacy.size())); } -TEST(CudaWeightManifestTest, ParsesVersionedManifest) { - const std::vector bytes = valid_manifest(); - cuda::CudaFqnWeightManifest manifest; +TEST(CudaWeightCacheTest, ParsesSerializedMetadata) { + const std::vector bytes = serialized_metadata(); + cuda::CudaWeightCache::Metadata metadata; ASSERT_EQ( - cuda::parse_cuda_fqn_weight_manifest( - bytes.data(), bytes.size(), manifest), + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), Error::Ok); - ASSERT_EQ(manifest.so_blob_key, "so-key"); - ASSERT_EQ(manifest.entries.size(), 1u); - const auto& entry = manifest.entries[0]; + ASSERT_EQ(metadata.so_blob_key, "so-key"); + ASSERT_EQ(metadata.entries.size(), 1u); + const auto& entry = metadata.entries[0]; EXPECT_EQ(entry.fqn, "model.weight"); EXPECT_EQ(entry.storage_key, "storage-key"); EXPECT_EQ(entry.storage_nbytes, 24u); @@ -84,36 +85,32 @@ TEST(CudaWeightManifestTest, ParsesVersionedManifest) { EXPECT_EQ(entry.strides, (std::vector{3, 1})); } -TEST(CudaWeightManifestTest, RejectsTruncationAndTrailingData) { - std::vector bytes = valid_manifest(); - cuda::CudaFqnWeightManifest manifest; +TEST(CudaWeightCacheTest, RejectsTruncationAndTrailingData) { + std::vector bytes = serialized_metadata(); + cuda::CudaWeightCache::Metadata metadata; ASSERT_GT(bytes.size(), 1u); EXPECT_EQ( - cuda::parse_cuda_fqn_weight_manifest( - bytes.data(), bytes.size() - 1, manifest), + cuda::CudaWeightCache::parse(bytes.data(), bytes.size() - 1, metadata), Error::InvalidProgram); bytes.push_back(0); EXPECT_EQ( - cuda::parse_cuda_fqn_weight_manifest( - bytes.data(), bytes.size(), manifest), + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), Error::InvalidProgram); } -TEST(CudaWeightManifestTest, RejectsUnsupportedDtype) { +TEST(CudaWeightCacheTest, RejectsUnsupportedDtype) { const std::vector bytes = - valid_manifest(7); // Double is unsupported. - cuda::CudaFqnWeightManifest manifest; + serialized_metadata(7); // Double is unsupported. + cuda::CudaWeightCache::Metadata metadata; EXPECT_EQ( - cuda::parse_cuda_fqn_weight_manifest( - bytes.data(), bytes.size(), manifest), + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), Error::InvalidProgram); } -TEST(CudaWeightManifestTest, RejectsUnsupportedDeviceType) { - const std::vector bytes = valid_manifest(6, 2); - cuda::CudaFqnWeightManifest manifest; +TEST(CudaWeightCacheTest, RejectsUnsupportedDeviceType) { + const std::vector bytes = serialized_metadata(6, 2); + cuda::CudaWeightCache::Metadata metadata; EXPECT_EQ( - cuda::parse_cuda_fqn_weight_manifest( - bytes.data(), bytes.size(), manifest), + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), Error::InvalidProgram); } diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 51c7be8d381..08e54eda63b 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -21,9 +21,9 @@ from executorch.backends.cuda.cuda_weight_collector import ( AOTI_DEVICE_TYPE_CPU, AOTI_DEVICE_TYPE_CUDA, - CUDA_FQN_WEIGHTS_MAGIC, + CUDA_WEIGHT_CACHE_MAGIC, CudaWeightCollector, - encode_cuda_weight_manifest, + encode_cuda_weight_metadata, ) from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStore @@ -95,10 +95,10 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: artifact.storages[artifact.entries[1].storage_key].to_bytes(), ) - manifest = encode_cuda_weight_manifest("so-key", artifact.entries) - self.assertTrue(manifest.startswith(CUDA_FQN_WEIGHTS_MAGIC)) - self.assertIn(b"first", manifest) - self.assertIn(b"second", manifest) + metadata = encode_cuda_weight_metadata("so-key", artifact.entries) + self.assertTrue(metadata.startswith(CUDA_WEIGHT_CACHE_MAGIC)) + self.assertIn(b"first", metadata) + self.assertIn(b"second", metadata) for storage in artifact.storages.values(): storage.close() From 29c08b539b702f5c5c3acf304904cb6194f39f34 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 11:48:41 -0700 Subject: [PATCH 08/12] cuda: simplify weight collector finalization Split named-data merging and FQN registration into focused collector helpers to satisfy complexity lint without changing serialization behavior.\n\nGenerated with Codex. --- backends/cuda/cuda_weight_collector.py | 88 ++++++++++++++++---------- 1 file changed, 53 insertions(+), 35 deletions(-) diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 4a6b315a5b5..52b3ce6be46 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -251,25 +251,12 @@ def materialize( trim_host_memory() return CudaWeightArtifact(entries=entries, storages=storages) - def add_preprocess_result( + def _merge_aoti_data( self, - result: PreprocessResult, - artifact: CudaWeightArtifact, - device_name: str, + parent_store: Any, + compatibility_blob_key: Optional[str], + keep_compatibility_blob: bool, ) -> None: - if result.data_store_output is None: - raise RuntimeError("CUDA AOTI preprocess returned no named data") - try: - parent_keys = result.processed_bytes.decode("utf-8").splitlines() - except UnicodeDecodeError as error: - raise RuntimeError("Malformed CUDA AOTI named-data payload") from error - if not parent_keys or not parent_keys[0]: - raise RuntimeError("CUDA AOTI payload is missing its shared-object key") - so_blob_key = parent_keys[0] - compatibility_blob_key = parent_keys[1] if len(parent_keys) > 1 else None - - parent_store = result.data_store_output - keep_compatibility_blob = not artifact.storages for key, entry in parent_store.pte_data.items(): if key != compatibility_blob_key or keep_compatibility_blob: self._store.add_named_data( @@ -289,26 +276,57 @@ def add_preprocess_result( tensor_layout=entry.tensor_layout, ) + def _add_weight( + self, + entry: CudaWeightEntry, + data: FileBackedData, + external_tag: str, + ) -> None: + previous = self._entries.get(entry.storage_key) + try: + if previous is not None and previous != entry: + raise ValueError( + f"Duplicate key {entry.storage_key} with different tensor " + "metadata." + ) + self._store.add_named_data( + entry.storage_key, + data, + alignment=1, + external_tag=external_tag, + ) + finally: + if previous is not None: + data.close() + self._entries.setdefault(entry.storage_key, entry) + + def add_preprocess_result( + self, + result: PreprocessResult, + artifact: CudaWeightArtifact, + device_name: str, + ) -> None: + if result.data_store_output is None: + raise RuntimeError("CUDA AOTI preprocess returned no named data") + try: + parent_keys = result.processed_bytes.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise RuntimeError("Malformed CUDA AOTI named-data payload") from error + if not parent_keys or not parent_keys[0]: + raise RuntimeError("CUDA AOTI payload is missing its shared-object key") + so_blob_key = parent_keys[0] + compatibility_blob_key = parent_keys[1] if len(parent_keys) > 1 else None + + parent_store = result.data_store_output + self._merge_aoti_data( + parent_store, + compatibility_blob_key, + keep_compatibility_blob=not artifact.storages, + ) + external_tag = f"aoti_{device_name}_blob" for entry in artifact.entries: - data = artifact.storages[entry.storage_key] - previous = self._entries.get(entry.storage_key) - try: - if previous is not None and previous != entry: - raise ValueError( - f"Duplicate key {entry.storage_key} with different tensor " - "metadata." - ) - self._store.add_named_data( - entry.storage_key, - data, - alignment=1, - external_tag=external_tag, - ) - finally: - if previous is not None: - data.close() - self._entries.setdefault(entry.storage_key, entry) + self._add_weight(entry, artifact.storages[entry.storage_key], external_tag) result.processed_bytes = encode_cuda_weight_metadata( so_blob_key, artifact.entries From 29939a0ee8bc52d82de9ebe7cdf63a45b922c9e2 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 11:57:20 -0700 Subject: [PATCH 09/12] cuda: format backend source list Generated with Codex. --- backends/cuda/CMakeLists.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 4d3a26582e5..c3e5d17809e 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -337,9 +337,9 @@ install( ) # CUDA backend implementation -set(_aoti_cuda_backend_sources runtime/cuda_backend.cpp - runtime/cuda_mutable_state.cpp - runtime/cuda_weight_cache.cpp +set(_aoti_cuda_backend_sources + runtime/cuda_backend.cpp runtime/cuda_mutable_state.cpp + runtime/cuda_weight_cache.cpp ) if(_cuda_is_msvc_toolchain) # MSVC links aoti_cuda_backend into portable_lib without relying on C++ From 905edadac4a8f468024c3bc896fbb28039a769f8 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 12:37:06 -0700 Subject: [PATCH 10/12] cuda: run weight cache test in CI Generated with Codex. --- .github/workflows/cuda.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 25684f8b32e..ede0cd591cd 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -273,9 +273,10 @@ jobs: export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON - cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state -j$(nproc) + cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache -j$(nproc) ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_mutable_state --output-on-failure -V + ctest --test-dir cmake-out -R test_cuda_weight_cache --output-on-failure -V test-model-cuda-e2e: name: test-model-cuda-e2e-${{ matrix.model.name }}-${{ matrix.quant }} From 915db8133f9aeebcb2e08c8e996800c92175886e Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 13:01:35 -0700 Subject: [PATCH 11/12] cuda: allow per-method views of shared weights Generated with Codex. --- backends/cuda/cuda_weight_collector.py | 17 +++++-------- backends/cuda/tests/test_cuda_partitioner.py | 26 ++++++++++++++------ 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 52b3ce6be46..4ebae200302 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -141,7 +141,6 @@ class CudaWeightCollector: def __init__(self) -> None: self._store = NamedDataStore() - self._entries: Dict[str, CudaWeightEntry] = {} self._results: List[PreprocessResult] = [] @contextlib.contextmanager @@ -282,23 +281,19 @@ def _add_weight( data: FileBackedData, external_tag: str, ) -> None: - previous = self._entries.get(entry.storage_key) + is_duplicate = entry.storage_key in self._store.key_to_buffer_idx try: - if previous is not None and previous != entry: - raise ValueError( - f"Duplicate key {entry.storage_key} with different tensor " - "metadata." - ) self._store.add_named_data( entry.storage_key, data, alignment=1, external_tag=external_tag, ) - finally: - if previous is not None: - data.close() - self._entries.setdefault(entry.storage_key, entry) + except Exception: + data.close() + raise + if is_duplicate: + data.close() def add_preprocess_result( self, diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 08e54eda63b..4cb9aa6be8a 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -192,7 +192,7 @@ def test_same_fqn_with_different_data_is_rejected_when_merged(self) -> None: collector.add_preprocess_result( self._parent_result("first_so"), first_artifact, "cuda" ) - with self.assertRaises(ValueError): + with self.assertRaisesRegex(ValueError, "different data"): collector.add_preprocess_result( self._parent_result("second_so"), second_artifact, "cuda" ) @@ -201,7 +201,7 @@ def test_same_fqn_with_different_data_is_rejected_when_merged(self) -> None: for storage in artifact.storages.values(): storage.close() - def test_same_fqn_with_different_metadata_is_rejected(self) -> None: + def test_same_fqn_with_different_metadata_shares_storage(self) -> None: tensor = torch.arange(4) view = tensor.reshape(2, 2) with ( @@ -215,13 +215,23 @@ def test_same_fqn_with_different_metadata_is_rejected(self) -> None: Weights({"weight": (tensor, TensorProperties(view))}), second_dir ) collector = CudaWeightCollector() - collector.add_preprocess_result( - self._parent_result("first_so"), first_artifact, "cuda" + first_result = self._parent_result("first_so") + second_result = self._parent_result("second_so") + collector.add_preprocess_result(first_result, first_artifact, "cuda") + collector.add_preprocess_result(second_result, second_artifact, "cuda") + collector.finish() + + self.assertIs( + first_result.data_store_output, second_result.data_store_output + ) + self.assertNotEqual( + first_result.processed_bytes, second_result.processed_bytes + ) + weight_key = first_artifact.entries[0].storage_key + self.assertIn( + weight_key, + first_result.data_store_output.external_data["aoti_cuda_blob"], ) - with self.assertRaisesRegex(ValueError, "different tensor metadata"): - collector.add_preprocess_result( - self._parent_result("second_so"), second_artifact, "cuda" - ) for artifact in (first_artifact, second_artifact): for storage in artifact.storages.values(): From fa7c98aa0a623546a8ce23eb4ba28f094a4718b4 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Tue, 25 Aug 2026 13:52:26 -0700 Subject: [PATCH 12/12] cuda: scope AOTI tensor constants per library Generated with Codex. --- backends/cuda/cuda_weight_collector.py | 29 +++++++++-- backends/cuda/tests/test_cuda_partitioner.py | 51 ++++++++++++++++++++ 2 files changed, 76 insertions(+), 4 deletions(-) diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 4ebae200302..ac47f7c7e2c 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -12,7 +12,7 @@ import struct import tempfile import threading -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple import torch @@ -93,16 +93,27 @@ def write_chunk(output, chunk) -> None: return digest.digest() -def _storage_key(fqn: str, device_type: int) -> str: +def _storage_key( + fqn: str, device_type: int, aoti_library_key: Optional[str] = None +) -> str: if device_type == AOTI_DEVICE_TYPE_CPU: device = "cpu" elif device_type == AOTI_DEVICE_TYPE_CUDA: device = "cuda" else: raise RuntimeError(f"Unsupported AOTI device type: {device_type}") + if aoti_library_key is not None: + fqn = f"{aoti_library_key}:{fqn}" return f"cuda_fqn_weight:{device}:{fqn}" +def _is_aoti_library_local_fqn(fqn: str) -> bool: + # PyTorch assigns this prefix to TensorConstant entries that do not have a + # model-level FQN. The numbering restarts in every independently compiled + # AOTI library, so the library key is part of their global identity. + return fqn.startswith("_tensor_constant") + + def encode_cuda_weight_metadata( so_blob_key: str, entries: List[CudaWeightEntry] ) -> bytes: @@ -320,11 +331,21 @@ def add_preprocess_result( ) external_tag = f"aoti_{device_name}_blob" + serialized_entries = [] for entry in artifact.entries: - self._add_weight(entry, artifact.storages[entry.storage_key], external_tag) + data = artifact.storages[entry.storage_key] + if _is_aoti_library_local_fqn(entry.fqn): + entry = replace( + entry, + storage_key=_storage_key( + entry.fqn, entry.device_type, aoti_library_key=so_blob_key + ), + ) + self._add_weight(entry, data, external_tag) + serialized_entries.append(entry) result.processed_bytes = encode_cuda_weight_metadata( - so_blob_key, artifact.entries + so_blob_key, serialized_entries ) self._results.append(result) diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 4cb9aa6be8a..153828b2cb0 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -237,6 +237,57 @@ def test_same_fqn_with_different_metadata_shares_storage(self) -> None: for storage in artifact.storages.values(): storage.close() + def test_aoti_local_fqns_are_scoped_by_library(self) -> None: + first = torch.zeros(4) + second = torch.ones(5) + + with ( + tempfile.TemporaryDirectory() as first_dir, + tempfile.TemporaryDirectory() as second_dir, + ): + first_artifact = self._materialize( + Weights( + { + "_tensor_constant0": ( + first, + TensorProperties(first), + ) + } + ), + first_dir, + ) + second_artifact = self._materialize( + Weights( + { + "_tensor_constant0": ( + second, + TensorProperties(second), + ) + } + ), + second_dir, + ) + first_result = self._parent_result("first_so") + second_result = self._parent_result("second_so") + collector = CudaWeightCollector() + collector.add_preprocess_result(first_result, first_artifact, "cuda") + collector.add_preprocess_result(second_result, second_artifact, "cuda") + collector.finish() + + external_data = first_result.data_store_output.external_data[ + "aoti_cuda_blob" + ] + first_key = "cuda_fqn_weight:cpu:first_so:_tensor_constant0" + second_key = "cuda_fqn_weight:cpu:second_so:_tensor_constant0" + self.assertIn(first_key, external_data) + self.assertIn(second_key, external_data) + self.assertIn(first_key.encode(), first_result.processed_bytes) + self.assertIn(second_key.encode(), second_result.processed_bytes) + + for artifact in (first_artifact, second_artifact): + for storage in artifact.storages.values(): + storage.close() + def test_methods_share_one_collected_named_data_store(self) -> None: tensor = torch.arange(4) weights = Weights({"weight": (tensor, TensorProperties(tensor))})