fhelium.execution.cuda_graph

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

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

CudaGraphCaptureStats

class View source

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

Name	Type	Default/value
warmup_iterations	int
warmup_seconds	float
capture_seconds	float
first_replay_seconds	float
memory_allocated_bytes	int
memory_reserved_bytes	int

CudaGraphProgram

class View source

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

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.execution.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 bound 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 exact 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 exact 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.

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

def capture(function: Callable[..., _OutputT], *, example_inputs: Sequence[object], warmup: int=3, check_input_liveness: bool=True) -> CudaGraphProgram[_OutputT]: ...
def capture(function: None=None, *, example_inputs: Sequence[object], warmup: int=3, check_input_liveness: bool=True) -> 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. This method intentionally does not accept dynamic keyword arguments or arbitrary Python control objects. Use a positional adapter for a dynamic keyword-only tensor or exact value.

Passing function directly is the canonical 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, exact 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 exact metadata specialize this program.
• warmup: Number of fresh-buffer side-stream evaluations before capture.
• check_input_liveness: Forwarded to torch.cuda.graph.

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, or representative storage aliases.
• CudaGraphCaptureError: If warmup, capture, output validation, or the first native replay fails.

device

property

device: torch.device

CUDA device on which this program was captured.

input_signature

property

input_signature: ValueTreeSignature

Device-independent signature of the positional argument tuple.

input_nbytes

property

input_nbytes: int

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

stats

property

stats: CudaGraphCaptureStats

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

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

output

property

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

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

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 exact state differs from capture.
• ExecutionError: If the program is closed.

replay_prepared

method

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 intentionally 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

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 exact FHElium metadata must match. The default output is borrowed; copy_output=True returns independent storage.

close

method

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

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.

device is intentionally absent. Shape, stride, dtype, layout, and requires_grad determine whether payload can be copied into fixed target storage; the target residency belongs to the buffer or program that owns that storage.

Parameters

• shape: Exact tensor dimensions.
• stride: Exact 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

Name	Type	Default/value
shape	tuple[int, ...]
stride	tuple[int, ...]
dtype	torch.dtype
layout	torch.layout
requires_grad	bool

from_tensor

method

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

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

Exact FHElium value state plus device-independent tensor topology.

A value signature fixes cryptographic metadata such as context, level, scale, plaintext representation, polynomial domain, modulus basis, residue representation, prime identities, schema version, and key identity. A value signature is storage-independent metadata. Model, user, request, and cache associations remain external.

Parameters

• type_name: Registered serialized exact-value class name.
• schema_version: Exact-value serialization schema version.
• context_id: Cryptographic context identity, if the value has one.
• metadata: Frozen, deterministically ordered non-tensor exact state.
• tensors: Ordered tensor names and their device-independent signatures.

Attributes

Name	Type	Default/value
type_name	str
schema_version	int
context_id	str | None
metadata	tuple[tuple[str, object], ...]
tensors	tuple[tuple[str, TensorSignature], ...]

from_value

method

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

Describe one serializable exact FHElium value.

Parameters

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

Returns

Device-independent exact-value signature.

ValueTreeSignature

class View source

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

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

Supported leaves are torch.Tensor and serializable FHElium fhelium.core.TensorResident values. Lists, tuples, and dictionaries may nest those leaves. Arbitrary Python scalars and control objects are intentionally unsupported; bind them statically in a callable or keep them in the application control plane.

ValueTreeSignature composes TensorSignature and ValueSignature rather than replacing them: tensor leaves need only tensor topology, while exact-value leaves additionally require cryptographic metadata.

Parameters

• kind: Node kind: tensor, exact 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

Name	Type	Default/value
kind	TreeKind
leaf	TensorSignature | ValueSignature | None	None
children	tuple[ValueTreeSignature, ...]	()
keys	tuple[object, ...]	()

from_value

method

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

Describe a supported tensor/exact-value tree.

Parameters

• value: Representative tensor/exact-value tree to describe.

Returns

Recursive device-independent structure and exact-state signature.

Raises

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

validate

method

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

Require value to have this exact structure and state.

Tensor devices may differ because signatures describe transfer compatibility, not residency. Validation completes for the full tree before fhelium.execution.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 exact-value metadata differs.