[cuda backend] share AOTI weights by FQN across methods

.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 }}

backends/cuda/BUCK

@@ -93,6 +93,7 @@ fbcode_target(
  name = "cuda_backend",
  srcs = [
  "cuda_backend.py",
+ "cuda_weight_collector.py",
  ],
  visibility = ["PUBLIC"],
  deps = [

backends/cuda/CMakeLists.txt

@@ -337,8 +337,9 @@ install(
 )
 
 # CUDA backend implementation
-set(_aoti_cuda_backend_sources runtime/cuda_backend.cpp
- runtime/cuda_mutable_state.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++
@@ -448,4 +449,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_cache SOURCES runtime/test/test_cuda_weight_cache.cpp
+ EXTRA_LIBS aoti_cuda_backend
+ )
+ target_compile_definitions(test_cuda_weight_cache PRIVATE CUDA_AVAILABLE=1)
 endif()

backends/cuda/cuda_backend.py

@@ -7,10 +7,7 @@
 
 import contextlib
 import copy
-import ctypes
 import functools
-import gc
-import hashlib
 import logging
 import os
 import shutil
@@ -21,6 +18,12 @@
 
 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,
 )
@@ -32,7 +35,7 @@
 )
 from executorch.exir._serialize._cord import FileBackedData
 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 torch._inductor.decomposition import conv1d_to_conv2d
 from torch.nn.attention import SDPBackend
@@ -66,12 +69,18 @@ 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 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
+ if tensor.device.type == "cpu":
+ return AOTI_DEVICE_TYPE_CPU
+ raise RuntimeError(
+ f"Unsupported AOTI constant device for CUDA export: {tensor.device}"
+ )
 
 
 @contextlib.contextmanager
@@ -195,10 +204,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 +289,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,47 +412,6 @@ 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)
- chunk_size = 8 * 1024 * 1024
- digest = hashlib.sha256()
-
- def write_chunk(output, chunk) -> None:
- digest.update(chunk)
- output.write(chunk)
-
- with open(blob_path, "wb") as output:
- for tensor in tensors:
- 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()
- return digest.digest()
-
-
 @final
 @experimental(
  "This API and all of cuda backend related functionality are experimental."
@@ -454,11 +423,6 @@ class CudaBackend(AotiBackend, BackendDetails):
  using the Executorch runtime.
  """
 
- # AOTI calls materialize_weights_blob immediately before load_weights_blob
- # for a given output path. A new materialization overwrites any digest left
- # behind by an export that aborted before the consumer ran.
- _materialized_blob_hashes: Dict[str, bytes] = {}
-
  @classmethod
  def get_device_name(cls) -> str:
  return "cuda"
@@ -557,33 +521,83 @@ def _setup_cuda_environment_for_fatbin() -> 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_with_weight_collector(
+ cls,
+ edge_program: Any,
+ compile_specs: List[CompileSpec],
+ collector: CudaWeightCollector,
+ ) -> PreprocessResult:
+ with collector.capture(_aoti_device_type_for_weight) as capture:
+ result = super().preprocess(edge_program, compile_specs)
+ if capture.artifact is None:
+ raise RuntimeError("CUDA AOTI did not return a structured Weights output")
+ collector.add_preprocess_result(result, capture.artifact, cls.get_device_name())
+ return result
+
+ @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
+ )
+ 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]
  ) -> 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.
  """
  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()
+ blob_data = FileBackedData.move_from(blob_path)
+ return blob_data, blob_data.sha256().hex()
 
  @classmethod
  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 +621,19 @@ 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 = CudaWeightCollector.current_capture()
+ capture.artifact = capture.collector.materialize(
+ weights[0], os.path.dirname(blob_path), capture.device_type_for_weight
  )
 
- # 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
+
+ # 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 +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,
- # 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 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
  ),
@@ -877,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()
@@ -892,11 +910,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-FQN metadata payload. This is backend-wide rather than a
+ # model/export-script option.
+ return "pickle_weights"
 
  @classmethod
  def move_program_to_device(