Engineering a CUDA Machine Learning Framework from Scratch: Microarchitectural Kernel Fusion, Vectorized Memory Pipelines, and Asymptotic Scaling

2. Motivation: Why Build an ML Framework from Scratch?

Modern deep learning engineering relies heavily on high-level frameworks like PyTorch, JAX, and TensorFlow. While these frameworks offer flexibility and comprehensive model coverage, their runtime architectures are engineered for generality rather than domain-specific hardware saturation. Writing neural network primitives directly in raw CUDA C++ uncovers three structural overheads inherent to general-purpose frameworks:

2.1 The Memory Bandwidth Bottleneck (FLOP/Byte Ratio)

Under the GPU Roofline Model, operations are partitioned into compute-bound (e.g., large GEMMs where arithmetic intensity $\gg 10\text{ FLOP/Byte}$) and memory-bandwidth bound (e.g., activations, normalization, Softmax, elementwise additions where arithmetic intensity $< 1\text{ FLOP/Byte}$). When executing standard PyTorch code:

Each isolated PyTorch operation writes its temporary output tensor back to global High-Bandwidth Memory (HBM/DRAM) over the $320\text{ GB/s}$ memory bus. In contrast, streaming multiprocessor (SM) 32-bit registers provide $>30\times$ higher aggregate bandwidth. By fusing activations, normalizations, and loss gradients into registers, intermediate DRAM traffic is eliminated entirely.

2.2 Host-Side CPU Launch Latency & Dynamic Autograd Graph Tape

PyTorch constructs a dynamic computational Directed Acyclic Graph (DAG) on the host CPU during every forward pass and traverses it in reverse during backward propagation. Each node allocation and CUDA driver dispatch incurs $3\text{--}8\,\mu\text{s}$ of host latency. For small-to-medium network architectures or unrolled recurrent sequence steps ($T \ge 64$), host-side graph recording accounts for $40\text{--}60\%$ of total wall-clock step time. Our custom CUDA C++ autograd bindings execute pre-allocated analytical backward pipelines with zero dynamic CPU graph allocations.

2.3 128-Bit Memory Vectorization and Register Warp Reductions

Standard un-vectorized kernels execute 32-bit single-precision memory loads (LDG.E.32). Aligning global memory pointers to 16 bytes and issuing 128-bit vector instructions (float4 / LDG.E.128) allows a single thread instruction to fetch 4 values simultaneously, saturating memory controller queues and hiding load latency. Intra-warp communication (such as row-wise maximums or sums) is executed via register shuffle intrinsics (__shfl_down_sync), completing reductions in 5 clock cycles without shared memory allocation or atomic lock contention.

3. Architecture Overview: 8 Core GPU Primitives

The computational engine consists of 8 standalone, hardware-optimized CUDA C++ primitives in kernels/. Each primitive is engineered to saturate GPU execution pipelines, serving as building blocks for 5 end-to-end trained deep learning architectures.

3.1 Mathematical Recurrence: Online FlashSoftmax

Standard Softmax requires three sequential passes over global memory: (1) find $\max(x)$, (2) compute normalizer $\sum e^{x - \max}$, and (3) normalize and store probabilities $P_i$. Our online FlashSoftmax kernel computes the running maximum and exponential sum concurrently in a single pass without writing intermediate exponentials to DRAM:

3.2 Mathematical Recurrence: Welford Single-Pass Normalization

Rather than performing separate passes over global memory for mean and variance, the Welford LayerNorm kernel updates moments concurrently inside warp register space:

4. Model Architecture Deep Dives & Empirical Parity

// Fused forward hypothesis & sigmoid activation kernel
__global__ void forward_kernel(const float* __restrict__ X,
                               const float* __restrict__ w,
                               const float* __restrict__ b,
                               float* __restrict__ y_hat,
                               int N, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float z = *b;
        #pragma unroll 4
        for (int j = 0; j < D; ++j) {
            z += X[idx * D + j] * w[j];
        }
        y_hat[idx] = 1.0f / (1.0f + __expf(-z));
    }
}

// Vectorized In-Place Adam Optimizer Kernel
__global__ void adam_kernel(float* __restrict__ param,
                            float* __restrict__ m,
                            float* __restrict__ v,
                            const float* __restrict__ grad,
                            float lr, float beta1, float beta2, float eps,
                            float bias_correction1, float bias_correction2,
                            float weight_decay, int size) {
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    if (idx < size) {
        float p = param[idx];
        float g = grad[idx];
        if (weight_decay != 0.0f) g += weight_decay * p;

        float m_val = beta1 * m[idx] + (1.0f - beta1) * g;
        float v_val = beta2 * v[idx] + (1.0f - beta2) * g * g;

        m[idx] = m_val;
        v[idx] = v_val;

        float m_hat = m_val / bias_correction1;
        float v_hat = v_val / bias_correction2;
        param[idx] = p - (lr * m_hat) / (sqrtf(v_hat) + eps);
    }
}

// Fused 4-Gate LSTM Forward Kernel in Register Space
__global__ void fused_lstm_gates_forward_kernel_torch(
    const float* __restrict__ d_gates_preact, // [N x 4H]
    const float* __restrict__ d_c_prev,       // [N x H]
    float* __restrict__ d_gates_act,          // [N x 4H] -> [i, f, g, o]
    float* __restrict__ d_c_next,             // [N x H]
    float* __restrict__ d_tanh_c,             // [N x H]
    float* __restrict__ d_h_next,             // [N x H]
    int N, int H) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N * H) {
        int n = idx / H;
        int h = idx % H;
        int base_4h = n * (4 * H);

        float z_i = d_gates_preact[base_4h + h];
        float z_f = d_gates_preact[base_4h + H + h];
        float z_g = d_gates_preact[base_4h + 2 * H + h];
        float z_o = d_gates_preact[base_4h + 3 * H + h];

        float i_val = 1.0f / (1.0f + __expf(-z_i));
        float f_val = 1.0f / (1.0f + __expf(-z_f));
        float g_val = tanhf(z_g);
        float o_val = 1.0f / (1.0f + __expf(-z_o));

        if (d_gates_act) {
            d_gates_act[base_4h + h] = i_val;
            d_gates_act[base_4h + H + h] = f_val;
            d_gates_act[base_4h + 2 * H + h] = g_val;
            d_gates_act[base_4h + 3 * H + h] = o_val;
        }

        float c_prev = d_c_prev ? d_c_prev[idx] : 0.0f;
        float c_next = f_val * c_prev + i_val * g_val;
        float tc = tanhf(c_next);
        float h_next = o_val * tc;

        d_c_next[idx] = c_next;
        if (d_tanh_c) d_tanh_c[idx] = tc;
        d_h_next[idx] = h_next;
    }
}

// Fused GRU Backward Kernel with In-Register dh Accumulation
__global__ void gru_step_backward_kernel(
    const float* __restrict__ dh_incoming,  // [N, H] -> dH_seq[t]
    const float* __restrict__ dh_recurrent, // [N, H] -> dh_next from step t+1
    const float* __restrict__ h_prev,       // [N, H]
    const float* __restrict__ gates_act,    // [N, 3H] -> [r, z, n]
    const float* __restrict__ g_hh,         // [N, 3H] -> raw recurrent projections
    const float* __restrict__ b_hh,         // [3H] (optional)
    float* __restrict__ dg_ih,              // [N, 3H] -> [dz_r, dz_z, dz_n]
    float* __restrict__ dg_hh,              // [N, 3H] -> [dz_r, dz_z, dG_hh_n]
    float* __restrict__ dh_prev_direct,     // [N, H] -> direct gradient (dh * z)
    int N, int H) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N * H) {
        int n = idx / H;
        int h = idx % H;
        int base_3h = n * (3 * H);

        // Fused in-register accumulation: dh = dh_incoming + dh_recurrent
        float dh = (dh_incoming ? dh_incoming[idx] : 0.0f) + 
                   (dh_recurrent ? dh_recurrent[idx] : 0.0f);
        float hp = (h_prev ? h_prev[idx] : 0.0f);

        float r_val = gates_act[base_3h + h];
        float z_val = gates_act[base_3h + H + h];
        float n_val = gates_act[base_3h + 2 * H + h];

        float g_hh_n_val = g_hh[base_3h + 2 * H + h] + (b_hh ? b_hh[2 * H + h] : 0.0f);

        float dn = dh * (1.0f - z_val);
        float dz = dh * (hp - n_val);

        float dz_z = dz * z_val * (1.0f - z_val);
        float dz_n = dn * (1.0f - n_val * n_val);
        float dr = dz_n * g_hh_n_val;
        float dz_r = dr * r_val * (1.0f - r_val);
        float dg_hh_n = dz_n * r_val;

        // Write out coalesced pre-activation gradients
        dg_ih[base_3h + h] = dz_r;
        dg_ih[base_3h + H + h] = dz_z;
        dg_ih[base_3h + 2 * H + h] = dz_n;

        dg_hh[base_3h + h] = dz_r;
        dg_hh[base_3h + H + h] = dz_z;
        dg_hh[base_3h + 2 * H + h] = dg_hh_n;

        if (dh_prev_direct) dh_prev_direct[idx] = dh * z_val;
    }
}

5. Master GPU Benchmark Suite (NVIDIA Tesla T4)

6. Limitations & What I'd Do Differently

A rigorous systems evaluation must document where the custom implementation encounters hardware boundaries and architectural limitations:

6.1 FP32 CUDA Cores vs. Mixed-Precision Tensor Cores

Our custom GEMM engines are written for single-precision (FP32) CUDA cores using double-buffered shared-memory tiling, reaching $85\text{--}92\%$ of cuBLAS FP32 throughput. However, modern production LLM engines leverage NVIDIA Tensor Cores (via WMMA or MMA PTX assembly) in FP16 / BF16 mixed precision, offering an order of magnitude higher theoretical FLOPS ($65\text{ TFLOPS}$ on T4 for FP16 vs $8.1\text{ TFLOPS}$ for FP32). Adding half-precision Tensor Core micro-kernels is the immediate next step for scaling to billion-parameter workloads.

6.2 Persistent RNN Megakernels at Small Hidden Sizes ($H \le 128$)

At small hidden dimensions ($H=128$), PyTorch's cuDNN implementation is $2\times$ faster than our custom kernel. cuDNN achieves this by executing a proprietary Persistent RNN Megakernel that keeps recurrent weights $\mathbf{W}_{hh}$ resident in on-chip SRAM across all $T$ timesteps, bypassing DRAM reload. Our current engine issues $T$ consecutive recurrent GEMMs. Implementing a persistent SRAM CTA loop across temporal steps will resolve this discrepancy for small models.

6.3 Asynchronous Stream Overlap & CUDA Graph Capture

While our C++ autograd bindings eliminate CPU Python dynamic tape construction, each backward operation is still launched into a single default stream. Leveraging `cudaGraph_t` capture to encapsulate the static recurrent execution graph into hardware dispatch queues would eliminate the remaining $1\text{--}2\,\mu\text{s}$ of driver overhead per step.

6.4 Multi-Head Attention & FlashAttention-2 Extension

The online FlashSoftmax algorithm implemented in kernels/src/softmax.cu can be extended directly into a 2D block-tiled causal Multi-Head Attention (FlashAttention-2) kernel, allowing this framework to scale from recurrent sequence models to full Transformer decoder stacks.

7. Code Repository & Reproducibility

The complete source code, CUDA kernels, Python extensions, and benchmarking suites are open-source and structured for reproduction:

# 1. Clone the repository
git clone https://github.com/sarimahsan/cuda-ml-from-scratch.git
cd cuda-ml-from-scratch

# 2. Install dependencies
pip install torch ninja torchvision pandas scikit-learn

# 3. Benchmark standalone CUDA Kernel Engine primitives vs PyTorch:
python kernels/benchmarks/benchmark_all_kernels.py

# 4. Run MLP training on MNIST (>98% test accuracy):
python 02_mlp/train_mnist.py

# 5. Run LSTM Shakespeare Language Model and Parity Suite:
python 04_lstm/benchmark.py
python 04_lstm/train_sequence.py

# 6. Run GRU 150-iteration Statistical Benchmark Suite:
python 05_gru/benchmark.py --warmup 50 --reps 150

!git clone https://github.com/sarimahsan/cuda-ml-from-scratch.git && cd cuda-ml-from-scratch && pip install ninja && python 04_lstm/benchmark.py

Repository: github.com/sarimahsan/cuda-ml-from-scratch
License: MIT Open Source License. Written for systems research, high-performance computing education, and GPU kernel engineering.

Read More Systems Research & Engineering