CUDA Graph matrix-vector multiplication

Example source: examples/11_cuda_graph_matrix_vector.py

This example captures a fixed packed CKKS matrix-vector evaluator $\mathbf{y}=A\mathbf{x}$ while keeping each request ciphertext as a dynamic input. The tutorial explains captured state, replay storage, and the steady-state measured region.

Run the example

python examples/11_cuda_graph_matrix_vector.py \
  --device cuda:0 \
  --preset slots8192-scale40-levels7-int64 \
  --size 8 \
  --capture-warmup 3 \
  --benchmark-warmup 10 \
  --runs 100

1. Identify static and dynamic state

The workload computes y = A @ x with cyclic diagonals. Its state divides into:

Static across replays	Dynamic for each replay
engine and CKKS context	encrypted input vector
encoded matrix diagonals	ciphertext payload
exact rotation keys	request-specific values
operation schedule	matching exact-value input signature

Encryption and decryption remain outside the graph.

2. Prepare operation-ready constants

diagonals = [
    engine.prepare_plaintext_for_multiplication(
        engine.encode(
            cyclic_diagonal_slots(matrix, step, engine.num_slots),
            level=0,
        )
    )
    for step in range(matrix.size(0))
]
rotation_keys = {
    step: engine.rotation_key(step)
    for step in range(1, matrix.size(0))
}

These values are captured as callable state. Changing their object identity, storage address, level, or shape after capture would invalidate the captured schedule.

3. Bind static state

from functools import partial

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

The resulting callable has one dynamic argument: the source ciphertext. This is preferable to a hidden global cache because the captured resources are visible through direct Python calls.

4. Capture from a prototype

program = CudaGraphProgram.capture(
    schedule,
    example_inputs=(prototype,),
    warmup=capture_warmup,
)

CudaGraphProgram performs side-stream warmup, allocates fixed dynamic-input storage, captures the evaluator, records the output storage, and derives an exact input signature.

The prototype determines structure, including:

• value-tree shape;
• exact value type;
• tensor shape and dtype;
• CKKS context and level;
• polynomial domain, modulus basis, residue representation, scale, and prime IDs.

5. Replay with changing ciphertexts

result = program.replay(
    encrypted,
    synchronize=True,
)

Replay validates the new value before staging its payload into the fixed input allocation. A mismatched level or representation is rejected rather than silently converted inside the graph wrapper.

The example uses three different encrypted vectors and verifies all three against matrix @ vector.

6. Understand borrowed output storage

borrowed_pointer = result.data.data_ptr()

The default result references graph-owned output storage. The next replay can overwrite it. This avoids an extra device copy when the caller consumes the output immediately.

Use:

owned = program.replay(encrypted, copy_output=True)

when a result must survive a later replay or escape into another asynchronous lifetime.

7. Separate construction and steady-state cost

The example reports:

• warmup time;
• graph capture time;
• first replay time;
• eager mean latency;
• graph replay mean latency.

Capture is a one-time program-construction cost. Compare steady-state replay only when the same static schedule will execute enough times to amortize it.

Close the program

program.close()

Closing releases graph-owned inputs, outputs, and capture state. Do not treat a captured program as an unbounded global singleton when different contexts, models, or input signatures require independent storage.

Application-owned serving policy

A serving extension assigns models and users, admits requests, manages per-user keys and eviction, and composes those policies around one or more fixed captured programs.

Complete runnable source

#!/usr/bin/env python3

"""Capture and replay packed CKKS matrix-vector multiplication.

The fixed evaluator computes ``y = A @ x`` with the cyclic-diagonal method.
The matrix, encoded diagonals, rotation keys, and operation schedule are static;
``functools.partial`` binds them into the captured callable.
each replay stages a newly encrypted vector into the program's dynamic input
buffer. Encryption and decryption remain outside capture.
"""

from __future__ import annotations

import argparse
from functools import partial

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

from fhelium import Ciphertext, CkksEngine, Plaintext, RotationKey
from fhelium.execution import CudaGraphProgram


def matrix_and_vector(
    size: int, *, seed: int
) -> tuple[torch.Tensor, torch.Tensor]:
    row = torch.arange(size, dtype=torch.float64).view(-1, 1)
    column = torch.arange(size, dtype=torch.float64).view(1, -1)
    matrix = 0.018 * torch.sin((row + 1) * (column + 2) * 0.17)
    matrix += 0.007 * torch.cos((row + column + 1) * 0.23)
    generator = torch.Generator().manual_seed(seed)
    vector = torch.randn(size, generator=generator, dtype=torch.float64) * 0.025
    return matrix, vector


def periodic_slots(values: torch.Tensor, num_slots: int) -> torch.Tensor:
    return values.repeat(num_slots // values.numel())


def cyclic_diagonal_slots(
    matrix: torch.Tensor,
    rotation_step: int,
    num_slots: int,
) -> torch.Tensor:
    """Return weights aligned with ``torch.roll(x, rotation_step)``."""

    size = matrix.size(0)
    row = torch.arange(num_slots) % size
    column = torch.remainder(row - rotation_step, size)
    return matrix[row, column]


def prepare_constants(
    engine: CkksEngine,
    matrix: torch.Tensor,
) -> tuple[list[Plaintext], dict[int, RotationKey]]:
    diagonals = [
        engine.prepare_plaintext_for_multiplication(
            engine.encode(
                cyclic_diagonal_slots(matrix, step, engine.num_slots), level=0
            )
        )
        for step in range(matrix.size(0))
    ]
    rotation_keys = {
        step: engine.rotation_key(step) for step in range(1, matrix.size(0))
    }
    return diagonals, rotation_keys


def matrix_vector(
    source: Ciphertext,
    *,
    engine: CkksEngine,
    diagonals: list[Plaintext],
    rotation_keys: dict[int, RotationKey],
) -> Ciphertext:
    """Evaluate one vector with statically bound program state."""

    rotated_values = []
    for step in range(len(diagonals)):
        rotated = (
            source
            if step == 0
            else engine.rotate_with_key(source, rotation_keys[step])
        )
        rotated_values.append(rotated)
    rotated_ntt = engine.coefficient_domain_to_ntt_domain(
        Ciphertext.stack_batch(rotated_values)
    )
    diagonal_batch = Plaintext.stack_batch(diagonals)
    weighted = engine.multiply_plaintext(rotated_ntt, diagonal_batch)
    return engine.rescale_to_next_level(
        engine.ntt_domain_to_coefficient_domain(
            engine.sum_ciphertext_batch(weighted)
        )
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_engine_args(
        parser,
        default_preset="slots8192-scale40-levels7-int64",
        default_device="cuda:0",
    )
    parser.add_argument("--size", type=int, default=8)
    parser.add_argument("--capture-warmup", type=int, default=3)
    parser.add_argument("--benchmark-warmup", type=int, default=10)
    parser.add_argument("--runs", type=int, default=100)
    args = parser.parse_args()

    engine = make_engine(args)
    if engine.device.type != "cuda":
        parser.error("this CUDA Graph example requires CUDA")
    if args.size <= 0 or engine.num_slots % args.size != 0:
        parser.error(
            f"--size must be positive and divide num_slots={engine.num_slots}"
        )
    if args.capture_warmup < 0 or args.benchmark_warmup < 0 or args.runs <= 0:
        parser.error("warmup counts must be non-negative and --runs positive")

    matrix, prototype_vector = matrix_and_vector(args.size, seed=500)
    diagonals, rotation_keys = prepare_constants(engine, matrix)
    schedule = partial(
        matrix_vector,
        engine=engine,
        diagonals=diagonals,
        rotation_keys=rotation_keys,
    )
    prototype = engine.encrypt_message(
        periodic_slots(prototype_vector, engine.num_slots)
    )
    program = CudaGraphProgram.capture(
        schedule,
        example_inputs=(prototype,),
        warmup=args.capture_warmup,
    )

    correctness_rows = []
    borrowed_pointer = None
    for replay_index, seed in enumerate((501, 502, 503)):
        _, vector = matrix_and_vector(args.size, seed=seed)
        encrypted = engine.encrypt_message(
            periodic_slots(vector, engine.num_slots)
        )
        result = program.replay(encrypted, synchronize=True)
        borrowed_pointer = result.data.data_ptr()
        actual = engine.decrypt_message(result, is_real=True)[: args.size]
        error = error_stats(actual, matrix @ vector)
        correctness_rows.append(
            [
                replay_index,
                f"{error['max_abs']:.3e}",
                f"0x{borrowed_pointer:x}",
            ]
        )

    print("Changing-input matrix-vector replay correctness:")
    print_table(
        ["replay", "max abs error", "borrowed output"],
        correctness_rows,
    )

    eager_stats, _ = time_ms(
        lambda: schedule(prototype),
        warmup=args.benchmark_warmup,
        runs=args.runs,
        device=engine.device,
    )
    graph_stats, _ = time_ms(
        lambda: program.replay(prototype),
        warmup=args.benchmark_warmup,
        runs=args.runs,
        device=engine.device,
    )
    print("\nFixed matrix-vector evaluator latency:")
    print_table(
        ["mode", "mean ms", "relative"],
        [
            ["eager", f"{eager_stats['mean_ms']:.4f}", "1.000x"],
            [
                "CUDA Graph",
                f"{graph_stats['mean_ms']:.4f}",
                f"{eager_stats['mean_ms'] / graph_stats['mean_ms']:.3f}x",
            ],
        ],
    )

    stats = program.stats
    print(
        "\nOne-time construction: "
        f"warmup={stats.warmup_seconds:.3f}s, "
        f"capture={stats.capture_seconds:.3f}s, "
        f"first_replay={stats.first_replay_seconds:.6f}s"
    )
    print(
        "Replay returns borrowed storage at "
        f"0x{borrowed_pointer:x}; pass copy_output=True when a result must "
        "survive the next replay."
    )
    program.close()


if __name__ == "__main__":
    main()

Related concepts and guides

• CUDA Graph execution model
• Exact signatures and buffers
• Capture a repeated evaluator