MLP-CUDA

High-performance CUDA implementations of neural network components, featuring:

• Batched MLP with Adam optimizer
• Multi-Head Attention for Transformer architectures

Overview

This project provides GPU-accelerated implementations of fundamental deep learning building blocks:

1. Multi-Layer Perceptron (MLP)

A 2-hidden-layer MLP (4 layers total) with:

• Batched forward pass for efficient inference
• Batched backward pass with automatic differentiation
• Adam optimizer for training
• Efficient CUDA kernels with tiled matrix multiplication
• Comprehensive test suite for correctness verification

Architecture:

Input Layer (h1) → Hidden Layer 1 (h2) → Hidden Layer 2 (h3) → Output Layer (h4)
                    [ReLU]                  [ReLU]              [Linear]

2. Multi-Head Attention

Transformer-style attention mechanism with:

• Scaled dot-product attention
• Multiple attention heads for parallel attention
• Self-attention and cross-attention support
• Efficient batched operations
• Custom CUDA kernels for reshape/transpose operations

Architecture:

Input [B, N, d_model]
  ↓
QKV Projections → [B, N, 3×d_model]
  ↓
Reshape to heads → [B, h, N, d_k]
  ↓
Scaled Dot-Product Attention → softmax(Q·K^T/√d_k)·V
  ↓
Concatenate heads → [B, N, d_model]
  ↓
Output Projection → [B, N, d_model]

Features

MLP Features

• ✅ Batched Operations: Process multiple samples simultaneously
• ✅ Tiled Matrix Multiplication: Efficient CUDA implementation with shared memory
• ✅ ReLU Activation: Fast element-wise activation with backward pass
• ✅ Adam Optimizer: Adaptive learning rate with momentum
• ✅ MSE Loss: Mean Squared Error for regression tasks
• ✅ Numerical Stability: Careful handling of edge cases
• ✅ Save/Load: Persist trained models to disk

Attention Features

• ✅ Multi-Head Attention: Parallel attention with configurable heads
• ✅ Self-Attention: Query, Key, Value from same input
• ✅ Cross-Attention: Separate Query and Key/Value inputs
• ✅ Batched Softmax: Optimized softmax with shared memory reduction
• ✅ Reshape/Transpose Kernels: Efficient head separation and concatenation
• ✅ Variable Sequence Lengths: Support for different sequence lengths

Common Features

• ✅ Comprehensive Tests: Unit and integration tests for all components
• ✅ Examples: Working demonstrations for both MLP and Attention
• ✅ Documentation: Mathematical derivations and implementation guides

Data Handling

Document Boundary Awareness with <|endoftext|>

The training pipeline includes intelligent handling of document boundaries using <|endoftext|> markers:

How it works:

1. Tokenizer Recognition: The <|endoftext|> marker is automatically recognized and converted to an EOS (End-of-Sequence) token during encoding
2. Document Splitting: Training data is split into separate documents at each <|endoftext|> marker
3. Boundary Respect: Training sequences never span across document boundaries, preventing the model from learning incorrect transitions between unrelated content

Example:

Input file:
once upon a time there was a cat<|endoftext|>the weather today is sunny<|endoftext|>

What the model sees:
Document 1: "once upon a time there was a cat" [EOS]
Document 2: "the weather today is sunny" [EOS]

✓ Training sequences stay within document boundaries
✗ No sequences mixing unrelated content like "...cat the weather..."

Benefits:

• ✅ Clean document separation in multi-document datasets
• ✅ Model learns when text should end (via EOS token)
• ✅ No contamination between unrelated documents during training
• ✅ Better generation quality with proper ending behavior

Usage: Simply include <|endoftext|> markers in your training data:

echo "First document text here<|endoftext|>Second document here<|endoftext|>" > train.txt
./train_transformer --data train.txt

Project Structure

MLP-cuda/
├── include/                       # Header files
│   ├── mlp.h                     # MLP class
│   ├── multi_head_attention.h    # Multi-Head Attention class
│   ├── matrix_ops.h              # Matrix operations
│   ├── attention_ops.h           # Attention-specific operations
│   ├── activations.h             # Activation functions
│   ├── loss.h                    # Loss functions
│   └── adam.h                    # Adam optimizer
├── src/                          # Implementation files
│   ├── mlp.cu                    # MLP implementation
│   ├── multi_head_attention.cu   # Attention implementation
│   ├── matrix_ops.cu             # Matrix kernels
│   ├── attention_ops.cu          # Attention kernels
│   ├── activations.cu            # Activation kernels
│   ├── loss.cu                   # Loss kernels
│   └── adam.cu                   # Adam kernels
├── tests/                        # Test suite
│   ├── test_matrix_ops.cu        # Matrix operation tests
│   ├── test_mlp.cu               # MLP tests
│   └── test_attention.cu         # Attention tests
├── examples/                     # Example programs
│   ├── train_regression.cu       # MLP regression example
│   └── attention_demo.cu         # Attention demonstration
├── docs/                         # Documentation
│   ├── mathematical_derivation.md      # MLP mathematics
│   ├── cuda_implementation_plan.md     # MLP CUDA plan
│   └── attention_design.md             # Attention design
├── CMakeLists.txt               # Build configuration
└── README.md

Requirements

• CUDA Toolkit: Version 10.0 or higher
• CMake: Version 3.18 or higher
• C++ Compiler: Supporting C++14
• GPU: NVIDIA GPU with compute capability 6.0 or higher

Building

# Create build directory
mkdir build && cd build

# Configure with CMake
cmake ..

# Build
make

# Run tests
ctest --verbose

Usage

Basic Example

#include "mlp.h"

int main() {
    // Define network architecture: [input, hidden1, hidden2, output]
    int layer_sizes[4] = {784, 256, 128, 10};
    int batch_size = 128;
    float learning_rate = 0.001f;

    // Create MLP
    MLP mlp(layer_sizes, batch_size, learning_rate);

    // Prepare data (arrays on host)
    float* train_X;  // Shape: [batch_size, 784]
    float* train_Y;  // Shape: [batch_size, 10]

    // Training step
    float loss = mlp.train_step(train_X, train_Y, batch_size);

    // Inference
    float* test_X;   // Shape: [batch_size, 784]
    float* output;   // Shape: [batch_size, 10]
    mlp.forward(test_X, output, batch_size);

    // Save model
    mlp.save_parameters("model.bin");

    return 0;
}

Running Examples

# Regression example
./train_regression

Expected output:

======================================
  MLP CUDA - Regression Example
======================================

Network Architecture:
  Input:    2 neurons
  Hidden 1: 32 neurons
  Hidden 2: 16 neurons
  Output:   1 neurons

Training for 500 epochs...
Epoch | Train Loss | Test Loss
------|------------|----------
    1 | 0.523187 | 0.518234
   50 | 0.045321 | 0.043876
  100 | 0.023456 | 0.024123
  ...
  500 | 0.008234 | 0.009123

Final - Train Loss: 0.008234, Test Loss: 0.009123

API Reference

MLP Class

Constructor

MLP(int layer_sizes[4], int batch_size,
    float learning_rate = 0.001f,
    float beta1 = 0.9f,
    float beta2 = 0.999f,
    float epsilon = 1e-8f)

Parameters:

• layer_sizes: Array of 4 integers [h1, h2, h3, h4] defining layer sizes
• batch_size: Maximum batch size
• learning_rate: Learning rate for Adam optimizer (default: 0.001)
• beta1: Adam beta1 parameter (default: 0.9)
• beta2: Adam beta2 parameter (default: 0.999)
• epsilon: Adam epsilon for numerical stability (default: 1e-8)

Methods

forward()

void forward(const float* h_X, float* h_output, int batch_size)

Performs forward pass (inference).

• h_X: Input batch on host [batch_size × h1]
• h_output: Output buffer on host [batch_size × h4]
• batch_size: Actual batch size (≤ max batch size)

train_step()

float train_step(const float* h_X, const float* h_Y, int batch_size)

Performs one training iteration (forward + backward + update).

• h_X: Input batch on host [batch_size × h1]
• h_Y: Target batch on host [batch_size × h4]
• Returns: Loss value

evaluate()

float evaluate(const float* h_X, const float* h_Y, int batch_size)

Computes loss without updating parameters.

save_parameters() / load_parameters()

void save_parameters(const char* filename)
void load_parameters(const char* filename)

Save/load model parameters to/from disk.

MultiHeadAttention Class

Constructor

MultiHeadAttention(
    int d_model,
    int num_heads,
    int max_seq_len,
    int max_batch_size
)

Parameters:

• d_model: Model dimension (must be divisible by num_heads)
• num_heads: Number of attention heads (e.g., 8)
• max_seq_len: Maximum sequence length
• max_batch_size: Maximum batch size

Methods

forward() - Self-Attention

void forward(
    const float* h_X,
    float* h_output,
    int batch_size,
    int seq_len,
    const float* h_mask = nullptr
)

Performs self-attention where Q = K = V = X.

• h_X: Input on host [batch_size × seq_len × d_model]
• h_output: Output buffer on host [batch_size × seq_len × d_model]
• batch_size: Actual batch size
• seq_len: Actual sequence length
• h_mask: Optional attention mask (1 = attend, 0 = ignore)

forward_cross() - Cross-Attention

void forward_cross(
    const float* h_Q,
    const float* h_KV,
    float* h_output,
    int batch_size,
    int seq_len_q,
    int seq_len_kv,
    const float* h_mask = nullptr
)

Performs cross-attention with separate query and key/value inputs.

• h_Q: Query input [batch_size × seq_len_q × d_model]
• h_KV: Key/Value input [batch_size × seq_len_kv × d_model]
• h_output: Output buffer [batch_size × seq_len_q × d_model]

Example Usage:

#include "multi_head_attention.h"

int main() {
    int d_model = 512;
    int num_heads = 8;
    int seq_len = 64;
    int batch_size = 32;

    MultiHeadAttention mha(d_model, num_heads, 128, 64);

    float* input;   // [32, 64, 512]
    float* output;  // [32, 64, 512]

    // Self-attention
    mha.forward(input, output, batch_size, seq_len);

    return 0;
}

Performance

Typical performance on NVIDIA RTX 3090 (example configuration):

Batch Size	Network Size	Forward (ms)	Backward (ms)	Total (ms)
128	784-256-128-10	0.8	1.6	3.2
256	784-256-128-10	1.2	2.4	4.8
512	784-256-128-10	2.1	3.8	7.2

Throughput: ~35,000 samples/second for batch size 128

Mathematical Details

See docs/mathematical_derivation.md for:

• Detailed derivation of forward pass
• Backpropagation equations
• Adam optimizer formulation
• Batched computation details

See docs/cuda_implementation_plan.md for:

• CUDA kernel design
• Memory management strategy
• Optimization techniques
• Performance analysis

Testing

Unit Tests

Test individual kernels:

./test_matrix_ops    # Matrix operations

Integration Tests

Test full MLP:

./test_mlp          # Full MLP tests including XOR and overfitting

Test Coverage

• ✅ Matrix multiplication (standard, transposed A, transposed B)
• ✅ Bias operations (add, gradient sum)
• ✅ Activation functions (ReLU forward/backward)
• ✅ Loss functions (MSE forward/gradient)
• ✅ Adam optimizer updates
• ✅ Forward pass correctness
• ✅ Backward pass (gradient checking)
• ✅ Overfitting on small datasets
• ✅ XOR problem (non-linear learning)
• ✅ Save/load functionality
• ✅ Multiple batch sizes

Implementation Highlights

Tiled Matrix Multiplication

Uses shared memory tiling (16×16 tiles) to reduce global memory access:

__shared__ float As[TILE_SIZE][TILE_SIZE];
__shared__ float Bs[TILE_SIZE][TILE_SIZE];

// Load tiles collaboratively
// Compute partial products
// Accumulate results

Benefits:

• Reduces global memory access by ~16x
• Coalesced memory access patterns
• High arithmetic intensity

Activation Functions

ReLU implementation with separate forward/backward kernels:

// Forward: y = max(0, x)
output[i] = fmaxf(0.0f, input[i]);

// Backward: dy/dx = 1 if x > 0, else 0
grad_input[i] = (input[i] > 0.0f) ? grad_output[i] : 0.0f;

Adam Optimizer

Efficient single-kernel update:

m = β₁·m + (1-β₁)·g
v = β₂·v + (1-β₂)·g²
m̂ = m / (1 - β₁ᵗ)
v̂ = v / (1 - β₂ᵗ)
θ = θ - α·m̂ / (√v̂ + ε)

Limitations

• Fixed architecture: 2 hidden layers (4 layers total)
• MSE loss only (can be extended to cross-entropy)
• Single GPU only
• Maximum batch size must be specified at construction

Future Enhancements

• Configurable number of layers
• Additional activation functions (sigmoid, tanh, GELU)
• Cross-entropy loss for classification
• Batch normalization
• Dropout regularization
• Multi-GPU support
• Mixed precision training (FP16)
• cuBLAS integration for matrix operations

License

See LICENSE file for details.

References

1. Kingma & Ba (2014). "Adam: A Method for Stochastic Optimization"
2. LeCun et al. (1998). "Gradient-Based Learning Applied to Document Recognition"
3. NVIDIA CUDA Programming Guide
4. Kirk & Hwu. "Programming Massively Parallel Processors"

Contributing

Contributions are welcome! Please ensure:

• All tests pass (ctest --verbose)
• Code follows existing style
• New features include tests
• Documentation is updated

Acknowledgments

This implementation demonstrates efficient GPU computing for deep learning, showcasing:

• CUDA kernel optimization
• Memory management strategies
• Numerical stability considerations
• Comprehensive testing methodology