Skip to content

Commit ce52563

Browse files
committed
[executorch][cuda] Lower off-graph KV storage through AOTI
Pull Request resolved: #22815 Lower the neutral `kvcache::update_and_attend` contract to runtime-bound private tensors, retain their logical metadata without serializing payloads, and allow the weight loader to defer those bindings. Add validation for manifests and storage omission. Authored with Codex. ghstack-source-id: 435427976 @exported-using-ghexport Differential Revision: [D119571985](https://our.internmc.facebook.com/intern/diff/D119571985/)
1 parent 3235f58 commit ce52563

8 files changed

Lines changed: 801 additions & 47 deletions

File tree

‎backends/cuda/BUCK‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,16 @@ fbcode_target(
7777
srcs = [
7878
"passes/__init__.py",
7979
"passes/move_cond_predicate_to_cpu.py",
80+
"passes/lower_offgraph_kv.py",
8081
"passes/replace_int64_floordiv.py",
8182
],
8283
visibility = [
8384
"//executorch/backends/cuda/...",
8485
],
8586
deps = [
87+
":triton_kernels",
8688
"//caffe2:torch",
89+
"//executorch/backends/transforms:utils",
8790
"//executorch/exir/dialects:lib",
8891
],
8992
)

‎backends/cuda/cuda_backend.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,11 @@
2424
CudaWeightCollector,
2525
trim_host_memory,
2626
)
27+
from executorch.backends.cuda.passes.lower_offgraph_kv import (
28+
LowerOffGraphKVPass,
29+
OFFGRAPH_KV_COMPILE_SPEC,
30+
parse_offgraph_kv_manifest,
31+
)
2732
from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
2833
MoveCondPredicateToCpuPass,
2934
)
@@ -777,6 +782,11 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
777782
)
778783
triton_kernel_mode = mode
779784
passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
785+
for spec in compile_specs:
786+
if spec.key == OFFGRAPH_KV_COMPILE_SPEC:
787+
passes.append(
788+
LowerOffGraphKVPass(parse_offgraph_kv_manifest(spec.value))
789+
)
780790
if triton_kernel_mode == "ON":
781791
passes.append(ReplaceEdgeOpWithTritonOpPass())
782792
return passes

‎backends/cuda/cuda_weight_collector.py‎

Lines changed: 65 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
1717

1818
import torch
19+
from executorch.backends.cuda.passes.lower_offgraph_kv import OFFGRAPH_KV_FQN_PREFIX
1920
from executorch.exir._serialize._cord import FileBackedData
2021
from executorch.exir._serialize._named_data_store import NamedDataStore
2122
from 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+
114151
def _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+
135176
def _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

Comments
 (0)