Explicit residency plans and CUDA leases

Example source: examples/13_explicit_residency.py

This example places live FHElium plaintext, key, and ciphertext values under a ResidencyManager. It performs residency transitions, dry-runs one named stage plan with a CUDA reservation, runs real CKKS evaluation on an indexed CUDA consumer stream, observes unbudgeted pageable accounting beside strict pinned/CUDA budgets, and ends each managed logical value.

Run the example

python examples/13_explicit_residency.py \
  --device cuda:0 \
  --preset slots8192-scale40-levels7-int64

CUDA is required. The workload rotates one encrypted message by one slot, multiplies it by a prepared 0.5 plaintext, rescales, decrypts, and compares with the cleartext result.

1. Start with live exact values

The example creates:

• an operation-ready Plaintext weight;
• a RotationKey for step +1;
• an encrypted request Ciphertext.

The stable ResidencyManager owns process-local managed values through opaque handles.

weight = engine.prepare_plaintext_for_multiplication(
    engine.encode(weight_message, level=0),
    modulus_basis="Q",
).cpu()
rotation_key = engine.create_rotation_key(1, engine.secret_key).cpu()
source = engine.encrypt_message(message)

2. Configure optional admission budgets

Locations are immutable identities. Host locations are canonical constants; CUDA locations require an device index.

device_location = cuda_location(engine.device)
residency = ResidencyManager(
    budgets={
        PINNED_HOST: pinned_capacity,
        device_location: cuda_budget,
    },
)

The example leaves pageable host unbudgeted while applying strict budgets to pinned host and CUDA. The manager still accounts every pageable materialization and lifetime. At the two budgeted locations, current materialization charges plus active MemoryReservation charges must fit the application-supplied byte limit.

ResidencyManager() with no arguments is the fully managed unbudgeted form. Valid host and indexed CUDA locations become part of its accounting lazily after their first successful managed use. A budgets entry adds a strict admission limit only to that location; other valid local locations remain unbudgeted. One manager can therefore use several indexed CUDA locations in the current process without predeclaring all of them.

The budgets mapping defines strict admission limits. Free-memory readings from PyTorch or NVML remain observations used by the application when selecting those values. A direct materialization or reservation that would exceed a configured budget raises ResidencyBudgetError. Plan explanation instead reports infeasibility, and scope entry raises ResidencyPlanError before executing an infeasible plan. Application code then chooses its next move, drop, workload-admission, or manager-budget decision as a named operation.

3. Adopt values and receive opaque handles

weight_handle = residency.adopt(
    weight,
    at=PAGEABLE_HOST,
    replica_mode=ReplicaMode.REPLICABLE,
)
source_handle = residency.adopt(
    source,
    at=device_location,
    replica_mode=ReplicaMode.EXCLUSIVE,
)
del weight, source

Every adopt call returns a fresh unique ResidencyHandle. The handle is an opaque process-local token that application code stores and passes back to the manager as a complete object. Moving or replicating materializations preserves that handle.

adopt transfers logical ownership under caller-enforced ownership rules. Python cannot destroy other aliases, so callers must stop using the input value and must not allow raw values obtained from a later lease to escape that lease.

The weight and rotation key are REPLICABLE, so they may have simultaneous host and CUDA materializations. The request ciphertext is EXCLUSIVE, so it has one steady materialization and changes location with move.

Adopted values have Recoverability.MUST_PRESERVE: their last materialization cannot be dropped accidentally. The separate register_source API accepts a synchronous ResidencySource only with Recoverability.RECONSTRUCTIBLE and likewise returns a fresh opaque handle; the returned handle identifies that registration to the manager.

4. Issue direct transitions

residency.ensure(weight_handle, PINNED_HOST)
residency.move(
    source_handle,
    PINNED_HOST,
    from_location=device_location,
)

ensure creates a replica and retains existing materializations. move creates the destination and removes the selected source. drop removes one unprotected replica, while discard ends the managed value and removes all of its unprotected state.

The application names each destination and removal.

5. Build and explain a stage plan

plan = ResidencyPlan(
    name="inference/rotate-scale/tile-0",
    enter=(
        EnsureResident(weight_handle, device_location),
        EnsureResident(key_handle, device_location),
        MoveResident(
            source_handle,
            device_location,
            from_location=PINNED_HOST,
        ),
    ),
    exit=(
        MoveResident(
            source_handle,
            PINNED_HOST,
            from_location=device_location,
        ),
        DropResident(key_handle, device_location),
        DropResident(weight_handle, device_location),
    ),
    reservations=(
        MemoryReservation(
            device_location,
            workspace_bytes,
            label="rotate/multiply outputs and workspace",
        ),
    ),
)
explanation = residency.explain(plan)

A plan is ordered low-level intermediate representation (IR). explain dry-runs reclaim, reservation, entry, and exit effects against current state, checks budgets and ownership constraints, resolves sources, and predicts managed storage peaks without loading a source or executing the plan.

The complete scope order is reclaim actions, reservation admission, entry actions, the application body, exit actions, and reservation release. This example has no reclaim prefix because its configured budget already admits the stage. A reservation is accounted headroom for unmanaged expansion. In this example the reservation reduces the remaining CUDA budget for the scope lifetime; it does not allocate the rotation output, multiplication output, or native workspace.

The string inference/rotate-scale/tile-0 is a diagnostic plan name. Stages and tiles are ordinary named, nestable program structure. Plan actions still refer to the opaque local handles. This is where the application can control FHE memory expansion: select only the current weights and keys, reserve measured output/workspace headroom, and release the window at a known completion point. scope(..., transfer_streams={device_location: transfer_stream}) can bind a copy stream per CUDA destination. A prepared caller may also pass expected_state_version=snapshot.state_version; scope entry rejects a stale version before mutation.

6. Protect a CUDA consumer stream

compute_stream = torch.cuda.Stream(device=engine.device)
scope = residency.scope(plan)
with scope:
    with residency.acquire(
        (source_handle, key_handle, weight_handle),
        at=device_location,
        consumer_stream=compute_stream,
    ) as resident:
        with torch.cuda.stream(compute_stream):
            rotated = engine.rotate_with_key(
                resident[source_handle],
                resident[key_handle],
            )
            output = engine.multiply_plaintext(
                engine.coefficient_domain_to_ntt_domain(rotated),
                resident[weight_handle],
            )
            output = engine.rescale_to_next_level(
                engine.ntt_domain_to_coefficient_domain(output)
            )

    released_snapshot = residency.snapshot()
    compute_stream.synchronize()

acquire borrows already-resident values; it does not materialize missing ones. The borrowed mapping checks that the lease remains active and must not be retained as a source of raw aliases after release.

CUDA acquisition requires the initial consumer_stream argument. This is a lifetime identity, not a performance hint: release may occur on another Python thread, whose ambient current stream is unrelated to the consumer.

At lease release, the manager records an event on compute_stream. If kernels are still reading the values, the manager retains pending protection until the event completes. It does not synchronize the full device solely to close the Python lease. The example synchronizes only the result-producing stream before host decryption and before plan exit removes or moves the inputs.

A consumer that launches reads on more streams must call lease.add_consumer_stream(...) before release. A ResidencyHold is different: it retains materializations across a longer application lifetime but exposes no values. Evaluation still requires a lease.

7. Choose manual or automatic admission

This example deliberately keeps placement actions and plan order under direct application control. Use the separate automatic residency admission workflow when the application should state exact working-set endpoints and headroom while a deterministic policy selects legal reclaim actions. Both workflows execute through the same manager authority and strict leases; automation does not change ResidencyManager.acquire() or introduce background movement.

8. Read the accounting layers correctly

ResidencySnapshot contains no tensors. For every materialization it reports:

• logical tensor payload bytes;
• actual unique backing-storage bytes and the fixed conservative charge;
• active use and hold counts;
• pending CUDA consumer-event count.

For each location it reports optional budget_bytes, optional remaining_budget_bytes, used and reserved bytes, peak materialization charge (peak_used_bytes), peak total charge including reservations (peak_charged_bytes), and aggregate protection counts. The two budget fields are None for pageable host in this example, while byte accounting remains present. CUDA location snapshots and transition reports also sample process-wide PyTorch allocator metrics where applicable. A transition report identifies the CUDA device represented by its allocator sample.

The example prints manager accounting beside:

torch.cuda.memory_allocated(engine.device)
torch.cuda.memory_reserved(engine.device)

These numbers answer different questions:

Metric	Meaning
Logical payload	Sum of declared tensor elements; shared/viewed fields still count logically.
Managed storage	Conservative backing-storage charge admitted by the manager.
Optional residency budget	Application-selected admission limit for one location; None means unbudgeted.
PyTorch allocated	Live allocations known to the process-local caching allocator.
PyTorch reserved	Blocks retained by that allocator, including reusable free blocks.
nvidia-smi / NVML	Broader device/process usage, including contexts and allocations outside manager accounting.

Offloading can lower manager used_bytes without immediately lowering PyTorch reserved bytes or NVML usage. Use the process-wide measurements to set appropriate application budgets and unmanaged headroom.

9. End values and close the manager

residency.discard(source_handle)
residency.discard(key_handle)
residency.discard(weight_handle)
residency.close()

discard ends each managed value. close rejects active or pending lifetimes unless the caller explicitly chooses its force escape hatch.

The evaluator output is not adopted in this example; it remains an ordinary application-owned ciphertext.

Complete runnable source

#!/usr/bin/env python3

"""Execute a named FHE stage with explicit residency plans and CUDA leases.

The example adopts one replicated plaintext weight, one replicated rotation
key, and one exclusive request ciphertext. It then performs host/CUDA
transitions with optional pinned/CUDA budgets, dry-runs a stage plan with a
CUDA reservation, protects asynchronous CUDA readers with a consumer-stream
event, and reports manager and allocator accounting.
"""

from __future__ import annotations

import argparse

import torch
from common import (
    add_engine_args,
    error_stats,
    format_bytes,
    make_engine,
    print_table,
)

from fhelium.residency import (
    PAGEABLE_HOST,
    PINNED_HOST,
    DropResident,
    EnsureResident,
    MemoryReservation,
    MoveResident,
    ReplicaMode,
    ResidencyLocationSnapshot,
    ResidencyManager,
    ResidencyPlan,
    ResidencySnapshot,
    cuda_location,
)


def location_rows(snapshot: ResidencySnapshot) -> list[list[str]]:
    """Format managed storage and optional budgets from one snapshot."""

    return [
        [
            item.location.name,
            format_budget(item),
            format_remaining_budget(item),
            format_bytes(item.used_bytes),
            format_bytes(item.reserved_bytes),
            format_bytes(item.peak_used_bytes),
            format_bytes(item.peak_charged_bytes),
            str(item.use_count),
            str(item.pending_event_count),
            (
                "n/a"
                if item.allocator_allocated_bytes is None
                else format_bytes(item.allocator_allocated_bytes)
            ),
            (
                "n/a"
                if item.allocator_reserved_bytes is None
                else format_bytes(item.allocator_reserved_bytes)
            ),
        ]
        for item in snapshot.locations
    ]


def format_budget(snapshot: ResidencyLocationSnapshot) -> str:
    """Format one optional strict admission budget."""

    if snapshot.budget_bytes is None:
        return "unbudgeted"
    return format_bytes(snapshot.budget_bytes)


def format_remaining_budget(snapshot: ResidencyLocationSnapshot) -> str:
    """Format remaining strict budget or the unbudgeted marker."""

    if snapshot.remaining_budget_bytes is None:
        return "unbudgeted"
    return format_bytes(snapshot.remaining_budget_bytes)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_engine_args(
        parser,
        default_preset="slots8192-scale40-levels7-int64",
        default_device="cuda:0",
    )
    args = parser.parse_args()

    engine = make_engine(args)
    if engine.device.type != "cuda":
        parser.error("this residency example requires CUDA")
    device_location = cuda_location(engine.device)

    positions = torch.arange(engine.num_slots, device="cpu")
    message = 0.01 * torch.sin(positions.to(torch.float64) * 0.017)
    weight_message = torch.full(
        (engine.num_slots,),
        0.5,
        dtype=torch.float64,
    )

    # Construct live exact FHElium values before transferring their logical
    # ownership to the manager.
    weight = engine.prepare_plaintext_for_multiplication(
        engine.encode(weight_message, level=0),
        modulus_basis="Q",
    ).cpu()
    rotation_key = engine.create_rotation_key(1, engine.secret_key).cpu()
    source = engine.encrypt_message(message)

    weight_bytes = max(weight.nbytes, weight.storage_nbytes)
    key_bytes = max(rotation_key.nbytes, rotation_key.storage_nbytes)
    ciphertext_logical_bytes = source.nbytes
    ciphertext_storage_bytes = source.storage_nbytes
    ciphertext_charge = max(
        ciphertext_logical_bytes,
        ciphertext_storage_bytes,
    )
    workspace_bytes = 3 * ciphertext_charge

    # Pageable host remains unbudgeted. Pinned host and CUDA use strict
    # application-selected admission budgets; the CUDA budget includes the
    # plan's reservation for unmanaged evaluator expansion.
    residency = ResidencyManager(
        budgets={
            PINNED_HOST: weight_bytes + ciphertext_charge,
            device_location: (
                weight_bytes + key_bytes + ciphertext_charge + workspace_bytes
            ),
        },
    )
    weight_handle = residency.adopt(
        weight,
        at=PAGEABLE_HOST,
        replica_mode=ReplicaMode.REPLICABLE,
    )
    key_handle = residency.adopt(
        rotation_key,
        at=PAGEABLE_HOST,
        replica_mode=ReplicaMode.REPLICABLE,
    )
    source_handle = residency.adopt(
        source,
        at=device_location,
        replica_mode=ReplicaMode.EXCLUSIVE,
    )
    # adopt() transfers logical alias ownership under caller-enforced rules. Concrete aliases
    # must not be retained or used outside a residency lease.
    del weight, rotation_key, source

    # Direct primitive transitions are application decisions. ensure() keeps a
    # replica; move() changes the sole location of an EXCLUSIVE ciphertext.
    residency.ensure(weight_handle, PINNED_HOST)
    residency.move(
        source_handle,
        PINNED_HOST,
        from_location=device_location,
    )

    # The plan name expresses this application's stage and tile. The plan is
    # ordered low-level IR, and the reservation is accounted headroom for
    # unmanaged evaluator outputs/workspace rather than a tensor allocation.
    plan = ResidencyPlan(
        name="inference/rotate-scale/tile-0",
        enter=(
            EnsureResident(weight_handle, device_location),
            EnsureResident(key_handle, device_location),
            MoveResident(
                source_handle,
                device_location,
                from_location=PINNED_HOST,
            ),
        ),
        exit=(
            MoveResident(
                source_handle,
                PINNED_HOST,
                from_location=device_location,
            ),
            DropResident(key_handle, device_location),
            DropResident(weight_handle, device_location),
        ),
        reservations=(
            MemoryReservation(
                device_location,
                workspace_bytes,
                label="rotate/multiply outputs and workspace",
            ),
        ),
    )
    explanation = residency.explain(plan)
    if not explanation.feasible:
        raise RuntimeError(
            f"residency plan is infeasible: {explanation.reason}"
        )

    compute_stream = torch.cuda.Stream(device=engine.device)
    scope = residency.scope(plan)
    with scope:
        with (
            residency.acquire(
                (source_handle, key_handle, weight_handle),
                at=device_location,
                consumer_stream=compute_stream,
            ) as resident,
            torch.cuda.stream(compute_stream),
        ):
            rotated = engine.rotate_with_key(
                resident[source_handle],
                resident[key_handle],
            )
            output = engine.multiply_plaintext(
                engine.coefficient_domain_to_ntt_domain(rotated),
                resident[weight_handle],
            )
            output = engine.rescale_to_next_level(
                engine.ntt_domain_to_coefficient_domain(output)
            )

        # Lease release records a completion event on compute_stream. The
        # manager retains any unfinished CUDA readers without synchronizing the
        # whole device. This snapshot may already reap a fast completed event.
        released_snapshot = residency.snapshot()

        # Synchronize only the result-producing stream before host decryption.
        # The plan's exit actions are therefore safe when the scope closes.
        compute_stream.synchronize()
        del rotated

    if scope.report is None:
        raise RuntimeError("residency plan scope did not produce a report")

    expected = 0.5 * torch.roll(message, shifts=1)
    actual = engine.decrypt_message(output, is_real=True)
    error = error_stats(actual, expected)
    final_snapshot = residency.snapshot()

    print(f"Plan: {explanation.plan_name}")
    print_table(
        ["location", "predicted managed peak"],
        [
            [location.name, format_bytes(nbytes)]
            for location, nbytes in explanation.predicted_peak_bytes
        ],
    )
    print("\nAccounting after the CUDA lease was released:")
    print_table(
        [
            "location",
            "budget",
            "remaining",
            "used",
            "reserved",
            "peak used",
            "peak charged",
            "uses",
            "pending events",
            "torch allocated",
            "torch reserved",
        ],
        location_rows(released_snapshot),
    )
    print("\nAccounting after stage exit:")
    print_table(
        [
            "location",
            "budget",
            "remaining",
            "used",
            "reserved",
            "peak used",
            "peak charged",
            "uses",
            "pending events",
            "torch allocated",
            "torch reserved",
        ],
        location_rows(final_snapshot),
    )
    print(
        "\nSource logical payload: "
        f"{format_bytes(ciphertext_logical_bytes)}; "
        "source unique storage: "
        f"{format_bytes(ciphertext_storage_bytes)}; "
        "managed source charge: "
        f"{format_bytes(ciphertext_charge)}; "
        f"torch allocated: {format_bytes(torch.cuda.memory_allocated(engine.device))}; "
        f"torch reserved: {format_bytes(torch.cuda.memory_reserved(engine.device))}; "
        f"plan transitions: {len(scope.report.transitions)}; "
        f"max error: {error['max_abs']:.3e}"
    )

    # End the managed logical values, then close the manager. The
    # output is an ordinary unmanaged ciphertext owned by this application.
    residency.discard(source_handle)
    residency.discard(key_handle)
    residency.discard(weight_handle)
    residency.close()

    if error["max_abs"] > 1e-5:
        raise RuntimeError(f"residency example error too large: {error}")


if __name__ == "__main__":
    main()

Related concepts and guides

• Residency lifetimes
• Ownership and runtime responsibilities
• Stream resources with bounded memory