Kestrel

A high-performance C++20 distributed runtime that coordinates sharded dataset prefetching and non-blocking checkpointing for large-scale ML training jobs.

I built this to experiment with lock-free concurrency, custom memory allocation strategies, and low-overhead network coordination to see how to prevent GPU starvation during heavy disk I/O.

Architecture

┌──────────────────────────────────────────────────────┐
│                    Public API                        │
│                 kestrel::Runtime                     │
├──────────────────────────────────────────────────────┐
│               Coordination Layer                     │
│     Coordinator · WorkerGroup · BarrierSync          │
├───────────────────┬──────────────────────────────────┤
│  Checkpoint Flow  │   Data Flow                      │
│  Serializer       │   ShardedDataset                 │
│  CheckpointWriter │   ShardRouter                    │
│  CheckpointStore  │   Prefetcher                     │
├───────────────────┴──────────────────────────────────┤
│               Network Transport                      │
│  Protocol · TcpTransport · ConnectionPool            │
├──────────────────────────────────────────────────────┤
│                Core Runtime                          │
│    ThreadPool · MPMCQueue · MemoryPool · AsyncIO     │
└──────────────────────────────────────────────────────┘

• Checkpoint Pipeline: Worker State → Binary Serializer → Lock-free Queue → Async Scatter-gather I/O → Atomic Commit.
• Data Prefetch Pipeline: Dataset Registry → Jump Consistent Hash Sharding → Async Read → Slab Allocator → Bounded Queue (Backpressure).

Building & Testing

# Configure & Build
cmake -B build -DCMAKE_BUILD_TYPE=Release -DKESTREL_BUILD_BENCHMARKS=ON
cmake --build build --config Release

# Run Unit Tests
cd build
ctest --output-on-failure -C Release

How It Works (Quick Example)

Coordinator Node

#include "kestrel/coord/coordinator.h"

int main() {
    kestrel::CoordinatorConfig config{
        .port = 9000,
        .expected_workers = 8,
        .num_shards = 256,
        .heartbeat_interval = std::chrono::milliseconds(500)
    };
    kestrel::Coordinator coordinator(config);
    coordinator.start();
    
    // Block or handle coordinate events
    while (coordinator.running()) {
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
}

Worker Node

#include "kestrel/api/runtime.h"

int main() {
    auto runtime = kestrel::Runtime::Builder()
        .num_workers(8)
        .checkpoint_dir("checkpoints")
        .checkpoint_interval_steps(100)
        .prefetch_depth(4)
        .build();

    // Register parameters to checkpoint
    std::vector<float> weights(1000000, 0.0f);
    runtime->register_checkpoint("model.weights", weights.data(), weights.size() * sizeof(float));

    runtime->start();

    // Simulated training loop
    for (uint64_t step = 1; step <= 1000; ++step) {
        // ... compute step ...
        
        if (runtime->should_checkpoint(step)) {
            runtime->trigger_checkpoint(); // Offloaded asynchronously
        }
    }
    
    runtime->shutdown();
}

Performance & Benchmarks

You can run ./build/benchmarks/ binaries to reproduce these results.

1. Slab Memory Allocator vs malloc (64KB blocks)

Standard heap allocations bypass thread-local heaps for larger buffers, triggering kernel page maps. The pre-allocated slab memory pool reuses blocks inside user-space to avoid this.

• 1 Thread: 2.03x speedup over malloc (34.27 Mops/s vs 16.85 Mops/s)
• 4 Threads: 7.40x speedup over malloc (21.93 Mops/s vs 2.96 Mops/s under lock contention)

2. Lock-free MPMC Queue vs std::mutex + std::deque

Uses Dmitry Vyukov's bounded ring buffer with atomic sequence indices and alignas(64) padding to prevent cache-line false sharing.

• Speedup: Stable ~1.1x to 1.2x throughput improvement under heavy contention across 2–16 threads.

3. Checkpoint Pipeline Interruption

Writing a 16MB state vector directly to disk blocks the execution thread. By offloading serialization to a thread pool and writes to a scatter-gather AsyncIOEngine, wait times are bypassed.

• Synchronous Write: 51.6 ms training loop freeze.
• Asynchronous Pipeline: 0.0 ms training loop freeze (offloaded within 20 microseconds).