A Hello-World Example of AI Transformer in C/C++ and CUDA

      

// transformer.cpp -- A hello-world example of transformer
// http://forejune.co/cuda/

#include <iostream>
#include <vector>
#include <cmath>
#include <random>
#include <algorithm>

using namespace std;

// Activation function: using ReLU (Rectified Linear Unit)
double f(double z)
{
  return max(0.0, z);
}






// softmax activation function, 2D input vector
vector<vector<double>> 
    softmax(const vector<vector<double>>& input) 
{
  int n = input.size(), m = input[0].size();// n x m matrix 

  // 2D output vector n x m y[n][m]
  vector< vector< double>> y(n, vector< double>(m));		 
        
  for (int i = 0; i < n; i++) {
    //maximum value of the row
    double maxVal=*max_element(input[i].begin(),input[i].end());
									     
    double sum = 0;
    // Subtract max for numerical stability
    for (int j = 0; j < m; j++) {
      y[i][j] = exp(input[i][j] - maxVal);
      sum += y[i][j];
    }
            
    // Normalize
    for (int j = 0; j < input[i].size(); j++) 
      y[i][j] /= sum;
  }

  return y;
}




 


// Matrix multiplication   C = A X B   A: nxm, B: mxr, C: nxr
vector<vector<double>> matmul(const vector<vector<double>>& A,
                              const vector<vector<double>>& B) 
{
  int n  = A.size();
  int m = A[0].size();
  int r = B[0].size();
        
  vector<vector<double>> C(n, vector<double>(r, 0));
        
  for (int i = 0; i < n; i++) 
    for (int j = 0; j < r; j++) 
      for (int k = 0; k < m; k++) 
        C[i][j] += A[i][k] * B[k][j];

   return C;
}
    
// Transpose matrix n x m, B = A^T : m x n
vector<vector<double>> transpose(const vector<vector<double>>& A) 
{
   int n = A.size();	//number of rows
   int m = A[0].size(); //number of columns
        
   vector<vector<double>> B(m, vector<double>(n));   // m x n
        
   for (int i = 0; i < n; i++) 
     for (int j = 0; j < m; j++)
       B[j][i] = A[i][j];

   return B;
}

// Generating Random weights
void randomWeights(vector<vector<double>> &w ) 
{
  int n = w.size();
  int m = w[0].size();
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < m; j++) {
      w[i][j] = (rand() % 1000) /1000.0;
      if ( rand() % 3 == 0 )
	 w[i][j] = -w[i][j];	
     }
   }
}



// Hello-word transformer class
class HelloWorldTransformer 
{
private:
    int d;  	 // Dimension of the model
    int nTokens; // Sequence length
    
    //Scaled dot-product attention, Q:Query, K:Key, V:Value
    vector<vector<double>> scaled_attention(
        const vector<vector<double>>& Q,
        const vector<vector<double>>& K,
        const vector<vector<double>>& V) 
    {
        
      // Q X K^T
      vector<vector<double>>QxKT = matmul(Q, transpose(K));
        
      // Scale by sqrt(d_k)
      double d = K[0].size();
      double scale = sqrt(d);
      for (int i = 0; i < QxKT.size(); i++)
	for(int j = 0; j < QxKT[i].size(); j++)
	    QxKT[i][j] /= scale;

      // Apply softmax
      vector<vector<double>> attention_weights = softmax(QxKT);
        
      // Multiply by V --> A x V
      return matmul(attention_weights, V);
    }





public:
    HelloWorldTransformer(int sequence_len, int model_dim) 
    { 
	nTokens = sequence_len;
        d = model_dim;
    }

    // Simple self-attention layer
    vector<vector<double>> self_attention(const vector<vector<double>>& Q0, 
	const vector<vector<double>>& K0, const vector<vector<double>>& V0) 
    {
        auto Q = Q0;
        auto K = K0;
        auto V = V0;
        
        //return scaled_dot_product_attention(Q, K, V);
        return scaled_attention(Q, K, V);
    }
};

void printMatrix(const vector<vector<double>> &A)
{
    for (const auto& row : A) {
        for (double val : row) {
            printf("%6.2f ", val);
        }
        cout << endl;
    }
    cout << endl;
}

int main() 
{
    srand(time (0));
    // Configuration
    int nTokens = 3;    // Number of tokens  (3 tokens)
    int d = 4;         // Embedding dimension
    
    // Create transformer
    HelloWorldTransformer transformer(nTokens, d);

     // 3 tokens with 4 features (dimensions)
     vector<vector<double>> X = {
        {1, 0, 1, 0},
        {0, 1, 0, 1},
        {1, 1, 1, 1}
    };

   // weight matrices
   vector<vector<double>> WQ(d, vector<double>(d));    
   vector<vector<double>> WK(d, vector<double>(d));    
   vector<vector<double>> WV(d, vector<double>(d));    
   randomWeights( WQ );
   randomWeights( WK );
   randomWeights( WV );

    // Compute Q, K, V
    auto Q = matmul(X, WQ);	//Q = X x WQ: nTokens x d
    auto K = matmul(X, WK);	//Q = X x WK: nTokens x d
    auto V = matmul(X, WV);	//Q = X x WV: nTokens x d
    
    cout << "Q, K, V embeddings:\n";
    printMatrix( Q );
    printMatrix( K );
    printMatrix( V );
    
    // Apply self-attention: Q, K, V
    auto output = transformer.self_attention(Q, K, V);
    
    cout << "Output after self-attention:\n";
    printMatrix( output );
   
    // Simple feed-forward network: FFN(x) = ReLU(W1*x)
    // W1[d][d]: weights for generating linear sums from attention output
    vector<vector<double>> W1 = {
        {-1,-1,0,1},
        {0,1,1,-1},
        {1,0,1,0},
        {1,1,1,0}
    };

    vector<vector<double>> z(nTokens, vector<double>(d));
    auto FFout = z;	// FFout has same dimensions as z 
    z = matmul(output, W1);
    for (int i = 0; i < nTokens; i++)
      for(int j = 0; j < d; j++) 
	FFout[i][j] = f(z[i][j]); //output of feed forward network

    cout << "Transformer Block Output:\n";
    printMatrix( FFout );

    return 0;
}