Accept AOTI's compact clones of view buffers in the CUDA weight collector - #23355
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23355
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 2 PendingAs of commit e1704e2 with merge base ecf5c39 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The truncation test must exercise the compact-clone path to protect its remaining bounds check against regressions.
Review effort: Balanced
Findings: 1
What changed in this PR
Updates the CUDA weight collector to accept AOTInductor’s compact clones of view buffers while retaining storage-bounds validation.
Changes:
- Detects matching clone layouts and uses the written tensor’s offset.
- Adds compact-clone and layout-fallback regression tests.
| File | Description |
|---|---|
| backends/cuda/tests/test_cuda_partitioner.py | Adds clone, truncation, and offset-fallback tests. |
| backends/cuda/cuda_weight_collector.py | Accepts compact clones while preserving view-bounds checks. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
dd3a0d7 to
0103560
Compare
0103560 to
51dcf12
Compare
51dcf12 to
d609fa9
Compare
| storages: Dict[str, FileBackedData] = {} | ||
|
|
||
| for fqn, (tensor, properties) in weights.items(): | ||
| is_offgraph_kv = _is_offgraph_kv_fqn(fqn) |
There was a problem hiding this comment.
Why would offgraph kv be in the weights fqn?
There was a problem hiding this comment.
That check isn't new in this change. It came with the off-graph KV lowering (#23292). That pass puts its cache buffers into the program as constants named __et_offgraph_kv_*, so they reach the weight collector together with the real weights. The collector skips them because their data is managed elsewhere. This change only moves that existing line up, so the new check can use it too.
There was a problem hiding this comment.
It is for removing kvcache buffer from weight: once we detected is_offgraph_kv we will remove the buffer from ptd to make runtime control it.
| getattr(properties, "storage_size", None) or 0 | ||
| ) | ||
| if not is_offgraph_kv and storage_nbytes < expected_storage_nbytes: | ||
| if ( |
There was a problem hiding this comment.
so basically iiuc if you have a fused weight aot like qkv aoti wants to split them into q k and v and there was some metadata getting trashed in this process that you fix? @shoumikhin
There was a problem hiding this comment.
Close, but it's about views rather than splitting. Take a buffer that is a view of a bigger tensor, like big[2:5]. AOTInductor copies it into a compact clone: 48 bytes, starting at offset 0. But the TensorProperties it reports still describe the original: 128 bytes at offset 8. The collector compared the two and rejected the export, even though the clone holds every byte the view needs. A fused QKV weight split into views is one way to get such a buffer, so yes, that's a real-world case.
The fix uses the clone's own size and offset when it's a compact clone with the same shape and strides. The bounds check still runs, so a truncated clone is still rejected.
| ) | ||
| required_nbytes = _required_view_nbytes( | ||
| fqn, sizes, strides, storage_offset, tensor.element_size() | ||
| ) |
There was a problem hiding this comment.
can we add a check here that the required_bytes is the same as the tensor.nbytes()?
There was a problem hiding this comment.
Added, slightly adjusted: a compact clone must now span exactly its storage. Checking against tensor.nbytes() would reject strided views. For example, a column slice like big[:, 1:3] of an 8x8 tensor has 64 bytes of elements, but its clone spans 232 bytes, because the clone keeps the gaps between rows. A new test covers a clone with extra bytes past its view, which is now rejected, and a strided view, which is still accepted.
d609fa9 to
d812bb0
Compare
d812bb0 to
5339a7a
Compare
5339a7a to
1542b00
Compare
…ctor
AOTInductor (AOTI) clones every buffer into compact storage that starts at
the view, but builds the buffer's TensorProperties from the original. When a
buffer is a view into a larger tensor, the CUDA weight collector judged the
clone by the original storage and rejected the export: a contiguous view
failed the storage size check ("cloned storage is smaller than its
TensorProperties"), and a strided view, whose TensorProperties has no storage
size, failed the bounds check ("requires 236 bytes from a 232-byte cloned
storage"). Yet the clone holds every byte the view needs.
When the value is such a clone (a different storage, with the same shape and
strides as its TensorProperties), skip the storage size check and take the
offset from the clone that is written. The clone must also span its storage
exactly. Any other value takes the existing path. The check that the view
fits inside the written bytes still runs, so a truncated clone is still
rejected.
1542b00 to
e1704e2
Compare
What is wrong today
When the CUDA backend compiles a model, AOTInductor hands back every constant as a pair: the tensor to save, and a
TensorPropertiesrecord that describes it (storage size, offset, shape, strides). The CUDA weight collector writes the tensor's storage to the.ptdfile and records that metadata so the runtime can rebuild the tensor.AOTInductor clones every buffer into compact storage first. The clone covers only the bytes the buffer needs and starts at offset 0. But
TensorPropertiesis built from the original buffer.That breaks when a buffer is a view into a larger tensor:
The clone is 48 bytes at offset 0, while
TensorPropertiessays 128 bytes at offset 8. The collector compares the two, and lowering withCudaPartitionerfails:A strided view fails in the bounds check instead, because its
TensorPropertieshas no storage size. A column slicebig[:, 1:3]of an 8x8 tensor fails with:In both cases the clone holds every byte the buffer needs. Only the metadata describes a different storage.
What this change does
When the tensor is AOTInductor's compact clone of the view (a different storage, with the same shape and strides as
TensorProperties), the collector skips the comparison with the original storage size and takes the offset from the clone it writes. Any other tensor goes through the existing path unchanged.The check that the view fits inside the bytes actually written still runs for every weight, so a truncated clone is still rejected. A compact clone is also checked to span its storage exactly, so a clone with bytes past its view is rejected too.
What was tested
test_cuda_partitioner.py: a clone ofbig[2:5]paired with the view'sTensorPropertiesis written as 48 bytes at offset 0; a clone of the stridedbig[:, 1:3]is written as its 232-byte storage, byte for byte; a clone with too little storage or with bytes past its view is rejected; a parameter that slices a fused tensor and keeps its storage still records its own offset and storage size; a value with the same shape but other strides is not taken for a clone; and a tensor with another layout keeps the offset itsTensorPropertiesrecords. The two clone tests fail before this change with the errors above.CudaPartitioner: before, the error above; after, it exports, and running the.ptewith the ExecuTorch runtime gives exactly the eager output.test_cuda_partitioner.pyandtest_cuda_weight_metadata.pypass with this change.