Nott

Nott is a ~35k-LoC C++20 framework that layers a strongly-typed, DAG-based API over LibTorch and compiles it once into a flat execution plan: mean step latency within a few percent of hand-written LibTorch, step-time jitter cut by up to two orders of magnitude, and CUDA graph capture, tensor-core paths and NHWC/NCHW memory formats kept explicit throughout.

Most ANN frameworks in C++ force a pick: LibTorch's raw module API, correct but requiring a register_module call and a hand-written forward() per network, or a header-only library built for classical ML (dlib, mlpack) that never touches deep learning. Neither closes the gap between optimized code and syntax you'd actually want to write. After years of building networks in C++/CUDA from scratch, re-growing the same staging, AMP and logging scaffolding each project with each copy carrying its own bugs, I wanted the terse call site and the optimized hot path to be the same artifact, not a scripting layer bolted on top.

That's the shape Nott is built around: every layer, loss, optimizer and scheduler is a builtin with the same call shape, so a network reads as a flat list of descriptors instead of a hand-rolled module tree, and training is one model.train(...) call instead of a loop rewritten per project. The "obvious" way to hand-write that loop (.pin_memory() wherever it seems needed, .item() the loss every step, std::function for anything callable) quietly costs redundant copies, hidden device syncs, and heap allocations on the hot path: the exact costs a builtin only pays once.

So I built Nott: a C++20 framework (~35k LoC) layering a strongly typed, descriptor-driven API over LibTorch. You describe a network as a DAG of typed descriptors; Nott compiles that description once into a flat execution plan, and the per-step hot path touches none of the machinery that produced it.

#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});

Same network, both APIs, side by side:

LibTorch: register_module calls, five conv layers of repeated options, a hand-written forward() chaining relu/pool in order. Nott: a flat list of model.add(...) calls, no forward() to write or get wrong, every option named instead of positional. Same compute, no .item()-per-step or pin_memory()-wherever footguns to avoid by hand.

What follows is the path from that call to the CUDA kernel, and why each choice was made.

01How it works

The architecture splits into a cold path (graph authoring, descriptor resolution, module registration, runs once) and a hot path (per-batch fetch -> stage -> forward -> loss -> backward -> step, runs hundreds of thousands of times). Organizing principle: all polymorphism resolves on the cold path; the hot path is flat function-pointer calls over pre-staged buffers.

Descriptors: pay for dispatch once, at build time

Every layer, optimizer, loss, and regularizer has its own Options struct plus a Descriptor bundling it with activation, initialization, and per-layer local config:

struct FCOptions {
    std::int64_t in_features{};
    std::int64_t out_features{};
    bool bias{true};
};

struct FCDescriptor {
    FCOptions options;
    Activation::Descriptor activation{Activation::Identity};
    Initialization::Descriptor initialization{Initialization::Default};
    LocalConfig local{};
};

Model::add() takes a std::variant over all descriptor types, std::visit-ed exactly once, when the graph compiles into RegisteredLayer entries. C++20 designated initializers ({.learning_rate = 1e-3, .weight_decay = 5e-4}) give Python-kwarg ergonomics with compile-time field checking: misspell a field and it's a build error, not a silently ignored dict key.

Type-erased forward dispatch without std::function

The natural way to store a Tensor-to-Tensor callable is std::function<torch::Tensor(torch::Tensor)>. Wrong on a hot path: potential heap allocation for captured state, double indirection through its vtable-like on every call. Nott's RegisteredLayer uses a hand-rolled binding instead, one raw function pointer plus one 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));
}

Not every LibTorch module fits directly, an RNN's forward() returns an (output, hidden_state) tuple, so an adapter functor unwraps it: a plain struct with operator(), whose entire state is visible member fields rather than opaque lambda captures:

struct ForwardFunctor {
    torch::nn::RNNImpl* module_ptr; // the entire state
    torch::Tensor operator()(torch::Tensor input) const {
        auto out = module_ptr->forward(std::move(input));
        return Detail::take_recurrent_output(out); // drop the hidden state
    }
};
registered_layer.bind_inline_forward(ForwardFunctor{module.get()});

Because the state is one trivially copyable pointer, adapters get an inline path: bind_inline_forward() memcpys the functor into alignas(std::max_align_t) std::byte storage[3 * sizeof(void*)] embedded in the layer, small-buffer optimization with size/triviality enforced by static_assert rather than discovered at runtime. An adapter capturing a std::string or shared_ptr fails to compile instead of corrupting memory. Copy/move constructors re-point forward.context at the new object's storage. Executing the graph is a linear walk over a contiguous std::vector<RegisteredLayer>: load two pointers, one indirect call, done.

The data path: pin once, stage into stable buffers, double-buffer the copies

The dataset pipeline is pure functions with one rule each. ensure_pinned() pins host memory exactly once, at the CPU-to-GPU boundary before training starts (the previous incarnation called pin_memory() in four different places). ensure_memory_format() only makes a tensor contiguous (or ChannelsLast for >=4-D inputs) if it isn't already, killing a class of silent GPU-to-GPU copies the old code triggered.

Staging to device goes through a reusable buffer with a stability flag:

// 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);

In steady state, fixed batch shape (the common case), every step is a single async copy_ into memory that already exists: no allocator traffic, no fragmentation, and the copy overlaps compute via a double-buffered prefetcher on a dedicated CUDA stream.

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{};
};

The layout is deliberate: per-iteration boolean flags share a 64-byte cache line at the front, while tensors, stream, and events (touched once per transfer) sit behind. The loop is a ping-pong: process slot s, schedule the H2D copy for the next batch into slot s ^ 1 on the side stream, record a CUDAEvent, have the compute stream wait() on it instead of a global sync. Batch i+1 streams over PCIe while batch i trains.

There are three iteration strategies (prefetch, buffered CPU-side deque with a buffer_vram + 1 look-ahead, plain loop), selected at runtime from a single TrainingPolicy struct. The previous design encoded these as template booleans, metastasized into 12 instantiated code paths; collapsing to 3 runtime-dispatched functions traded a perfectly-predicted branch for a 4x reduction in paths to test.

One more sync killer: loss logging. Calling .item() on a GPU loss tensor forces a device sync every step, purely so a progress bar can print a float. Nott's DeferredScalar issues a non-blocking D2H copy, records a CUDAEvent, and only materializes the value (polling event->query() first, syncing only if genuinely not ready) when something actually reads it. The stall becomes overlap.

CUDA graph capture as a state machine

For small models, per-step latency is dominated by the CPU launching kernels, not the kernels themselves. Nott exposes LibTorch's CUDA graph support behind a GraphMode {Disabled, Capture, Replay} enum, coordinated by a small state machine (NeverCaptured -> Pending -> Ready). The coordinator hashes each batch against a stored shape/dtype/device signature; a match replays the captured graph (one launch replays the entire step's kernel sequence), a mismatch resets and re-captures. Graph mode also tightens error handling: a regularization penalty changing device or dtype mid-run is silently coerced in eager mode but throws under capture, since static shapes are the contract that makes replay legal.

02Trade-offs

03Benchmarks

Methodology: MNIST (60k samples, 28x28), 100 epochs at batch 64, warm runs only. The first 100 to 1000 steps are discarded, then samples are filtered with a Tukey fence (k = 0.98) to strip OS-scheduling and thermal outliers. Timing is std::chrono::high_resolution_clock around the full step, with 70k to 91k retained steps per runner. Three runners: Nott's prebuilt Train(), a custom training loop on Nott's primitives, and a raw hand-written LibTorch loop. Full harness: `examples/speedtest.cpp`.

Unified I/O (async pinned memory enabled for all three runners: the fair fight):

The abstraction costs +3.5% mean latency against hand-written LibTorch, and buys a coefficient of variation ~3x lower than the raw loop and ~10x lower than the custom loop. The P98/P50 spread for the prebuilt runner is 0.4%.

Mixed I/O (async pinned memory only in Nott's Train(): deliberately favorable, isolating what the I/O configuration alone is worth):

Here the prebuilt loop is 5.7% faster than raw LibTorch with ~100x less jitter, a measurement of how much a naive data path costs, not of Nott beating the backend it sits on. The wrapper's tax is a few percent; the data-path discipline it enforces is worth more.

04Where it went sideways

The most instructive bugs were the ones the refactor surfaced rather than introduced. The original training loop had grown organically, and pin_memory() ended up called at four separate points, each a no-op when already pinned, except in one path a fresh contiguity pass triggered a silent GPU-to-GPU copy. Nothing was wrong in the correctness sense; the profiler just showed device-to-device copy_ kernels with no business existing. The fix was structural: dataset_pipeline.hpp is now the single point of truth, every helper is a pure function, pinning happens at exactly one boundary.

The second was the step executor. graph_train_step_impl() had become a 230-line switch with the actual step body buried in an inner lambda, and under CUDA graph capture that body had different legality rules than in eager mode. A regularization penalty whose dtype drifted (a float64 accumulator meeting a float16 AMP loss) would be silently .to()-converted in eager mode and corrupt a captured graph in replay mode. Extracting the step body into a standalone template made the divergence explicit: eager mode coerces, graph mode throws naming the exact violated invariant. When two execution modes share code, the shared code must know which contract it's operating under, or the stricter mode fails in the least debuggable way possible.

05What I'd do better next time

My starting mental model, "the wrapper's job is to not add latency," aimed at the wrong axis. The mean was never the problem; hand-written LibTorch loops already hit a good mean. The real constraint was variance, and variance comes from exactly the places a framework can own outright: allocation on the hot path, hidden synchronization points, I/O that isn't overlapped with compute. The prebuilt Train() beats loops written by its own author not because it's cleverer per-step, but because it's incapable of the small inconsistencies a hand-rolled loop accumulates.

The next iteration follows from that. The shape-signature machinery built for CUDA graph re-capture is most of what's needed for multi-shape graph caching: keep N captured graphs keyed by signature instead of re-capturing on every change. The DeferredScalar pattern wants to grow into a general deferred-metrics pipeline so evaluation never syncs either. And the latency harness itself, Tukey-fenced distributions over 80k steps, turned out to be the most reusable artifact of the whole project: it's what made every claim above checkable, and what I'd now build first.

06Further reading

• Nott on GitHub: source, full docs (docs/), and the benchmark harness.
• CUDA Graphs in PyTorch: the capture/replay mechanism Nott's GraphModeCoordinator orchestrates.
• How to Optimize Data Transfers in CUDA C/C++: pinned memory and async copy overlap, the basis of the prefetch design.
• Nott docs, Research References: per-module citations for every optimizer, loss, and transformer variant implemented (Muon, Sophia, Mamba, SWAG, and ~60 others).