Research software

Framework-Free C++/CUDA Chronos and Moirai-MoE

Language: C++20 · CUDA · Python
Scale: ~28,600 lines of C++/CUDA · 111 device kernels
Source: Not publicly released
Status: Pre-release · fidelity suite partly run

Overview

Two pretrained time-series foundation-model families are reimplemented from scratch: Chronos, a T5 encoder–decoder over quantized value tokens, and Moirai-MoE, a sparse mixture-of-experts any-variate encoder. No tensor framework is linked into the runtime — the only dependency is the CUDA runtime and optionally cuBLAS, with no libtorch, ONNX Runtime or TensorRT at any point. Weights are converted ahead of time into a self-describing artifact that the C++ loader memory-maps by parameter name, so the runtime never sees an upstream checkpoint format. Almost every forward kernel has a matching backward kernel, which means these models fine-tune natively rather than only running inference.

The two families were chosen because they contrast on the axes that matter for a reimplementation study: discrete value tokenization against continuous patching, encoder–decoder against encoder-only, and dense against sparse feed-forward. They therefore stress different parts of the kernel library — mean-scale quantization and T5 relative attention bias on one side, any-variate patching, partial rotary embedding, SwiGLU experts and centroid-distance top-k routing on the other.

The harder problem is not writing the kernels but establishing that they are right, and that is what most of the project is. Bitwise equality with the reference is neither achievable nor checkable: different reduction orders, split-K decompositions, fused epilogues that keep intermediates in registers, and TF32 defaults on recent hardware all guarantee divergence in the last bits. So the useful question is not whether two outputs are identical but whether they agree more closely than any genuine implementation difference could produce — and every claim is scoped to what was actually compared, with each tolerance justified by the error mode it excludes rather than by the result it permits.

Implementation

Families: Chronos — a T5 encoder–decoder over quantized value tokens; Moirai-MoE — a sparse mixture-of-experts any-variate encoder
Checkpoints: Four, from 46.2M to 935.7M parameters and 132 to 1,286 weight tensors, all at a 512-step context
Runtime dependencies: The CUDA runtime, optionally cuBLAS. No libtorch, ONNX Runtime or TensorRT is linked
Kernels: 111 device kernels, with a matching *_backward_kernel for almost every forward one
Build: CMake ≥ 3.24 and C++20, with a CUDA release preset and a CPU stub preset that compiles without a toolkit and reports a non-CUDA execution class so it cannot be mistaken for evidence
Weight conversion: safetensors, PyTorch state dicts and Orbax-PAX msgpack parsed with NumPy and repacked into a memory-mapped artifact
Verification: Per-tensor comparison against the reference implementation on converted checkpoint weights and real inputs, recording error distributions rather than a pass/fail verdict

Components

src/core/: Status and result plumbing, hashing, logging, and the artifact store
src/pretrained/backbone_artifact.cpp: The artifact loader and the parameter-name binding the runtime resolves weights through
src/pretrained/full_cuda_transformer_block.cpp: The shared dense transformer substrate, with family-specific behaviour parameterized rather than duplicated
src/pretrained/full_cuda_graph_backbones.cpp: The two model graphs — the largest single source file in the tree
cuda/kernels/cuda_runtime_kernels.cu: The device kernel library: T5 relative attention bias, partial rotary embedding, RMSNorm, SwiGLU, mean-scale quantization, patch extraction and projection, expert routing, dispatch, GEMM and combine
cuda/kernels/cuda_runtime_kernels_stub.cpp: The toolkit-free stub build, which reports its degraded execution class rather than pretending to run
cuda/memory/, cuda/ops/: Device allocation and op dispatch
include/cuda/cuda_backward_plan.hpp: The backward composition, mirroring the forward graph
include/pretrained/backbone_training_workspace.hpp: Saved activations and the per-parameter gradient buffers the backward pass fills
apps/parity_probe_main.cpp: The native executable the fidelity suites drive
python/tsfm_native/convert/: Checkpoint readers, parameter mapping, dtype policy, claim policy, and the conversion CLI
tests/fidelity/: The comparison runners, the comparator, the tolerance table, and the gate-matrix reporter
docs/: Architecture, conversion, claim policy, input contracts, benchmark protocol, and porting notes
paper/: The manuscript, built by a script that refuses to compile against a missing generated table

Workflow

convert Parse an upstream checkpoint and repack it into the self-describing native artifact.
build Configure and build the CUDA release preset; the CPU stub preset compiles but cannot produce evidence.
capture Dump upstream reference activations and gradients, taking inputs from the reference package’s own preprocessing path so the comparison isolates the model graph.
probe Run parity_probe against the converted artifact to produce the native side of each case.
compare Compare per tensor and per parameter name against gates fixed in advance, recording error distributions rather than a verdict.
report Emit the machine-written gate matrix, then the manuscript tables.

Reproducibility and validation

Every gate names the error mode it excludes — a shifted quantization bin edge, a misplaced end-of-sequence token, a wrong relative-attention bucket count — so a tolerance is evidence rather than decoration.
Gates were fixed before the comparisons ran and were not relaxed afterwards. Where a comparison failed, the implementation was corrected and the gate left alone.
Gating is on max-abs error and on error as a fraction of the reference tensor’s RMS. Elementwise relative error is recorded as a diagnostic and never gated, because it is unbounded wherever the reference approaches zero.
The Chronos gate is tighter than the Moirai-MoE gate — 1e-4 against 1e-3 — because Chronos encoder states sit an order of magnitude lower in absolute terms, so a gate copied across from the other family would have been roughly ten times looser in relative terms.
Every checkpoint appears in the generated matrix whether or not it was compared, so a case that was never executed is visible as not_run rather than quietly absent.
Comparison is per tensor and by name, with explicit coverage accounting, so anything the harness silently skipped surfaces as missing coverage and never as a pass.
Configurations whose memory requirement exceeds the benchmark machine are marked out_of_profile ahead of execution and kept in the machine-readable matrix, so a configuration never attempted cannot be mistaken for one attempted and lost.
The manuscript build fails when any machine-written table is missing, so the paper cannot compile against stale or absent evidence.

Where it stands

Forward agreement with the reference implementation is established for the small variant of each family. Chronos-T5-small agrees to 1.71 × 10⁻⁷ max-abs over 98,816 elements against a 10⁻⁴ gate. On values of order 10⁻² that is float32 round-off and nothing else, which is what makes it worth reporting: it establishes the achievable noise floor empirically, and so licenses treating a discrepancy above that floor as a defect rather than as a consequence of changing language or library. Moirai-MoE-1.0-R-small agrees to 1.66 × 10⁻⁵ against a 10⁻³ gate across smooth, volatile and multivariate inputs.

That is the extent of what this page claims. The two base variants report not_run because their reference fixtures had not been captured, and the small results do not stand in for them — the variants share the graph but not the widths, and at least one defect found during this work was width-dependent. How the implementation holds up beyond the forward pass, and how fast it runs, are open questions here: the benchmark suite has not been run, and its protocol commits in advance to comparing against compiled and graph-captured baselines rather than eager execution alone, and to reporting the configurations where the native implementation loses.

Availability

Not publicly released. The tree is Apache-2.0 licensed in place and carries a citation file and a pre-release licensing checklist, but it is explicitly pre-release: the fidelity matrix is incomplete and the performance suite has not run, so it should not be cited as evidence of numerical equivalence yet.

Model weights are not redistributed. They are downloaded from the upstream distribution points and converted locally, and converted artifacts are derived works that inherit the upstream licence — which is why they are excluded from version control. All measurements come from a single consumer GPU on Windows; no multi-GPU, multi-node or non-NVIDIA execution was attempted, and the implementation targets CUDA specifically rather than a portable backend.

The sources were extracted verbatim from the private research platform where they were developed, without renaming namespaces, macros or include paths, so the tree stays identical to the code that produced any earlier validation.

Related research

This software has no report of its own. These studies used the model families it implements.