Build an optimization workflow you can audit.

Usage

Optimize the serving distribution.

A single signature remains supported. For production selection, supply weighted cases for the shapes, batches, dtypes, and keyword forms that actually receive traffic.

from custom_dl_optimizer import (
    OptimizationConfig, Optimizer,
    WorkloadCase, WorkloadProfile,
)

optimizer = Optimizer(
    device="cuda",
    config=OptimizationConfig(
        enable_compile=True,
        plan_cache_dir=".custom-dl-cache",
        min_speedup=1.02,
    ),
)
profile = WorkloadProfile(
    name="serving",
    expected_calls=10_000,
    cases=(
        WorkloadCase("b1", args=(batch_1,), weight=80),
        WorkloadCase("b8", args=(batch_8,), weight=20),
    ),
)
result = optimizer.optimize_workload(model.eval(), profile)
result.save_bundle("artifacts/optimization")

Read the decision

Each candidate records per-case first calls and parity, serial P50/P90/P95/P99 samples, mean confidence bounds, incremental CUDA allocation, constraints, projected total time, break-even calls, and failures. Cached winners are revalidated before reuse.

Correctness boundary

Candidates must preserve output structure and satisfy torch.allclose using configured tolerances. Always run task-level accuracy evaluation before deployment.

Candidate providers

Compare maintained runtimes under the same policy.

First-party providers cover Torch-TensorRT, ONNX Runtime, and TorchAO. Private backends can still implement the provider protocol.

from custom_dl_optimizer import (
    ONNXRuntimeProvider, Optimizer,
    TorchAOQuantizationProvider,
    TorchTensorRTProvider,
)

result = Optimizer(
    device="cuda",
    providers=(
        TorchTensorRTProvider(),
        ONNXRuntimeProvider(),
        TorchAOQuantizationProvider(),
    ),
).optimize(model, sample)

Fair comparison

Keep transfers and precision consistent, separate engine-build cost from steady-state latency, warm caches, expose provider configuration, and never silently run another backend under the same provider name.

Agent toolkit

Register objects locally; expose names remotely.

Live modules and tensors stay in the host process. The toolkit exposes only runtime inspection, workload listing, optimization, and report retrieval.

from custom_dl_optimizer import OptimizationAgentToolkit, Optimizer

toolkit = OptimizationAgentToolkit(Optimizer(device="cuda"))
toolkit.register_workload(
    "encoder-b16",
    encoder,
    tokens,
    attention_mask=mask,
)

schemas = toolkit.tool_schemas()
report = toolkit.invoke(
    "custom_dl_optimize",
    {"workload": "encoder-b16"},
)

Security boundary

There is no tool argument for a file path, import, download, shell command, Python expression, model weight, or credential. The host remains responsible for deciding which workload names an agent may access.

Migration

Move from a tuple to one result object.

Version 1 constructed AutoOptimizer around a model and returned a module/report tuple. Version 2 constructs reusable optimizer policy and returns OptimizationResult.

# Version 1
engine = AutoOptimizer(model, config=config)
module, report = engine.optimize_with_report(sample)

# Version 2
optimizer = Optimizer(config=config)
result = optimizer.optimize(model, sample)
module = result.module
report = result.report

AutoOptimizer remains available for one transition release. New integrations should use the version 2 result contract.