Metadata-Version: 2.4
Name: piecewise_cuda_graphs
Version: 0.1.0rc0
Summary: Piecewise CUDA graphs for PyTorch
Author: Meta Platforms, Inc.
License: BSD 3-Clause License
        
        Copyright (c) 2026, Meta Platforms, Inc. and affiliates.
        All rights reserved.
        
        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions are met:
        
        1. Redistributions of source code must retain the above copyright notice, this
           list of conditions and the following disclaimer.
        
        2. Redistributions in binary form must reproduce the above copyright notice,
           this list of conditions and the following disclaimer in the documentation
           and/or other materials provided with the distribution.
        
        3. Neither the name of the copyright holder nor the names of its
           contributors may be used to endorse or promote products derived from
           this software without specific prior written permission.
        
        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
        AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
        IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
        DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
        FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
        DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
        SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
        CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
        OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
        OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
        
Project-URL: Homepage, https://github.com/meta-pytorch/piecewise_cuda_graphs
Project-URL: Repository, https://github.com/meta-pytorch/piecewise_cuda_graphs
Project-URL: Issues, https://github.com/meta-pytorch/piecewise_cuda_graphs/issues
Keywords: pytorch,torch,cuda,cuda-graphs
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: BSD License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Dynamic: license-file

# Piecewise CUDA Graphs

Piecewise CUDA Graphs extends the standard PyTorch CUDA graph workflow.

With `torch.cuda.graph`, one context manager captures one CUDA graph. That works
well when the whole region is capture-compatible. `piecewise-cuda-graphs` keeps
the same capture/replay shape, but lets one `with piecewise_graph(...)` block
produce a sequence of CUDA graph segments separated by explicit eager breaks. If
no eager breaks occur, the block is captured as a single graph segment.

Use this when most of a workload should run under CUDA graphs, but some sections
are not CUDA-graph-compatible or otherwise need to run eagerly. Mark those
sections with `@no_graph`, and make their CUDA inputs/outputs obey the
constraints below.

## Quick start

Install from source into an environment with PyTorch and CUDA support:

```bash
git clone https://github.com/meta-pytorch/piecewise-cuda-graphs.git
cd piecewise-cuda-graphs
pip install -e .
```

The capture flow is the same as regular CUDA graphs: allocate static buffers,
warm up on a side stream, capture, then replay by updating the static inputs.
The only addition is `@no_graph`, which marks functions that should run eagerly
between captured graph segments.

```python
import torch
from piecewise_cuda_graphs import CUDAGraphSequence, no_graph, piecewise_graph

@no_graph
def dynamic_scale(x: torch.Tensor) -> None:
    # Cannot be captured: reads a value back to the CPU.
    if x.sum().item() > 0:
        x.clamp_(min=0)

# Pre-allocate static buffers.
static_input = torch.empty(1024, device="cuda")
result = torch.empty(1024, device="cuda")

def workload(src: torch.Tensor, dst: torch.Tensor) -> None:
    dst.copy_(src * 2)
    dynamic_scale(dst)  # ends the current graph segment and runs eagerly
    dst.add_(1.0)

# Warm up on a side stream, as required by CUDA graphs.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        static_input.fill_(1.0)
        workload(static_input, result)
torch.cuda.current_stream().wait_stream(s)

# Capture.
seq = CUDAGraphSequence()
with piecewise_graph(seq):
    workload(static_input, result)

# Replay with new data by overwriting the static input buffer.
static_input.fill_(5.0)
seq.replay()
```

## Constraints

- `@no_graph` functions must not return CUDA tensors. Write CUDA outputs into
  pre-allocated buffers passed as arguments.
- Side streams must be joined back to the capturing stream before entering an
  `@no_graph` function or leaving the `piecewise_graph` context. Set
  `PIECEWISE_CUDA_GRAPHS_DEBUG=1` to add still-unjoined stream id(s) to the
  resulting error.
- Usual CUDA graph constraints still apply: replay uses the same tensor addresses
  captured during warmup/capture.

## Additional usage

### Explicit split points

`force_no_graph()` inserts a graph break with no eager work. This can be useful
for debugging or isolating capture regions.

```python
from piecewise_cuda_graphs import CUDAGraphSequence, force_no_graph, piecewise_graph

seq = CUDAGraphSequence()
with piecewise_graph(seq):
    a = step1(x)
    force_no_graph()
    b = step2(a)
```

### Sharing memory pools

All graph segments within a sequence share the same CUDA graph memory pool. You
can also share pools across sequences:

```python
seq1 = CUDAGraphSequence()
with piecewise_graph(seq1):
    workload_a(buf_a, src_a)

seq2 = CUDAGraphSequence(pool=seq1.pool())
with piecewise_graph(seq2):
    workload_b(buf_b, src_b)
```

## Reference

- **`CUDAGraphSequence(pool=None)`**: captured graph/eager segment sequence.
  Methods: `replay()`, `reset()`, `pool()`.
- **`piecewise_graph(seq, stream=None, capture_error_mode="global")`**: capture
  context, analogous to `torch.cuda.graph`.
- **`@no_graph` / `@no_graph(enable=...)`**: mark functions that run eagerly
  inside `piecewise_graph`.
- **`force_no_graph()`**: explicit split point with no eager work.

For implementation details, see [DESIGN.md](DESIGN.md).

## License

BSD 3-Clause License. See [LICENSE](LICENSE) for details.
