Mohammad Mohammad

All work

TitanInfer

A neural-network inference library in C++, written without BLAS, Eigen or LAPACK. It loads a model from a binary file and runs it on the CPU.

matmul: C += A × B, 16 × 16 floatsblocked order
step
0 / 512
lines loaded
0
cache hits
0
whole run, lines loaded
unblocked –, blocked –
A model, not a measurement: 16 × 16 matrices split into 4 × 8 and 8 × 4 blocks (the real kernel uses 64 × 256 × 64) and a 24-line cache that evicts the least recently used line. Same arithmetic in both orders; only the order of the loops changes.
Role
Solo
Dates
Feb 8 to Mar 28, 2026 (29 commits)
Language
C++20 with CMake (the cluster controller uses std::jthread). Python for the PyTorch exporter.
Tests
346 GoogleTest cases (294 TEST, 52 TEST_F) in 20 test executables. GitHub Actions runs on Ubuntu 22.04 and Windows.
Dependencies
The core library links only the system threads library. GoogleTest and Google Benchmark are fetched for tests. gRPC is optional and off by default.
Code
github.com/HamoodiM/TitanInfer

Worth a look

What I built

  • An N-dimensional tensor that owns 32-byte-aligned memory, with copy and move semantics.
  • A scalar matrix multiply and a cache-blocked AVX2 version that uses FMA when the compiler enables it. The AVX2 kernel is tested and benchmarked on its own, but the layers don't call it yet.
  • Dense, ReLU, Sigmoid, Tanh, Softmax, Conv2D, max and average pooling, and flatten layers in a Sequential container.
  • The .titan binary model format, a writer and a parser, and a Python script that exports PyTorch models to it.
  • An inference engine that allocates one output buffer per layer at load time, with warmup runs and per-layer timing.
  • A ModelHandle guarded by a mutex, an exception hierarchy, a logger, and a C API (titaninfer_load, titaninfer_predict, titaninfer_free and friends) with a C example.
  • INT8 quantization (q = clamp(round(x / scale) + zero_point, -128, 127)) with an SSE kernel that widens to 16 bits (_mm_cvtepi8_epi16) and multiply-adds with _mm_madd_epi16, and fusion of Dense followed by ReLU or Sigmoid into one layer.
  • A thread pool and a request batcher that groups concurrent inference calls.
  • A multi-model server with versioned routes (/v1/models/{name}/versions/{ver}/predict), LRU model eviction and per-tenant quotas.
  • A cluster controller that routes with power-of-two random choices over a load score (queue depth, CPU, p95 latency against the SLO), a per-node circuit breaker, and autoscaling signals, with optional gRPC binaries, Docker files and a chaos test.

An aligned tensor

AVX2 loads work on 32-byte registers, so the tensor allocates its storage on a 32-byte boundary. The byte size is rounded up to a multiple of the alignment because std::aligned_alloc requires it, and each platform gets its own allocator:

src/tensor.cpp, lines 141–158
#ifdef _WIN32
    // Windows: _aligned_malloc
    ptr = static_cast<float*>(_aligned_malloc(aligned_size, ALIGNMENT));
    if (!ptr) {
        throw std::bad_alloc();
    }
#elif defined(__APPLE__)
    // macOS: posix_memalign
    if (posix_memalign(reinterpret_cast<void**>(&ptr), ALIGNMENT, aligned_size) != 0) {
        throw std::bad_alloc();
    }
#else
    // Linux: aligned_alloc (C++17)
    ptr = static_cast<float*>(std::aligned_alloc(ALIGNMENT, aligned_size));
    if (!ptr) {
        throw std::bad_alloc();
    }
#endif

The tensor is RAII: the destructor frees through the matching platform call, and the move constructor takes the pointer and leaves the source empty, so returning tensors by value doesn't copy the data. In practice the matmul kernel below still uses unaligned loads (_mm256_loadu_ps), because a row in the middle of a matrix isn't guaranteed to start on a 32-byte boundary.

Blocked AVX2/FMA matrix multiply

The SIMD matmul tiles the loops into 64 × 256 × 64 blocks (MC, KC, NC) so each block's working set stays in cache, then runs an 8-wide inner loop along K. A horizontal add at the end collapses the 8 lanes into one value, and a scalar loop handles the remainder. The model at the top of this page runs the same loop order on 16 × 16 matrices with a 24-line cache, so you can watch which lines get reused.

src/ops/matrix_ops_simd.cpp, lines 124–144
                        for (; kk + 8 <= k + kb; kk += 8) {
                            // Load 8 elements from A[ii, kk:kk+8]
                            __m256 a_vec = _mm256_loadu_ps(A_data + ii * K + kk);

                            // Load 8 elements from B[kk:kk+8, jj]
                            // Note: B is row-major, so we need to gather elements
                            __m256 b_vec =
                                _mm256_set_ps(B_data[(kk + 7) * N + jj], B_data[(kk + 6) * N + jj],
                                              B_data[(kk + 5) * N + jj], B_data[(kk + 4) * N + jj],
                                              B_data[(kk + 3) * N + jj], B_data[(kk + 2) * N + jj],
                                              B_data[(kk + 1) * N + jj], B_data[(kk + 0) * N + jj]);

#ifdef __FMA__
                            // Use fused multiply-add
                            sum_vec = _mm256_fmadd_ps(a_vec, b_vec, sum_vec);
#else
                            // Manual multiply-add
                            __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                            sum_vec     = _mm256_add_ps(sum_vec, prod);
#endif
                        }

The kernel isn't on the inference path yet: the Dense, fused and Conv2D layers still call the scalar ops::matmul and ops::matvec. Routing them through the AVX2 kernel is the first next step.

The kernel's own weak point is the load of B. B is row-major, so 8 consecutive K values for one column are N floats apart, and the kernel gathers them with _mm256_set_ps from 8 scalar loads. The benchmark below points at that gather as the bottleneck, though I haven't profiled it yet. The next step is to pack each block of B into a contiguous, transposed buffer once, then compute several output columns per pass so each A load is reused.

Measured with matrix_benchmark (Release, MSVC /arch:AVX2, Google Benchmark, median of 3) on an Intel Core i5-14400F, Nov 2026:

SizeNaiveAVX2 kernelSpeed-up
64 × 6470.6 µs62.5 µs1.13×
128 × 1280.90 ms0.73 ms1.23×
256 × 2568.92 ms8.40 ms1.06×
512 × 51269.6 ms69.1 ms1.01×

So the kernel barely beats the naive loop, and the gain shrinks as the matrices grow. That fits the gather: past 128 × 128 the strided loads of B dominate, and blocking can't help while every 8-wide FMA waits on 8 scalar loads. Packing B is the fix I'd make first, then re-run this table.

The .titan model format

A .titan file starts with the magic bytes TITN and a format version, followed by a layer count and, for each layer, a type id, its shape parameters and its float weights. The parser refuses files from a newer format version than it knows instead of misreading them:

include/titaninfer/io/format.hpp, lines 8–12
/// Magic number for .titan files: "TITN"
static constexpr char TITAN_MAGIC[4] = {'T', 'I', 'T', 'N'};

/// Current format version
static constexpr uint32_t TITAN_FORMAT_VERSION = 2;
src/io/model_parser.cpp, lines 50–56
// Read and validate version
uint32_t version = read_value<uint32_t>(in);
if (version > TITAN_FORMAT_VERSION) {
    throw std::runtime_error("ModelParser: unsupported format version " +
                             std::to_string(version) +
                             " (max supported: " + std::to_string(TITAN_FORMAT_VERSION) + ")");
}

When the engine loads a model it walks the layers once, computes each output shape, and allocates an output buffer per layer. A forward pass then reads from buffer i - 1 and writes into buffer i, so layer outputs reuse the same memory on every call:

src/engine/inference_engine.cpp, lines 216–227
// Layers 1..N-1: buffers_[i-1] -> buffers_[i]
for (size_t i = 1; i < model_->size(); ++i) {
    if (profiling_enabled_) {
        auto start = clock::now();
        model_->layer(i).forward(buffers_[i - 1], buffers_[i]);
        auto end = clock::now();
        stats_.layer_times_ms[i] +=
            std::chrono::duration<double, std::milli>(end - start).count();
    } else {
        model_->layer(i).forward(buffers_[i - 1], buffers_[i]);
    }
}

That isn't the same as allocation-free inference. The Dense layer's batched (2-D) path allocates a transposed copy of its weights on every call, and predict() returns a deep copy of the last buffer.

Known issues

  • The layers use the scalar matmul, not the AVX2 kernel.
  • The Dense layer's 2-D path allocates and transposes its weights on every call. Storing the weights pre-transposed at load time would remove that.
  • predict() copies its output. An overload that writes into a caller-owned tensor would remove that allocation.
  • The PyTorch exporter writes format version 1 while the C++ writer uses version 2. The parser accepts both, but the exporter only handles nn.Linear and activations, not Conv2D.
  • ModelHandle uses one mutex for predict, so concurrent calls on the same handle run one at a time.
  • No benchmark results are committed yet, and the README's test counts are out of date.

How I built this

I designed and built TitanInfer over February and March 2026. The final cleanup commits were made with an AI assistant and are marked as co-authored. The cache model on this page was written for this site.

Back to all projects