mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-28 18:56:51 +08:00
main
5
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 PyTorch MergeBot
|
21756fd7f3 |
Revert "[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)"
This reverts commit
|
||
|
Natalia Gimelshein
|
e506c0307c |
[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)
## Summary (human written only) This PR is another approach to handling materialization of lazy clones. Instead of previously used cudaMemcpy it 1) allocates the materialized copy storage on the original stream where _lazy_clone was called, to have the same storage behavior 2) runs the copy on allocation stream, we are guaranteed that the data is correct on allocations so no sync to the current stream is needed 3) syncs to the current stream so that it has correct data To enable "stealing" of the last remaining reference, this approach requires users to mark the last use of any of lazy-cloned siblings `torch._C._cow_record_done(tensor)`. If we determine that we have the last remaining reference of the lazy clone, and all previous uses have a recorded event, we are free to reuse the last remaining reference, after syncing to all the recorded events. This way, we are syncing to the earliest possible point on other streams after the data is no longer needed, so the stalls will be minimal. <!-- Add your own summary and explain the relevance of the agent details below. --> ## Agent details (if applicable) AI-generated description (Codex), covering both commits on this branch. CUDA COW materialization used a synchronous cudaMemcpy on the legacy default stream. That copy was unordered with respect to nonblocking operation streams and was absent from CUDA graph replay, allowing replay to read stale data. Enqueue materializing copies on the original allocations reuse stream. Allocate replacements on the original allocation's reuse stream and establish event dependencies in the direction of current op stream. The same event operations provide eager ordering and captured graph dependencies, without recordStream or host synchronization. Retain replaced COW references with their logical storage to protect reads submitted before materialization. Add `torch._C._cow_record_done(tensor)` so callers can explicitly complete those prior uses. The caller first joins all relevant uses, including aliases and earlier pointer accesses, onto the current stream. Completion records an event and releases retained references. An active COW reference remains alive until destruction or materialization; later data access cancels its completion. Each sibling has its own completion record, and renewed completion preserves its earlier event dependencies. When every departed sibling completed explicitly, the last sibling can steal the original allocation after waiting on the recorded events, without fencing unrelated later work on the allocation stream. Releasing retained references also enables subsequent lazy clones to share the new allocation. Callers that omit completion retain the conservative copying behavior. Capture boundaries remain explicit: completion and shared materialization must stay within the context's capture. Creating COW storage for an external graph input during capture is rejected. Private-pool outputs cloned after capture, unsupported sharing streams, and allocators without allocation metadata use eager copies. Unique materialization preserves the allocation's address. Compatibility impact is limited to explicit CUDA `_lazy_clone()` users: unsupported capture crossings now raise errors, and some lazy clones fall back to eager copies. CPU behavior is unchanged. Deferred materialization can introduce dependencies at the later operation, so this does not promise the exact stream behavior of an eager `clone()` at the original `_lazy_clone()` call. Validation: built on GB200 with CUDA 13.1. Passed 16 C++ COW tests, 43 native CUDA COW tests, 96 existing tensor lazy-clone tests, and two cudaMallocAsync fallback tests (157 passed total). Tests cover independent readers, copies and stealing on different streams, allocator reuse, completion renewal, chained clones, capture boundaries, repeated graph replay with changing inputs, and reduce-overhead compilation. Targeted `spin lint` and pre-commit `lintrunner -a` passed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200800 Approved by: https://github.com/ezyang |
||
|
PyTorch MergeBot
|
2f8c573ca4 |
Revert "[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)"
This reverts commit
|
||
|
Natalia Gimelshein
|
20f82b619b |
[COW] track lifetimes of cow tensors via explicit completion events, harden materialization (#200800)
## Summary (human written only) This PR is another approach to handling materialization of lazy clones. Instead of previously used cudaMemcpy it 1) allocates the materialized copy storage on the original stream where _lazy_clone was called, to have the same storage behavior 2) runs the copy on allocation stream, we are guaranteed that the data is correct on allocations so no sync to the current stream is needed 3) syncs to the current stream so that it has correct data To enable "stealing" of the last remaining reference, this approach requires users to mark the last use of any of lazy-cloned siblings `torch._C._cow_record_done(tensor)`. If we determine that we have the last remaining reference of the lazy clone, and all previous uses have a recorded event, we are free to reuse the last remaining reference, after syncing to all the recorded events. This way, we are syncing to the earliest possible point on other streams after the data is no longer needed, so the stalls will be minimal. <!-- Add your own summary and explain the relevance of the agent details below. --> ## Agent details (if applicable) AI-generated description (Codex), covering both commits on this branch. CUDA COW materialization used a synchronous cudaMemcpy on the legacy default stream. That copy was unordered with respect to nonblocking operation streams and was absent from CUDA graph replay, allowing replay to read stale data. Enqueue materializing copies on the original allocations reuse stream. Allocate replacements on the original allocation's reuse stream and establish event dependencies in the direction of current op stream. The same event operations provide eager ordering and captured graph dependencies, without recordStream or host synchronization. Retain replaced COW references with their logical storage to protect reads submitted before materialization. Add `torch._C._cow_record_done(tensor)` so callers can explicitly complete those prior uses. The caller first joins all relevant uses, including aliases and earlier pointer accesses, onto the current stream. Completion records an event and releases retained references. An active COW reference remains alive until destruction or materialization; later data access cancels its completion. Each sibling has its own completion record, and renewed completion preserves its earlier event dependencies. When every departed sibling completed explicitly, the last sibling can steal the original allocation after waiting on the recorded events, without fencing unrelated later work on the allocation stream. Releasing retained references also enables subsequent lazy clones to share the new allocation. Callers that omit completion retain the conservative copying behavior. Capture boundaries remain explicit: completion and shared materialization must stay within the context's capture. Creating COW storage for an external graph input during capture is rejected. Private-pool outputs cloned after capture, unsupported sharing streams, and allocators without allocation metadata use eager copies. Unique materialization preserves the allocation's address. Compatibility impact is limited to explicit CUDA `_lazy_clone()` users: unsupported capture crossings now raise errors, and some lazy clones fall back to eager copies. CPU behavior is unchanged. Deferred materialization can introduce dependencies at the later operation, so this does not promise the exact stream behavior of an eager `clone()` at the original `_lazy_clone()` call. Validation: built on GB200 with CUDA 13.1. Passed 16 C++ COW tests, 43 native CUDA COW tests, 96 existing tensor lazy-clone tests, and two cudaMallocAsync fallback tests (157 passed total). Tests cover independent readers, copies and stealing on different streams, allocator reuse, completion renewal, chained clones, capture boundaries, repeated graph replay with changing inputs, and reduce-overhead compilation. Targeted `spin lint` and pre-commit `lintrunner -a` passed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200800 Approved by: https://github.com/ezyang |
||
|
Simon Layton
|
1be50d8e09 |
[native] Share RMSNorm forward and backward kernel launchers (#197008)
Human Notes:
Rework the launchers, reduce repeated code.
Codex:
> Replace the RMSNorm adapter's per-call CuTe conversions and ATen weight-gradient
> sum/cast with general native CuTeDSL launchers. Compile against fake tensor
> descriptors and pass read-only tensors through TVM FFI for every eligible width,
> including mixed gradient dtypes. Remove the older vendored launch path.
>
> Review the kernel specification and launch policy, then the adapter. Reuse the
> shared dtype, hardware, capability, and launch utilities from main. Hoist tensor
> preparation, const pointer access, and partial-row reduction into shared modules.
> Type the new interfaces and add the missing Tensor.const_data_ptr stub definition.
> The tests cover layouts, alignment, copy-on-write, output masks, fake tensors,
> and fallback routing without duplicating the later persistent-grid cases.
>
> Performance
> -----------
> Measured on NVIDIA B200 with CUDA 13.2 on 2026-10-05 after rebasing onto
>
|