FHElium

Appearance

fhelium.execution ​

Execution signatures, reusable value buffers, and CUDA Graphs.

CopyHandle ​

class View source

python
CopyHandle(*, event: torch.cuda.Event | None, device: torch.device, source_tensors: tuple[torch.Tensor, ...], bytes_copied: int, target_token: object)

Future-like handle for one copy enqueued into a reusable buffer.

A handle owns the exact submitted tensor leaves until the CUDA event reports completion. This prevents a mutable source tree from releasing or replacing pinned host or device storage while an asynchronous copy may still read it. wait_on inserts a stream dependency without blocking the CPU; synchronize blocks the caller.

CopyHandle is intentionally not an asyncio Future and is not awaitable. A future extension may bridge CUDA events to an async scheduler, but the core object currently exposes CUDA stream/event ordering only.

Instances are returned by ReusableValueBuffer.copy_from; direct construction is not part of the public API. Internally, event tracks completion, device identifies the target, source_tensors keep the submitted storage alive, bytes_copied records payload size, and target_token binds the handle to its originating buffer.

event ​

property

python
event: torch.cuda.Event | None

CUDA completion event, or None for a synchronous CPU copy.

device ​

property

python
device: torch.device

Target device of the copy.

bytes_copied ​

property

python
bytes_copied: int

Logical tensor payload bytes submitted by this copy.

done ​

method

python
def done() -> bool: ...

Return whether the copy has completed without blocking.

wait_on ​

method

python
def wait_on(stream: torch.cuda.Stream | None=None) -> CopyHandle: ...

Make stream wait for this copy and return self.

The default is the current stream on the copy target device. This method enqueues an event wait and does not synchronize the CPU.

Parameters

  • stream: Consumer CUDA stream. Defaults to the current stream on this handle's target device.

Returns

This handle, allowing fluent ordering before consumption.

Raises

  • ValueError: If the supplied stream belongs to another device.

synchronize ​

method

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

Block the CPU until the copy completes, then release the source.

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

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)

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

python
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

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

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

ReusableValueBuffer ​

class View source

python
ReusableValueBuffer(*, node: _BufferNode, signature: ValueTreeSignature, device: torch.device, pinned: bool)

Bases: Generic[_ValueT]

Fixed-structure storage whose exact-value payload may be replaced.

Construct a buffer from a representative tensor/exact-value tree with like. The representative defines container structure, cryptographic metadata, and tensor topology. The buffer separately owns tensors on one target device. copy_from accepts a matching tree on CPU or CUDA, validates the complete source before copying anything, and preserves every target tensor address.

The buffer identifies payloads solely by exact structure. Applications or extensions associate it with model weights, user keys, request ciphertexts, and cache/prefetch/eviction policy.

The value property reconstructs ordinary tensors and exact FHElium values around the fixed storage, so an eager callable can consume it without a buffer-specific evaluator API. Do not retain a reconstructed tree as an ownership or readiness signal; use the buffer and the CopyHandle returned by copy_from for lifetime and stream ordering.

Examples

Reuse one CUDA allocation for changing plaintext tiles::

    buffer = ReusableValueBuffer.like(
        prototype_tile,
        device="cuda:0",
    )
    copied = buffer.copy_from(
        pinned_cpu_tile,
        stream=transfer_stream,
        non_blocking=True,
    )

    with torch.cuda.stream(compute_stream):
        copied.wait_on(compute_stream)
        output = evaluate_tile(ciphertext, buffer.value)

Double buffering is built by creating two independent buffers and
copying into one while an eager program reads the other. Choosing the
next tile and recording the prior consumer event remain application
policy.

Construct buffers with like; the direct constructor accepts internal allocation-tree state and is not part of the public API.

like ​

method

python
def like(example: _ValueT, *, device: torch.device | str | None=None, pin_memory: bool=False) -> ReusableValueBuffer[_ValueT]: ...

Allocate independent storage with the structure of example.

Parameters

  • example: Tensor/exact-value tree defining structure, metadata, and initial payload. Supported containers are lists, tuples, and dictionaries.
  • device: Target residency. If omitted, all example tensors must already share one device, which becomes the target.
  • pin_memory: Allocate pinned host tensors. This requires a CPU target and is useful before non-blocking host-to-device copies.

Returns

A buffer initialized with a copy of example.

Raises

  • TypeError: If the tree contains an unsupported leaf.
  • ValueError: If there are no tensor leaves, the implicit target is ambiguous, pin_memory is requested for a non-CPU target, or representative leaves alias the same storage.

signature ​

property

python
signature: ValueTreeSignature

Device-independent structure and exact-value state of this buffer.

device ​

property

python
device: torch.device

Target residency of the owned storage.

is_pinned ​

property

python
is_pinned: bool

Whether all owned CPU tensors use pinned host memory.

nbytes ​

property

python
nbytes: int

Logical bytes in the fixed target tensors.

value ​

property

python
value: _ValueT

A fresh ordinary value tree around the fixed target tensors.

Container and exact-value wrapper objects may be newly reconstructed on each access; their tensor addresses remain fixed for the buffer lifetime. Payload readiness is governed by the relevant CopyHandle or caller stream ordering.

copy_from ​

method

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

Copy a matching source tree into the fixed target storage.

Full structural and exact-state validation happens before the first Tensor.copy_. Tensor source devices may differ from the target and from the representative used by like.

Parameters

  • source: Matching tensor/exact-value tree whose payload replaces the buffer contents.
  • stream: Caller-owned CUDA stream on which target writes are enqueued. Defaults to the current target-device stream. CPU buffers do not accept a stream.
  • non_blocking: Forwarded to Tensor.copy_. Actual asynchronous host-to-device overlap requires pinned CPU source tensors.
  • wait_for: One event or sequence of events that the copy stream must wait for before overwriting this buffer, for example an event recorded after the prior eager consumer finished reading it.

Returns

A future-like CopyHandle. Keep or wait on this handle before consuming the new payload. The buffer also retains in-flight handles so dropping the returned Python object cannot prematurely release source storage.

Raises

  • ExecutionInputError: If source structure, metadata, or tensor topology differs from the buffer.
  • ValueError: If stream/event arguments are incompatible with a CPU target or the stream belongs to another CUDA device.
  • ExecutionError: If the buffer is closed.

Note

This method orders writes but cannot infer when an arbitrary eager consumer has finished reading the previous payload. Record a CUDA event after that consumer and pass it as wait_for before reusing the buffer.

wait_for ​

method

python
def wait_for(handle: CopyHandle, stream: torch.cuda.Stream | None=None) -> None: ...

Wait on a copy produced by this buffer without blocking the CPU.

Parameters

  • handle: Copy handle previously returned by this buffer.
  • stream: Consumer CUDA stream, or the current target stream when omitted.

Raises

  • ExecutionInputError: If handle belongs to another buffer.

close ​

method

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

Wait for outstanding copies and release target-storage references.

Closing is idempotent. The caller remains responsible for ensuring that arbitrary eager or captured consumers have finished reading the target tensors before closing; the buffer can track its own writes but cannot infer external read completion.

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.

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 ​

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, 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 ​

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

from_value ​

method

python
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

python
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 ​

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

from_value ​

method

python
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

python
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.

pin_value_tree ​

function View source

python
def pin_value_tree(value: _ValueT) -> _ValueT: ...

Clone a supported value tree into pinned CPU storage.

The returned ordinary tree owns pinned tensors and has the same device-independent signature as value. This helper chooses no cache or persistence policy; callers retain and release the returned tree normally.

Parameters

  • value: Supported tensor/exact-value tree to clone. Source leaves may be on CPU or CUDA.

Returns

Structurally equivalent ordinary value tree backed by pinned CPU tensors.

value_tree_nbytes ​

function View source

python
def value_tree_nbytes(value: object) -> int: ...

Return logical tensor bytes in one supported execution value tree.

Parameters

  • value: Tensor/exact-value tree whose tensor payload is counted.

Returns

Sum of numel * element_size over every tensor leaf.

Towards full-stack encrypted execution infrastructures

Released under the MIT License. Copyright VisualDust and FHElium contributors