From b7d877f8494bb90c137317644dbc98bea17956a3 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 13 Aug 2026 11:30:04 -0700 Subject: [PATCH] cuda: share AOTI weights by FQN across methods --- backends/cuda/CMakeLists.txt | 6 + backends/cuda/cuda_backend.py | 399 ++++++++++++--- backends/cuda/runtime/cuda_backend.cpp | 479 +++++++++++++++++- backends/cuda/runtime/cuda_delegate_handle.h | 48 ++ 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, 1298 insertions(+), 111 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 06990692428..40b21764b41 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -260,4 +260,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 5b6ae5427d2..a3321288d37 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 @@ -28,9 +31,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 @@ -58,6 +63,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) @@ -118,10 +151,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) @@ -186,7 +220,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 prev_active = getattr(_CPU_CLONE_GUARD, "active", False) _CPU_CLONE_GUARD.active = True try: @@ -307,11 +341,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 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() @@ -319,33 +350,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 - 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( + 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( + 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() @@ -483,7 +760,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 @@ -502,13 +779,24 @@ def materialize_weights_blob( if isinstance(path, str) and path.endswith(".wrapper.so") ) 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) @@ -610,10 +898,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 ), @@ -766,11 +1052,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 8666fc9c098..c7d2acb4b2b 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -17,15 +17,18 @@ #include #include +#include #include #include #include #include +#include #include #include #include #include #include +#include #include // Include SlimTensor headers for CUDA backend @@ -45,6 +48,7 @@ #include #include #include +#include #include #include #include @@ -164,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); } @@ -176,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); \ @@ -225,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( @@ -316,10 +324,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()); @@ -393,6 +412,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 @@ -400,7 +427,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 { @@ -888,9 +915,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}; // --------------------------------------------------------------- @@ -1002,6 +1029,431 @@ 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< + uint32_t, + std::vector> + 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. // @@ -1237,6 +1689,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 ee360531c47..8bd057e1c54 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -17,6 +18,43 @@ 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 +186,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 +211,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 94fa08f1c4d..75e9572438d 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -127,6 +127,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, @@ -179,6 +180,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 b5be7a0df5c..c72249e9376 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)), + } + ) - 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, set()) + self.assertEqual(1, len(artifact.storages)) + self.assertEqual( + artifact.entries[0].storage_group, + artifact.entries[1].storage_group, ) - 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, 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)), + } + ) + + with tempfile.TemporaryDirectory() as directory: + artifact = _materialize_fqn_weights( + weights, directory, mutated_fqns={"first", "second"} + ) + 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_blob_stays_file_backed(self) -> None: data = b"cuda weights"