1616from typing import Any , Callable , Dict , Iterator , List , Optional , Tuple
1717
1818import torch
19+ from executorch .backends .cuda .passes .lower_offgraph_kv import OFFGRAPH_KV_FQN_PREFIX
1920from executorch .exir ._serialize ._cord import FileBackedData
2021from executorch .exir ._serialize ._named_data_store import NamedDataStore
2122from executorch .exir .backend .backend_details import PreprocessResult
@@ -111,6 +112,42 @@ def write_chunk(output, chunk) -> None:
111112 return digest .digest ()
112113
113114
115+ def _materialize_storage (tensor : torch .Tensor , directory : str ) -> FileBackedData :
116+ fd , storage_path = tempfile .mkstemp (
117+ prefix = ".cuda_weight_" , suffix = ".storage" , dir = directory
118+ )
119+ os .close (fd )
120+ try :
121+ digest = _write_tensor_storage (tensor , storage_path )
122+ return FileBackedData .move_from (storage_path , sha256 = digest )
123+ except Exception :
124+ with contextlib .suppress (OSError ):
125+ os .remove (storage_path )
126+ raise
127+
128+
129+ def _required_view_nbytes (
130+ fqn : str ,
131+ sizes : Tuple [int , ...],
132+ strides : Tuple [int , ...],
133+ storage_offset : int ,
134+ element_size : int ,
135+ ) -> int :
136+ if (
137+ len (sizes ) != len (strides )
138+ or storage_offset < 0
139+ or any (size < 0 for size in sizes )
140+ or any (stride < 0 for stride in strides )
141+ ):
142+ raise RuntimeError (f"AOTI view { fqn !r} has invalid tensor metadata" )
143+ if any (size == 0 for size in sizes ):
144+ return 0
145+ last_element = storage_offset + sum (
146+ stride * (size - 1 ) for size , stride in zip (sizes , strides )
147+ )
148+ return (last_element + 1 ) * element_size
149+
150+
114151def _storage_key (
115152 fqn : str , device_type : int , aoti_library_key : Optional [str ] = None
116153) -> str :
@@ -132,6 +169,10 @@ def _is_aoti_library_local_fqn(fqn: str) -> bool:
132169 return fqn .startswith ("_tensor_constant" )
133170
134171
172+ def _is_offgraph_kv_fqn (fqn : str ) -> bool :
173+ return fqn .startswith (OFFGRAPH_KV_FQN_PREFIX )
174+
175+
135176def _validate_cuda_aoti_variant (
136177 variant : CudaAotiVariant , has_fallback : bool , regular_sms : set [int ]
137178) -> None :
@@ -362,34 +403,23 @@ def materialize(
362403 storage_nbytes = storage .nbytes ()
363404 del storage
364405 device_type = device_type_for_weight (tensor )
406+ is_offgraph_kv = _is_offgraph_kv_fqn (fqn )
365407 expected_storage_nbytes = int (
366408 getattr (properties , "storage_size" , None ) or 0
367409 )
368- if storage_nbytes < expected_storage_nbytes :
410+ if not is_offgraph_kv and storage_nbytes < expected_storage_nbytes :
369411 raise RuntimeError (
370412 "AOTI cloned storage is smaller than its TensorProperties "
371413 f"({ storage_nbytes } < { expected_storage_nbytes } bytes)"
372414 )
373415
374- fd , storage_path = tempfile .mkstemp (
375- prefix = ".cuda_weight_" , suffix = ".storage" , dir = directory
376- )
377- os .close (fd )
378- try :
379- digest = _write_tensor_storage (tensor , storage_path )
380- data = FileBackedData .move_from (storage_path , sha256 = digest )
381- except Exception :
382- try :
383- os .remove (storage_path )
384- except OSError :
385- pass
386- raise
387-
388416 storage_key = _storage_key (fqn , device_type )
389- if storage_key in storages :
390- data .close ()
391- raise RuntimeError (f"Duplicate CUDA FQN weight key for { fqn !r} " )
392- storages [storage_key ] = data
417+ if not is_offgraph_kv :
418+ data = _materialize_storage (tensor , directory )
419+ if storage_key in storages :
420+ data .close ()
421+ raise RuntimeError (f"Duplicate CUDA FQN weight key for { fqn !r} " )
422+ storages [storage_key ] = data
393423
394424 sizes = tuple (
395425 int (size ) for size in getattr (properties , "shape" , tensor .shape )
@@ -398,24 +428,20 @@ def materialize(
398428 int (stride ) for stride in getattr (properties , "stride" , tensor .stride ())
399429 )
400430 storage_offset = int (getattr (properties , "offset" , tensor .storage_offset ()))
401- if (
402- len (sizes ) != len (strides )
403- or storage_offset < 0
404- or any (size < 0 for size in sizes )
405- or any (stride < 0 for stride in strides )
406- ):
407- raise RuntimeError (f"AOTI view { fqn !r} has invalid tensor metadata" )
408- required_nbytes = 0
409- if all (size != 0 for size in sizes ):
410- last_element = storage_offset + sum (
411- stride * (size - 1 ) for size , stride in zip (sizes , strides )
412- )
413- required_nbytes = (last_element + 1 ) * tensor .element_size ()
414- if required_nbytes > storage_nbytes :
431+ required_nbytes = _required_view_nbytes (
432+ fqn , sizes , strides , storage_offset , tensor .element_size ()
433+ )
434+ if not
435+ if not is_offgraph_kv and required_nbytes > storage_nbytes :
415436 raise RuntimeError (
416437 f"AOTI view { fqn !r} requires { required_nbytes } bytes from a "
417438 f"{ storage_nbytes } -byte cloned storage"
418439 )
440+ if is_offgraph_kv :
441+ # Preserve the AOTI view contract; the runtime supplies storage.
442+ storage_nbytes = max (
443+ storage_nbytes , expected_storage_nbytes , required_nbytes
444+ )
419445 entries .append (
420446 CudaWeightEntry (
421447 fqn = fqn ,
@@ -500,21 +526,24 @@ def add_preprocess_result(
500526 self ._merge_aoti_data (
501527 parent_store ,
502528 compatibility_blob_key ,
503- keep_compatibility_blob = not artifact .storages ,
529+ keep_compatibility_blob = not artifact .entries ,
504530 )
505531
506532 external_tag = f"aoti_{ device_name } _blob"
507533 serialized_entries = []
508534 for entry in artifact .entries :
509- data = artifact .storages [ entry .storage_key ]
535+ data = artifact .storages . get ( entry .storage_key )
510536 if _is_aoti_library_local_fqn (entry .fqn ):
511537 entry = replace (
512538 entry ,
513539 storage_key = _storage_key (
514540 entry .fqn , entry .device_type , aoti_library_key = so_blob_key
515541 ),
516542 )
517- self ._add_weight (entry , data , external_tag )
543+ if data is not None :
544+ self ._add_weight (entry , data , external_tag )
545+ elif not _is_offgraph_kv_fqn (entry .fqn ):
546+ raise RuntimeError (f"Missing CUDA FQN weight storage for { entry .fqn !r} " )
518547 serialized_entries .append (entry )
519548
520549 result .processed_bytes = encode_cuda_aoti_metadata (
0 commit comments