FHElium

fhelium.runtime.cuda_graph ​

CUDA Graph capture and replay for fixed-schedule FHElium evaluators.

CudaGraphProgram stages dynamic positional tensor and FHElium-value inputs through reusable fixed-address buffers. Callers bind static resources and schedule before capture; captured callables use fixed Python control flow.

CudaGraphCaptureStats ​

class View source

python
CudaGraphCaptureStats(warmup_iterations: int, warmup_seconds: float, capture_seconds: float, first_replay_seconds: float, memory_allocated_bytes: int, memory_reserved_bytes: int)

One-time construction costs and device memory after capture.

Parameters

  • warmup_iterations: Number of uncaptured warmup evaluations requested.
  • warmup_seconds: Host-observed total warmup duration.
  • capture_seconds: Host-observed CUDA Graph capture duration.
  • first_replay_seconds: Host-observed duration of the synchronized first native graph replay used for validation.
  • memory_allocated_bytes: CUDA allocator bytes live after capture.
  • memory_reserved_bytes: CUDA allocator reserved bytes after capture.

Attributes ​

NameTypeDefault/value
warmup_iterationsint
warmup_secondsfloat
capture_secondsfloat
first_replay_secondsfloat
memory_allocated_bytesint
memory_reserved_bytesint

CudaGraphPool ​

class View source

python
CudaGraphPool(device: torch.device | str='cuda')

Own shared temporary storage for sequential CUDA Graph programs.

Programs captured with this pool keep their inputs and retained outputs outside its scratch storage. Their graphs may therefore replay in a different order without overwriting another program's retained output. An event chain serializes scratch use across caller-selected streams. Persistent resources are prepared outside capture; captured intermediates must not escape into application state outside the returned output tree.

One host caller controls a pool. Concurrent workers use separate pools. Retained output leaves have independent storage and no autograd history; input/output and output/output aliases are not preserved. The callable executes in the caller's grad mode. Closing a pool synchronizes its work and closes its live member programs without extending their lifetimes.

device ​

property

python
device: torch.device

Device owning the pool's scratch storage.

close ​

method

python
def close() -> None: ...

Wait for scratch users and close this pool's live member programs.

CudaGraphProgram ​

class View source

python
CudaGraphProgram(*, function: Callable[..., _OutputT], device: torch.device, graph: torch.cuda.CUDAGraph, input_buffer: ReusableValueBuffer[tuple[object, ...]], static_inputs: tuple[object, ...], output: _OutputT, stats: CudaGraphCaptureStats, pool: CudaGraphPool | None=None)

Bases: Generic[_OutputT]

CUDA Graph adapter for an eager callable with reusable dynamic inputs.

The original callable remains independently valid for eager execution. Internally, dynamic arguments use a fhelium.runtime.ReusableValueBuffer; CUDA Graph kernels capture the buffer's fixed target addresses while replay changes only its payload.

Construct programs with capture; direct construction is not public. Parameters captured by a closure, functools.partial, or a callable object are static program state. Positional example_inputs define the reusable dynamic argument tree. Supported leaves are tensors and serializable FHElium values nested in lists, tuples, and dictionaries.

Dynamic source values may later reside on CPU or CUDA as long as their device-independent structure and metadata match capture. CPU-to-CUDA overlap requires pinned sources and the advanced copy_inputs_from plus replay_prepared path. Ordinary replay remains a same-call convenience.

With a CudaGraphPool, specializations share temporary storage while retaining separate input/output buffers. The pool orders scratch use and closes its live members; replay order may differ from capture order.

Replay returns a borrowed output tree whose storage is overwritten by the next replay. Pass copy_output=True when the result must outlive that replay. One program instance owns one input/output storage set and does not execute replays concurrently; create independent instances for concurrent workers.

Examples

A normal evaluator is still callable eagerly::

schedule = partial(
    matrix_vector,
    engine=engine,
    diagonals=diagonals,
    rotation_keys=rotation_keys,
)
eager_output = schedule(source)

Capture it for repeated execution::

program = CudaGraphProgram.capture(
    schedule,
    example_inputs=(prototype_source,),
)
borrowed = program.replay(next_source)
owned = program.replay(next_source, copy_output=True)

Separate transfer from replay with a caller-owned stream::

copied = program.copy_inputs_from(
    pinned_cpu_source,
    stream=transfer_stream,
    non_blocking=True,
)
with torch.cuda.stream(compute_stream):
    output = program.replay_prepared(copy_handle=copied)

capture ​

method

python
def capture(function: Callable[..., _OutputT], *, example_inputs: Sequence[object], warmup: int=3, check_input_liveness: bool=True, pool: CudaGraphPool | None=None) -> CudaGraphProgram[_OutputT]: ...
def capture(function: None=None, *, example_inputs: Sequence[object], warmup: int=3, check_input_liveness: bool=True, pool: CudaGraphPool | None=None) -> Callable[[Callable[..., _OutputT]], CudaGraphProgram[_OutputT]]: ...

Warm up and capture one deterministic evaluator schedule.

example_inputs defines both the dynamic positional argument tree and the fixed CUDA target residency. Warmup uses fresh reusable buffers, so backend JIT and lazy materialization happen without mutating the buffer retained by the captured graph.

Bind static parameters through a closure, functools.partial, or a callable object. Dynamic inputs use positional Tensor or FHElium-value arguments. Use a positional adapter for a dynamic keyword-only argument.

Passing function directly is the primary API. Omitting it returns a decorator shorthand that captures at function-definition time::

@CudaGraphProgram.capture(example_inputs=(prototype,))
def captured(source):
    return evaluator(source)

The decorated name is the ready CudaGraphProgram, not the original Python callable. All CUDA state, static resources, and example inputs must therefore already exist when the definition executes.

Parameters

  • function: Optional callable evaluator whose Python control flow and CUDA work are fixed at capture. Output must be a CUDA tensor, FHElium value, nested list/tuple/dict of those leaves, or None. Omit it only for the decorator shorthand.
  • example_inputs: CUDA-resident positional prototypes. Their structure and metadata specialize this program.
  • warmup: Number of fresh-buffer side-stream evaluations before capture.
  • check_input_liveness: Forwarded to torch.cuda.graph.
  • pool: Optional owner of shared temporary storage. Requires at least one warmup to allocate retained output storage outside the pool. Retained output leaves have cloned storage without autograd history or input/output and output/output aliases. The callable retains the caller's grad mode. Programs in a pool execute sequentially across streams.

Returns

A ready CUDA Graph program with one reusable input buffer, or a decorator producing one when function is omitted.

Raises

  • TypeError: If function or an input/output leaf is unsupported.
  • ValueError: If warmup is negative, examples have no tensor leaves, are not on one CUDA device, representative storage aliases, or shared-pool capture lacks a warmup or uses another device.
  • CudaGraphCaptureError: If warmup, capture, output validation, or the first native replay fails.

device ​

property

python
device: torch.device

CUDA device on which this program was captured.

input_signature ​

property

python
input_signature: ValueTreeSignature

Device-independent signature of the positional argument tuple.

input_nbytes ​

property

python
input_nbytes: int

Logical bytes in the program's fixed dynamic-input storage.

stats ​

property

python
stats: CudaGraphCaptureStats

One-time construction timings and post-capture memory snapshots.

Memory fields are process allocator snapshots, not graph-exclusive incremental footprints.

output ​

property

python
output: _OutputT

Borrowed output tree retained by the captured graph.

Its storage is overwritten by the next replay. Use replay(..., copy_output=True) when independent ownership is needed.

cuda_graph ​

property

python
cuda_graph: torch.cuda.CUDAGraph

Underlying PyTorch CUDA Graph for low-level inspection.

Calling its replay() directly bypasses input validation, payload copying, transfer-event waits, and overwrite protection. It reuses the payload currently held by the input buffer.

copy_inputs_from ​

method

python
def copy_inputs_from(*inputs: object, stream: torch.cuda.Stream | None=None, non_blocking: bool=True, wait_for: torch.cuda.Event | Sequence[torch.cuda.Event] | None=None) -> CopyHandle: ...

Copy matching positional inputs without replaying the graph.

This is the advanced prefetch mechanism. The caller decides when and on which stream to submit the copy. If a previous replay was launched, its completion event is automatically added as an overwrite dependency, so a transfer stream cannot replace input storage still read by the graph.

Source values may live on pinned CPU or CUDA. Actual asynchronous H2D overlap requires pinned CPU tensors.

Returns

A CopyHandle accepted by replay_prepared.

Raises

  • CudaGraphInputError: If structure or value state differs from capture.
  • ExecutionError: If the program is closed.

replay_prepared ​

method

python
def replay_prepared(*, copy_handle: CopyHandle | None=None, copy_output: bool=False, synchronize: bool=False) -> _OutputT: ...

Replay the payload already held by the input buffer.

Parameters

  • copy_handle: Optional handle returned by the latest copy_inputs_from. The current compute stream waits on its completion event without blocking the CPU. Omit this only to replay the buffer's current payload again.
  • copy_output: Clone the output after replay instead of returning borrowed graph storage.
  • synchronize: Synchronize the capture device before returning.

Returns

Borrowed retained output, or an independently owned clone.

Raises

  • CudaGraphInputError: If copy_handle belongs to another buffer or is not the most recent prepared payload.
  • ExecutionError: If the program is closed.

replay ​

method

python
def replay(*inputs: object, copy_output: bool=False, synchronize: bool=False) -> _OutputT: ...

Copy matching inputs, replay, and return the graph output.

This convenience path performs validation, payload copy, and replay on the caller's current CUDA stream. For transfer/compute overlap, use copy_inputs_from on a transfer stream followed by replay_prepared on a compute stream.

Tensor source devices may differ from capture, but structure, tensor topology, and FHElium metadata must match. The default output is borrowed; copy_output=True returns independent storage.

close ​

method

python
def close() -> None: ...

Synchronize outstanding work and release graph-owned references.

Closing is idempotent. The program, borrowed output, input buffer, and direct cuda_graph access must not be used afterward. CUDA caching allocator reserved memory may remain available for process reuse.

TensorSignature ​

class View source

python
TensorSignature(shape: tuple[int, ...], stride: tuple[int, ...], dtype: torch.dtype, layout: torch.layout, requires_grad: bool)

Device-independent tensor topology accepted by one execution buffer.

Shape, stride, dtype, layout, and requires_grad determine copy compatibility with fixed target storage. The owning buffer or program records that storage's device.

Parameters

  • shape: Tensor dimensions.
  • stride: Element strides required by fixed target storage.
  • dtype: Tensor scalar dtype.
  • layout: PyTorch tensor layout, normally torch.strided.
  • requires_grad: Autograd flag expected by the reusable payload.

Attributes ​

NameTypeDefault/value
shapetuple[int, ...]
stridetuple[int, ...]
dtypetorch.dtype
layouttorch.layout
requires_gradbool

from_tensor ​

method

python
def from_tensor(tensor: torch.Tensor) -> TensorSignature: ...

Build a copy-compatibility signature for tensor.

Parameters

  • tensor: Representative tensor whose device-independent topology is captured.

Returns

Signature excluding the tensor's current device.

ValueSignature ​

class View source

python
ValueSignature(type_name: str, schema_version: int, metadata: tuple[tuple[str, object], ...], tensors: tuple[tuple[str, TensorSignature], ...])

FHElium value state plus device-independent tensor topology.

The signature combines the serialized value type and schema version, frozen value metadata, and named Tensor signatures. Matching compares these recorded fields.

Parameters

  • type_name: Registered serialized-value class name.
  • schema_version: Value-serialization schema version.
  • metadata: Frozen, deterministically ordered non-tensor state.
  • tensors: Ordered tensor names and their device-independent signatures.

Attributes ​

NameTypeDefault/value
type_namestr
schema_versionint
metadatatuple[tuple[str, object], ...]
tensorstuple[tuple[str, TensorSignature], ...]

from_value ​

method

python
def from_value(value: TensorResident) -> ValueSignature: ...

Describe one serializable FHElium value.

Parameters

  • value: Resident value whose type, metadata, and tensor topology are captured.

Returns

Device-independent value signature.

ValueTreeSignature ​

class View source

python
ValueTreeSignature(kind: TreeKind, leaf: TensorSignature | ValueSignature | None = None, children: tuple[ValueTreeSignature, ...] = (), keys: tuple[object, ...] = ())

Structure of tensors and FHElium values in a reusable execution payload.

Supported leaves are torch.Tensor and serializable FHElium fhelium.values.TensorResident values. Lists, tuples, and dictionaries may nest those leaves. Bind Python scalars and control objects statically in the callable.

Each Tensor leaf uses a TensorSignature; each FHElium-value leaf uses a ValueSignature containing its value metadata and Tensor topology.

Parameters

  • kind: Node kind: tensor, FHElium value, list, tuple, or dictionary.
  • leaf: Tensor/value signature for a leaf node; None for containers.
  • children: Ordered signatures nested by a container node.
  • keys: Dictionary keys in the same order as children; empty for all other node kinds.

Attributes ​

NameTypeDefault/value
kindTreeKind
leafTensorSignature | ValueSignature | NoneNone
childrentuple[ValueTreeSignature, ...]()
keystuple[object, ...]()

from_value ​

method

python
def from_value(value: object) -> ValueTreeSignature: ...

Describe a supported tensor/value tree.

Parameters

  • value: Representative tensor/value tree to describe.

Returns

Recursive device-independent structure and state signature.

Raises

  • TypeError: If a leaf is not a tensor or serializable FHElium value, or if a container is not a list, tuple, or dictionary.

validate ​

method

python
def validate(value: object, *, path: str='value') -> None: ...

Require value to have this structure and state.

Tensor devices may differ because signatures describe transfer compatibility, not residency. Validation completes for the full tree before fhelium.runtime.ReusableValueBuffer copies any payload.

Parameters

  • value: Candidate tree to compare with this signature.
  • path: Root label used in mismatch diagnostics.

Raises

  • ExecutionInputError: If structure, tensor topology, or value metadata differs.