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, which is correct but means 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 at all. Neither closes the gap between optimized code and syntax you'd actually want to write by hand. 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 a framework where the terse call site and the optimized hot path were the same artifact, not a scripting layer bolted on top of one.

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 (call .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 has to pay once.

So I built Nott: a C++20 framework (~35k LoC) that layers a strongly typed, descriptor-driven API over LibTorch (PyTorch backend). 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:

The LibTorch side is register_module calls, five conv layers of repeated options, and a hand-written forward() chaining relu/pool calls in order. The Nott side is a flat list of model.add(...) calls, no forward() to write or get wrong, and every option named instead of positional. That's the whole pitch: same compute, no .item()-per-step or pin_memory()-wherever-it-seems-needed footguns to avoid by hand.

What follows is the path between that call and 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 that runs once) and a hot path (per-batch fetch -> stage -> forward -> loss -> backward -> step that runs hundreds of thousands of times). The organizing principle: all polymorphism is resolved 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 that bundles 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() (to add a new layer inside the network) takes a std::variant over all descriptor types. That variant is std::visit-ed exactly once, when the graph is compiled into RegisteredLayer entries. C++20 designated initializers ({.learning_rate = 1e-3, .weight_decay = 5e-4}) give the ergonomics of Python kwargs 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 something callable that maps Tensor-to-Tensor is std::function<torch::Tensor(torch::Tensor)>. It's also the wrong way on a hot path: potential heap allocation for the captured state, and a double indirection through its internal 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 that binding directly, an RNN's forward() returns an (output, hidden_state) tuple, so something has to unwrap it before the next layer. That something is an adapter functor: a plain struct with an operator(), whose entire state is visible member fields (here, a single module pointer) rather than the opaque captures of a lambda:

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, these 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 itself, small-buffer optimization with the size and triviality requirements enforced by static_assert rather than discovered at runtime. Write an adapter that captures a std::string or a shared_ptr and it fails to compile instead of corrupting memory. Copy/move constructors re-point forward.context at the new object's storage, because a memcpy'd context pointer into a moved-from buffer is exactly the kind of bug type erasure invites. Executing the graph is then 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 built from pure functions with one rule each. ensure_pinned() pins host memory exactly once, at the CPU-to-GPU boundary before training starts, the doc comment is candid that 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, which is the overwhelmingly common case, every step is a single async copy_ into memory that already exists. No allocator traffic, no fragmentation, and the copy overlaps compute because of the third piece: 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: the per-iteration boolean flags share a 64-byte cache line at the front of the struct, while the tensors, stream, and events (each touched once per transfer) sit behind them. The loop itself is the classic ping-pong: process slot s, schedule the H2D copy for the next batch into slot s ^ 1 on the side stream, record a CUDAEvent, and have the compute stream wait() on that event instead of a global sync. Batch i+1 streams over PCIe while batch i trains.

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

One more sync killer: loss logging. Calling .item() on a GPU loss tensor forces a device synchronization every step: a stall that exists 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 is converted into overlap.

CUDA graph capture as a state machine

For small models, per-step latency is dominated not by kernels but by the CPU launching them. Nott exposes LibTorch's CUDA graph support behind a GraphMode {Disabled, Capture, Replay} enum, coordinated by a small per-run state machine (NeverCaptured -> Pending -> Ready). The coordinator hashes each batch against a stored shape/dtype/device signature; a match means replay the captured graph (a single launch replays the entire step's kernel sequence), and a mismatch resets the model's graph cache 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, because static shapes are the contract that makes replay legal.

02Trade-offs

03Benchmarks

Methodology first, since the intro complained about untrustworthy numbers: 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, which is enough to report stable percentiles rather than a mean and a prayer. Three runners: Nott's prebuilt Train(), a custom training loop written 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, which is really a measurement of how much a naive data path costs, not of Nott beating the backend it sits on. Both readings matter: the wrapper's tax is a few percent; the data-path discipline it enforces is worth more than that tax.

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() had ended up called at four separate points, each a no-op when the tensor was already pinned, except for the pinning check itself and, in one path, a fresh contiguity pass that triggered a silent GPU-to-GPU copy. Nothing was wrong in the correctness sense; the profiler just showed device-to-device copy_ kernels that had no business existing. The fix was structural, not local: dataset_pipeline.hpp now exists as the single point of truth, every helper is a pure function, and 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 with a message naming the exact violated invariant. The lesson generalizes: when two execution modes share code, the shared code must know which contract it's operating under, or the stricter mode will fail 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," was 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, and 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 the instrument that made every claim above checkable, and it's 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).