Nott

Nott is a ~35k-LoC C++20 framework: a typed, DAG-based API over LibTorch, compiled once into a flat execution plan. Mean step latency within a few percent of hand-written LibTorch. Step-time jitter cut up to two orders of magnitude. CUDA graph capture, tensor-core paths and NHWC/NCHW formats stay explicit.

Nott: layers, losses, optimizers and schedulers are builtins with one call shape. No register_module, no hand-written forward(), no per-project training loop.

#include <Nott>

Nott::Model model("demo");
model.use_cuda(torch::cuda::is_available());

model.add(Nott::Layer::FC({784, 256, true}, Nott::Activation::GeLU));
model.add(Nott::Layer::Dropout({0.1}));
model.add(Nott::Layer::FC({256, 10, true}, Nott::Activation::Softmax));

model.set_optimizer(Nott::Optimizer::AdamW({.learning_rate = 1e-3}),
                    Nott::LrScheduler::CosineAnnealing({.T_max = 50}));
model.set_loss(Nott::Loss::MSE({}));
model.train(train_x, train_y, {.epoch = 10, .batch_size = 64});

01How it works

Cold path: graph authoring, descriptor resolution, module registration. Runs once. Hot path: fetch -> stage -> forward -> loss -> backward -> step. Runs 10^5+ times. All polymorphism resolves cold. The hot path is flat function-pointer calls over pre-staged buffers.

Model::add() takes a std::variant over all descriptors; std::visit runs once at graph compile, producing RegisteredLayer entries. A misspelled designated initializer is a build error.

Forward dispatch without std::function

std::function: possible heap allocation, double indirection per call. RegisteredLayer stores a raw function pointer plus a context pointer.

struct ForwardBinding {
    using Invoker = torch::Tensor (*)(void*, torch::Tensor);
    Invoker invoke{nullptr};
    void*   context{nullptr};
};

template <class Module>
static torch::Tensor dispatch_module(void* context, torch::Tensor input) {
    return static_cast<Module*>(context)->forward(std::move(input));
}

Other signatures (RNN returns (output, hidden)) bind an adapter functor, memcpyd into alignas(std::max_align_t) std::byte storage[3 * sizeof(void*)] in the layer. Trivially copyable, <= 3 pointers, static_assert-enforced.

Execution walks a contiguous std::vector<RegisteredLayer>: two loads, one indirect call.

Data path

ensure_pinned() pins host memory once, at the CPU-to-GPU boundary. ensure_memory_format() makes a tensor contiguous (ChannelsLast for >=4-D) only if needed.

// stage_to_device(): reuse the device-side buffer across steps when
// shape/dtype/layout are unchanged; allocate only on signature change.
if (buffer_stable && buffer.defined() && !buffer.sizes().equals(tensor.sizes()))
    buffer_stable = false;

if (!buffer_stable) {
    buffer = torch::empty(tensor.sizes(), options, fmt);   // cold path
} else {
    buffer_stable = true;                                  // steady state
}
buffer.copy_(tensor, /*non_blocking=*/true);

Fixed batch shape is the common case: no allocator traffic after step one.

struct alignas(64) PrefetchState {
    // Hot: read every iteration, packed into one cache line
    std::array<bool, 2> pending{false, false};
    std::array<bool, 2> input_stable{false, false};
    std::array<bool, 2> target_stable{false, false};

    // Cold: touched once per transfer
    torch::cuda::CUDAStream stream;
    std::array<torch::Tensor, 2>       inputs{};
    std::array<torch::Tensor, 2>       targets{};
    std::array<at::cuda::CUDAEvent, 2> events{};
};

Ping-pong: process slot s, queue the next H2D copy into s ^ 1 on a side stream, record a CUDAEvent, compute stream wait()s on it. No global sync.

Three iteration strategies (prefetch, buffered CPU deque with buffer_vram + 1 look-ahead, plain loop), chosen at runtime from one TrainingPolicy. Template booleans gave 12 instantiated paths; runtime dispatch costs a predicted branch, 4x fewer paths.

.item() on a GPU loss syncs the device every step to print a float. DeferredScalar copies D2H non-blocking, records an event, materializes on read.

CUDA graph capture

Small models are CPU-launch-bound. GraphMode {Disabled, Capture, Replay} sits behind a NeverCaptured -> Pending -> Ready state machine. Each batch hashes against a stored shape/dtype/device signature: match replays, mismatch re-captures. Static shapes make replay legal, so a penalty changing device or dtype mid-run throws under capture, where eager coerces.

02Trade-offs

03Benchmarks

MNIST 60k, 28x28. 100 epochs, batch 64, warm runs. First 100-1000 steps dropped, then a Tukey fence (k = 0.98). 70k-91k steps retained per runner. Harness: `examples/speedtest.cpp`.

Unified I/O, async pinned memory for all three runners:

Abstraction cost: +3.5% mean latency. CV ~3x lower than the raw loop, ~10x lower than the custom loop. Prebuilt P98/P50 spread: 0.4%.

Mixed I/O, async pinned memory only in Train():

5.7% faster than raw LibTorch, ~100x less jitter. Measures a naive data path's cost, not Nott beating its backend.

04Further reading

• Nott on GitHub: source, docs, harness.
• CUDA Graphs in PyTorch
• How to Optimize Data Transfers in CUDA C/C++
• Nott docs, Research References: ~60 per-module citations.