Sign up

Sign in

Sign up

Sign in

7 Step Optimization of Parallel Reduction with CUDA

Rimika Dhara

Follow

--

Listen

Share

In this post, I aim to take a simple yet popular algorithm — Parallel Reduction — and optimize its performance as much as possible. This effort is inspired by an NVIDIA webinar hosted by Mark Harris, from which I’m not only recreating the optimizations but also attempting to simplify the concepts for better understanding. Alongside this article, I have provided a GitHub implementation of these methods. While here I only showcase the final kernels with their methods, my GitHub repository includes more detailed information that will assist you in recreating this work.

Let’s get started!

Understanding the Algorithm

Let’s start by exploring what the parallel reduction algorithm entails. It’s a data-parallel primitive that’s straightforward to implement in CUDA. To keep it simple, parallel reduction aims to reduce a vector, matrix, or tensor in parallel by leveraging a GPU’s thread hierarchy. This reduction is achieved through operations like sum(), min(), max(), or avg() to aggregate and reduce the data. We will be using sum() to reduce our dataset. Despite their simplicity, these operations are versatile and crucial for many applications, requiring high optimizations to avoid becoming bottlenecks. Although these computations may seem simple, they can be time-consuming if not efficiently handled.

When parallelizing, the algorithm can be thought of as a tree-based approach, spreading across each thread block of our GPU. A critical question arises: “How can we communicate partial results between thread blocks?” The most straightforward solution might seem to be Global Synchronization — allow each block to compute, then synchronize them all and continue recursively. However, CUDA does not support global synchronization because it is costly in terms of hardware (HW) and would constrain the programmer to using only a few blocks to avoid deadlocks, thus reducing overall efficiency.

A more practical approach to communicating partial results while computing in each thread block is to decompose our kernel into multiple kernels. Kernel decomposition involves breaking down a large kernel task into smaller, manageable sub-tasks, which can be executed independently across different threads or blocks. This method minimizes hardware and software overhead. This allows for more flexible and efficient use of GPU resources, reducing the need for synchronization and improving overall computational speed.

Our Metrics

Our algorithm’s performance hinges on two critical metrics: time and bandwidth. These metrics gauge whether our GPU is being fully utilized, essentially measuring if it’s achieving peak performance. We aim for GPU peak performance with our metrics reflected in terms of compute (GFLOP/s) and memory (GB/s).

To optimize these metrics, we need to focus on two main aspects: data access and computational bottlenecks. This translates to assessing

REDUCE-0: Interleaved Addressing

This method serves as our base. A very naive approach to parallelizing reduction is to decide on a pattern for accessing address spaces where our elements are stored, retrieving those elements, combining those elements by summing them, and recursively repeating this process on different threads to parallelize our operation. In fact, this is what we will do for the first 3 methods of optimization.

Interleaved Addressing refers to accessing and combining address spaces that are positioned halfway to the segment that our current thread is dealing with. Consider an array of 1024 integers. If we use 256 threads per block, each thread starts with a different point and skips every 256 elements. For instance, thread 0 would process elements 0, 256, 512, and 768, each time combining its current element with another positioned halfway to the end of the array segment it’s responsible for. So, thread 0 would combine element 0 with element 128, element 256 with 384, element 512 with 640, and element 768 with 896. Then, this would continue recursively until a final result is reached.

This method not only simplifies synchronization among threads but also ensures that all threads are actively reducing data in parallel, leading to a more balanced load and efficient reduction.

// REDUCTION 0 – Interleaved Addressing__global__ void reduce0(int *g_in_data, int *g_out_data){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;    sdata[tid] = g_in_data[i];    __syncthreads();    // Reduction method -- occurs in shared memory because that's where sdata is stored    for(unsigned int s = 1; s < blockDim.x; s *= 2){        if (tid % (2 * s) == 0) {            sdata[tid] += sdata[tid + s];           }        __syncthreads();    }    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

Let’s go through our first implementation step by step:

Results

Problems with this method

Although this method is a great foundation for parallel programming, it still has its issues. Let’s bring back our metrics and assess where our code might be inefficient in terms of compute and memory.

First, let’s focus on compute-related problems with our next optimization.

REDUCE-1 : Interleaved Addressing 2.0

This method doesn’t change much from our previous method. The addressing is the same, but this time we construct our reduction function without the use of % operator or the divergent condition. By restructuring the index calculation (int index = 2 * s * tid;), REDUCE-1 ensures that each thread consistently performs its operation without checking the position relative to its stride, thereby removing divergence within the warp. This adjustment means all threads in a warp follow the same execution path, significantly improving the warp efficiency. The removal of the % operator further enhances performance by avoiding costly modulo operations, which are slow on GPUs due to their reliance on division.

// REDUCTION 1 – Interleaved Addressing without branch divergence and % operation__global__ void reduce1(int *g_in_data, int *g_out_data){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;    sdata[tid] = g_in_data[i];    __syncthreads();    // Reduction method -- occurs in shared memory    for(unsigned int s = 1; s < blockDim.x; s *= 2){        // note the stride as s *= 2 : this causes the interleaving addressing        int index = 2 * s * tid;    // now we don't need a diverging branch from the if condition        if (index + s < blockDim.x)        {            sdata[index] += sdata[index + s];   // s is used to denote the offset that will be combined        }        __syncthreads();    }    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

Results

Problems with this method

While REDUCE-1 improves on the computational efficiency and execution coherence over REDUCE-0, it introduces a new problem: shared memory bank conflicts. These conflicts occur when multiple threads attempt to access data from the same memory bank simultaneously. These bank conflicts lead to serialization of what could otherwise be parallel memory accesses.

From REDUCE-0 to REDUCE-1, we increased the computational efficiency of our algorithm. However, we did not solve the memory-related issues. In fact, we caused more memory-related issues by switching to strides. It is a bit hard to visualize, but essentially the stride method causes the threads to try and access the same shared memory addresses. REDUCE-0 spread out threads in intervals that acted like “boundaries” and kept thread accesses within those boundaries, reducing the chances of conflicts. But REDUCE-1 relies on strides and removes these boundaries, causing bank conflicts and serialization of processes.

Each bank can only service one access per cycle, so when multiple accesses are directed to the same bank, they must be serialized, effectively reducing the throughput of memory operations. This serialization negates some of the performance gains achieved by eliminating warp divergence and can become a significant bottleneck, especially in larger blocks where the probability of bank conflicts increases. Let’s try to solve this problem now.

REDUCE-2: Sequential Addressing

This method employs a different addressing technique that is more efficient. Instead of threads accessing elements spaced widely apart (interleaved addressing), this method employs sequential addressing where each thread deals with consecutive elements.

Let’s break that down. Bringing back our 1024 element example with 256 threads per block, thread 0 would try and access elements 0, 1, 2, 3 instead of 0, 256, 512, 768 which are spaced far apart. Thread 0 combines elements 0 and 1, then element 2, and so on recursively. What this does is take advantage of spatial locality and avoids bank conflicts by being cache efficient. The algorithm is also linear and minimizes the need for synchronization that increases wait times.

This change significantly enhances memory access patterns by aligning them more closely with the GPU’s preference for coalesced memory accesses. By accessing adjacent memory locations, REDUCE-2 reduces the likelihood of cache misses and memory bank conflicts, making the memory bandwidth usage more efficient and improving overall performance of the reduction operation.

// REDUCTION 2 – Sequence Addressing__global__ void reduce2(int *g_in_data, int *g_out_data){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;    sdata[tid] = g_in_data[i];    __syncthreads();    // Reduction method -- occurs in shared memory    for(unsigned int s = blockDim.x/2; s > 0; s >>= 1){        // REDUCE2 -- check out the reverse loop above        if (tid < s){   // then, we check threadID to do our computation            sdata[tid] += sdata[tid + s];        }        __syncthreads();    }    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

Let’s dive into the algorithm a bit. The major changes in this method include the replacement of strided indexing with a reversed loop structure coupled with threadID-based indexing. This fundamentally modifies how data is handled during reduction.

Results

Problems with this method

This method is mostly conflict free. At this point, we have employed obvious changes to resolve compute and memory issues. We should now try to make our algorithm smarter and find ways to construct it so we can make it faster.

One major problem is that half of the threads are idle in the first loop iteration which is wasteful and underutilizes our GPU’s compute. Although this was the case in our previous techniques, we had bigger fish to fry before getting to idle threads. Going with our 1024 element example, in the first iteration of the loop where s=blockDim.x/2, or s=512 in this case, the condition if (tid < s) restricts active computation to only the first 512 threads of the block. This condition means that while these 512 threads are actively summing pairs of elements (for example, sdata[tid] with sdata[tid + 512]), the remaining 512 threads are idle, contributing nothing to the computation. This pattern of halving the number of active threads in each subsequent iteration continues until the reduction completes; from 512, to 256, then 128, 64, 32 and so on. This rapid halving leads to a significant underutilization of the GPU's capabilities, especially in the initial iterations where only a fraction of the available threads are used.

Let’s solve this problem by doing our first computations when we load our data onto shared memory.

REDUCE-3: First Add During Load

To make use of our idle threads and make our computation smarter, we will do our first computation while we are loading our elements from global memory to shared memory. This will help us load and reduce two elements to one and halve the number of blocks we need to deal with.

More concretely put, in our array of 1024 elements with 256 threads, each thread would load the sum of their first two elements onto shared memory (e.g., thread 0 processes elements 0 and 1, thread 1 processes elements 2 and 3, and so forth). Meaning, that we would halve the number of blocks and the length of our shared memory — in this example, to 512. The rest of the code works exactly as it did before in REDUCE-2. This means that our first iteration would still activate 512 threads to start reducing our elements because of s=blockDim.x/2 = 512. Evidently, this would put more threads to work and avoid any slackers!

// REDUCTION 3 – First Add During Load__global__ void reduce3(int *g_in_data, int *g_out_data){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x*(blockDim.x*2) + threadIdx.x;    sdata[tid] = g_in_data[i] + g_in_data[i+blockDim.x];    __syncthreads();    // Reduction method -- occurs in shared memory    for(unsigned int s = blockDim.x/2; s > 0; s >>= 1){        // check out the reverse loop above        if (tid < s){   // then, we check tid to do our computation            sdata[tid] += sdata[tid + s];        }        __syncthreads();    }    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

During implementation, we see this method unveil as three major changes to our prior code (changes highlighted in bold). We first do our initial reduction step while loading the elements from global memory to shared memory: sdata[tid] = g_in_data[i] + g_in_data[i+blockDim.x] . Then, we make a two changes to accommodate this change:

Result

Problems with this method

Our current approach works great! But, can we make it faster and smarter? Let’s check in with our two metrics. With around 41 GB/s bandwidth usage on Tesla T4, we are definitely not reaching or exhausting our bandwidth. On the other hand, reduction has low arithmetic intensity, meaning we are not compute bound either.

Introducing our new villain…

Because we are not bandwidth bound or compute bound, there is one more bottleneck we can still check for: Instruction Overhead. This includes all the operations, or ancillary instructions, the GPU performs that are not directly related to loading data, storing data, or executing the primary arithmetic operations of the reduction. In other words, these include address arithmetic (calculating with address space to load next) and loop overhead (handling loops, loop conditions, and loop iterations).

Our strategy for this bottleneck would be Loop Unrolling.

REDUCE-4: Unroll Last Warp

Let’s discuss what’s happening in REDUCE-3 first to understand the need for this. With 1024 elements example, After the initial loading where each thread loads and adds pairs of elements, 256 threads work on 512 elements. Here, the reduction begins with each thread working on single elements moving forward. This means,

// Adding this function to help with unrolling__device__ void warpReduce(volatile int* sdata, int tid){  // the aim is to save all the warps from useless work   sdata[tid] += sdata[tid + 32];  sdata[tid] += sdata[tid + 16];  sdata[tid] += sdata[tid + 8];  sdata[tid] += sdata[tid + 4];  sdata[tid] += sdata[tid + 2];  sdata[tid] += sdata[tid + 1];}// REDUCTION 4 – Unroll Last Warp__global__ void reduce4(int *g_in_data, int *g_out_data){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x*(blockDim.x*2) + threadIdx.x;    sdata[tid] = g_in_data[i] + g_in_data[i+blockDim.x];    __syncthreads();    // only changing the end limit to stop before s = 32    for(unsigned int s = blockDim.x/2; s > 32; s >>= 1){          // check out the reverse loop above        if (tid < s){   // then, we check tid to do our computation            sdata[tid] += sdata[tid + s];        }        __syncthreads();    }    // Adding this to use warpReduce when s = 32    if (tid < 32){      warpReduce(sdata, tid);    }    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

The implementation is straightforward enough. We stop our loop before s = 32 and call the kernel warpReduce, with our handwritten 6 iterations, that runs only on __device__. We also need to use the keyword volatile in order for our implementation to still be correct.

Result

Problems with this method

There are definitely no problems with this; we’re getting great speedup! But, why stop our unrolling journey here when we have so many more loops to unroll!!

REDUCE-5: Completely Unroll

In order to continue our unrolling, we would need to know the total number of iterations of our loops at compile time. Luckily for us, the block size is limited to 512 threads by the GPU and we tend to stick to power-of-2 blocks. We know that we can easily unroll for a fixed block size, we just need to be generic. To help with this CUDA supports and provides C++ template parameters in device and host functions.

Templates in C++ allow us to write flexible, generic programs by letting us define functions or classes with placeholders that can be later substituted with specific types provided at compile time. We use this to make account for the potential variations in blockSize that would change unrolling requirements. Depending on the block size, different switch cases are prepared to handle the specific unrolling requirements. This complete unrolling eliminates unnecessary loops and conditions for the majority of the reduction phases, minimizing computational overhead.

By compiling different versions of the kernel tailored to specific block sizes (such as 512, 256, and 128), we optimize each variant for its particular scenario, stripping away any unnecessary operations and maximizing both memory and compute resource efficiency. In this specific implementation, I’ve chosen to set the blockSize to 256 in the main function, simplifying our approach. However, I’ve included switch cases for block sizes of 512, 256, and 128 to demonstrate this method’s flexibility and to highlight how effectively CUDA can leverage template parameters to enhance performance across different configurations.

// Adding this function to help with unrolling and adding the Templatetemplate <unsigned int blockSize>__device__ void warpReduce(volatile int* sdata, int tid){    if(blockSize >= 64) sdata[tid] += sdata[tid + 32];    if(blockSize >= 32) sdata[tid] += sdata[tid + 16];    if(blockSize >= 16) sdata[tid] += sdata[tid + 8];    if(blockSize >= 8) sdata[tid] += sdata[tid + 4];    if(blockSize >= 4) sdata[tid] += sdata[tid + 2];    if(blockSize >= 2) sdata[tid] += sdata[tid + 1];}// REDUCTION 5 – Completely Unrolltemplate <unsigned int blockSize>__global__ void reduce5(int *g_in_data, int *g_out_data){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x*(blockDim.x*2) + threadIdx.x;    sdata[tid] = g_in_data[i] + g_in_data[i+blockDim.x];    __syncthreads();    // Perform reductions in steps, reducing thread synchronization    if (blockSize >= 512) {        if (tid < 256) { sdata[tid] += sdata[tid + 256]; } __syncthreads();    }    if (blockSize >= 256) {        if (tid < 128) { sdata[tid] += sdata[tid + 128]; } __syncthreads();    }    if (blockSize >= 128) {        if (tid < 64) { sdata[tid] += sdata[tid + 64]; } __syncthreads();    }    if (tid < 32) warpReduce<blockSize>(sdata, tid);    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

Also, we should change the way our kernel is called to implement unrolling:

// Needed for Complete unrolling// Launch Kernel and Synchronize threadsswitch (blockSize) {    case 512:        reduce6<512><<<num_blocks, 512, 512 * sizeof(int)>>>(dev_input_data, dev_output_data, n);        break;    case 256:        reduce6<256><<<num_blocks, 256, 256 * sizeof(int)>>>(dev_input_data, dev_output_data, n);        break;    case 128:        reduce6<128><<<num_blocks, 128, 128 * sizeof(int)>>>(dev_input_data, dev_output_data, n);        break;}

The implementation doesn’t change much from REDUCE-4; we simply try to now feed in blockSize as a template parameter that is determined at compile time. Like before, we include if statements to tend to different values of blockSize and switch statements to call kernels based on those values.

Result

Problems with this method

While Reduce5 enhances efficiency by fully unrolling loops for known block sizes, we can’t use this method flexibly and scale it up. Specifically, the full unrolling technique relies heavily on compile-time optimizations that restrict the kernel to fixed block sizes. This approach can lead to inefficiencies in scenarios where the data size does not perfectly match the block configurations, potentially underutilizing GPU resources. Additionally, the complexity of managing multiple versions of the kernel for each block size increases the development overhead and limits dynamic adaptability to varying workloads, making it less practical for general-purpose applications where input sizes can vary greatly.

So, let’s try to get inspiration from First-Add-During-Load from REDUCE-3 and try to do as many Adds as possible instead of just the first one.

REDUCE-6: Multiple Adds / Threads

Transitioning to Reduce6 addresses the rigidity and scalability issues seen in Reduce5 by introducing a more dynamic approach termed “algorithm cascading”. In this method, each thread performs multiple additions within a broader range of block sizes, effectively reducing the dependency on specific block configurations. This flexibility allows the algorithm to adapt more fluidly to varying data sizes, optimizing resource utilization across a wider array of scenarios. By combining both sequential and parallel reductions, Reduce6 minimizes latency and maximizes throughput, particularly in environments with high kernel launch overheads and diverse workload sizes. The strategic distribution of work across threads, as per Brent’s theorem, ensures that each thread contributes optimally throughout the reduction process, maintaining cost-efficiency while scaling effectively with the hardware capabilities.

For example, rather than each thread processing a single pair of elements, it might process multiple pairs before any synchronization barrier, thereby amortizing the cost of synchronization across more computation and improving the overall performance.

Final Optimized Kernel

// Adding this function to help with unrolling and adding the Templatetemplate <unsigned int blockSize>__device__ void warpReduce(volatile int* sdata, unsigned int tid){    if(blockSize >= 64) sdata[tid] += sdata[tid + 32];    if(blockSize >= 32) sdata[tid] += sdata[tid + 16];    if(blockSize >= 16) sdata[tid] += sdata[tid + 8];    if(blockSize >= 8) sdata[tid] += sdata[tid + 4];    if(blockSize >= 4) sdata[tid] += sdata[tid + 2];    if(blockSize >= 2) sdata[tid] += sdata[tid + 1];}// REDUCTION 6 – Multiple Adds / Threadstemplate <int blockSize>__global__ void reduce6(int *g_in_data, int *g_out_data, unsigned int n){    extern __shared__ int sdata[];  // stored in the shared memory    // Each thread loading one element from global onto shared memory    unsigned int tid = threadIdx.x;    unsigned int i = blockIdx.x*(blockSize*2) + tid;    unsigned int gridSize = blockDim.x * 2 * gridDim.x;    sdata[tid] = 0;    while(i < n) {       sdata[tid] += g_in_data[i] + g_in_data[i + blockSize];       i += gridSize;     }    __syncthreads();    // Perform reductions in steps, reducing thread synchronization    if (blockSize >= 512) {        if (tid < 256) { sdata[tid] += sdata[tid + 256]; } __syncthreads();    }    if (blockSize >= 256) {        if (tid < 128) { sdata[tid] += sdata[tid + 128]; } __syncthreads();    }    if (blockSize >= 128) {        if (tid < 64) { sdata[tid] += sdata[tid + 64]; } __syncthreads();    }    if (tid < 32) warpReduce<blockSize>(sdata, tid);    if (tid == 0){        g_out_data[blockIdx.x] = sdata[0];    }}

Peep the while loop where each thread performs multiple additions directly in shared memory. This loop is designed to aggregate two data elements per thread for each iteration, effectively halving the number of necessary operations and interactions with global memory. The thread loads data from the global memory, adds it to a previously loaded value, and then jumps forward by the total number of threads times two, ensuring that it processes another pair of elements on the next iteration. This pattern significantly reduces the total amount of data each thread needs to handle at any one time, maximizing the use of available bandwidth and minimizing latency.

FINAL RESULTS

Comparing it with NVIDIA’s Performance metrics

One of the main differences between my implementation and NVIDIA’s is in GPU. For the webinar, they use GeForce 8800, while I used Tesla T4. This made my initial implementation a lot better right away than theirs, because of a more optimized architecture. However, it also left very little space for improvement in speedup. While I am not able to match the dramatic speedups, I am able to showcase continuous optimization and increasing GPU peak performance.

Generalizing Optimization Techniques

I’m simply going to list down my key takeaways while optimizing a CUDA kernel:

I hope this was helpful!

Source: https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf

My GitHub Implementation: https://github.com/rimikadhara67/Parallel-Reduction?tab=readme-ov-file

--

--

Written by Rimika Dhara

Aspiring SWE and ML Engineer. Current SWE Intern at Oracle.

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Understanding CUDA Memory Usage: A Practical Guide

Hey Amit

Follow

--

Listen

Share

If you think you need to spend $2,000 on a 180-day program to become a data scientist, then listen to me for a minute.

I understand that learning data science can be really challenging, especially when you’re just starting out, because you don’t know what you need to know.

But it doesn’t have to be this way.

That’s why I spent weeks creating the perfect roadmap to help you land your first data science job.

Here’s what it contains:

If this sounds exciting, you can grab it right now by clicking here.

Now let’s get back to the blog:

Introduction

Purpose: Why Understanding CUDA Memory Usage Matters

Let’s face it: when it comes to deep learning or scientific computing, GPUs are like rocket fuel. But here’s the thing — no matter how powerful your GPU is, poor memory management can leave you stuck in first gear.

CUDA memory is more than just a tool for moving data back and forth; it’s the backbone of your computational performance.

Efficient memory usage doesn’t just save milliseconds; it can save hours, dollars, and even your sanity on a tight project deadline.

If you’ve ever run into out-of-memory errors or struggled to squeeze more speed out of your GPU, this guide is for you.

I’m not going to hand you surface-level tips that you’ve heard a hundred times.

Instead, we’ll dive into the nuances of CUDA memory types, debugging, and optimization — things you can actually apply to your real-world projects to gain measurable improvements.

Reader Takeaway

By the end of this guide, you’ll not only understand CUDA’s memory architecture but also how to use it to unlock the full potential of your GPU.

Whether you’re optimizing neural networks or accelerating scientific simulations, the insights here will help you avoid costly mistakes and make smarter design choices.

Quick Overview

Here’s what we’re diving into:

Pro Tip:

This might surprise you: inefficient memory usage can make a high-end GPU perform worse than a mid-tier one.

Studies have shown that optimizing memory access patterns can reduce training times by up to 30%. Imagine what that means for your projects — whether it’s faster results or slashed cloud bills.

2. CUDA Memory Types and Architecture Overview

Memory Types: What You Need to Know

Think of CUDA memory like a multi-tiered storage system. Each tier has its own strengths and weaknesses, and understanding these differences is key to designing efficient kernels. Let’s break it down:

1. Global Memory

Code Example:Here’s a basic kernel where threads read from global memory:

__global__ void globalMemoryExample(float *input, float *output, int size) {    int idx = blockIdx.x * blockDim.x + threadIdx.x;    if (idx < size) {        output[idx] = input[idx] * 2.0f;  // Simple operation    }}

Why this matters: If threads access global memory inefficiently (e.g., uncoalesced access), you’re leaving performance on the table.

2. Shared Memory

Real-World Application: Accelerating matrix multiplication by tiling submatrices into shared memory.

Code Example:

__global__ void sharedMemoryExample(float *input, float *output, int N) {    __shared__ float tile[BLOCK_SIZE][BLOCK_SIZE];  // Shared memory array    int tx = threadIdx.x, ty = threadIdx.y;    int row = blockIdx.y * blockDim.y + ty;    int col = blockIdx.x * blockDim.x + tx;    if (row < N && col < N) {        tile[ty][tx] = input[row * N + col];        __syncthreads();  // Synchronize all threads in the block        output[row * N + col] = tile[ty][tx] * 2.0f;    }}

Pro Tip: Always use __syncthreads() when accessing shared memory to prevent race conditions.

3. Local Memory

4. Constant Memory

Code Example:

__constant__ float kernel[16];  // Constant memory array__global__ void constantMemoryExample(float *input, float *output, int size) {    int idx = threadIdx.x + blockIdx.x * blockDim.x;    if (idx < size) {        output[idx] = input[idx] * kernel[threadIdx.x % 16];    }}

5. Texture and Surface Memory

6. Unified Memory

Memory Hierarchy Visualization

Here’s the deal: the farther the memory is from your threads, the slower it is. A diagram comparing latencies and bandwidths (e.g., shared memory vs. global memory) will drive this point home.

Code Snippet: Performance Comparison

To truly understand the difference between memory types, let’s measure their performance in a simple kernel:

__global__ void memoryComparison(float *globalData, float *sharedData, int size) {    __shared__ float sharedBuffer[BLOCK_SIZE];    int idx = threadIdx.x + blockIdx.x * blockDim.x;    // Global memory access    if (idx < size) {        globalData[idx] *= 2.0f;      }    // Shared memory access    if (threadIdx.x < BLOCK_SIZE) {        sharedBuffer[threadIdx.x] = sharedData[threadIdx.x] * 2.0f;        __syncthreads();    }}

3. Real-World Project: Optimizing CUDA Memory in Image Processing

Context

Let’s optimize a real-world problem: image convolution. It’s a common operation in computer vision but often suffers from slow execution due to poor memory handling.

Step 1: Set Up the Problem

Suppose you’re processing a high-resolution image for edge detection. Each pixel involves a weighted sum of its neighbors using a kernel matrix.

Challenges:

Step 2: Naive CUDA Implementation

Here’s a basic kernel:

__global__ void naiveConvolution(float *input, float *output, int width, int height, float *kernel, int kernelSize) {    int x = blockIdx.x * blockDim.x + threadIdx.x;    int y = blockIdx.y * blockDim.y + threadIdx.y;    int halfK = kernelSize / 2;    if (x >= halfK && x < width - halfK && y >= halfK && y < height - halfK) {        float sum = 0.0f;        for (int i = -halfK; i <= halfK; i++) {            for (int j = -halfK; j <= halfK; j++) {                sum += input[(y + i) * width + (x + j)] * kernel[(i + halfK) * kernelSize + (j + halfK)];            }        }        output[y * width + x] = sum;    }}

Profiling the Naive Kernel

Use tools like Nsight Systems to measure memory usage and identify inefficiencies. Expect to see:

In the next section, we’ll optimize this kernel with shared memory to reduce global memory traffic and improve performance. Stay tuned!

4. Debugging and Profiling CUDA Memory

When it comes to CUDA programming, debugging memory issues and profiling performance can feel like untangling a stubborn knot. But let me tell you — if you don’t, that knot can strangle your performance. Here, I’ll show you how to tackle CUDA memory debugging and profiling like a pro.

Profiling Tools: Finding Bottlenecks with Nsight Systems and nvprof

Before we fix any issues, we need to find them. This is where profiling tools shine.

Nsight Systems and Nsight ComputeThese tools are your best friends when profiling CUDA memory. They let you analyze memory transfers, kernel execution, and thread synchronization. Let’s walk through a basic workflow:

Run your application with Nsight Systems:

nsys profile --output=my_profile ./my_cuda_app

This command generates a detailed report showing kernel execution timelines and memory transfer patterns.

Dive deeper with Nsight Compute:Once you identify a problematic kernel, analyze its memory usage:

ncu --set full ./my_cuda_app

Nsight Compute provides metrics like memory throughput, shared memory utilization, and L2 cache hit rates.

nvprofIf you’re working on older GPUs or prefer a CLI-first approach, nvprof is another powerful tool.

nvprof --print-gpu-trace ./my_cuda_app

This command outputs a timeline of kernel launches and memory transfers.

To isolate memory operations:

nvprof — metrics dram_read_throughput,dram_write_throughput ./my_cuda_app

Pro Tip: Use these tools early and often during development to catch bottlenecks before they snowball.

Debugging Memory Issues

Let’s move from profiling to debugging common CUDA memory issues.

1. Out-of-Memory (OOM) Errors

You’ve probably hit an OOM error if your training crashes mid-run. These errors often occur when VRAM is exhausted due to excessive memory allocations.

Solution:

Monitor memory usage with cudaMemGetInfo:

size_t freeMem, totalMem;cudaMemGetInfo(&freeMem, &totalMem);std::cout << "Free memory: " << freeMem << ", Total memory: " << totalMem << std::endl;

2. Misaligned Memory Access

Misaligned access happens when threads access non-contiguous memory addresses. It can kill performance and lead to subtle bugs.

Solution:

Ensure data structures are aligned using cudaMallocPitch for 2D arrays:

float *d_array;size_t pitch;cudaMallocPitch(&d_array, &pitch, width * sizeof(float), height);

3. Debugging Shared Memory Bank Conflicts

Bank conflicts occur when multiple threads access the same memory bank simultaneously, causing serialization.

Example of a Problematic Kernel:

__global__ void sharedMemoryBankConflict(float *input, float *output, int N) {    __shared__ float sharedArray[32];    int idx = threadIdx.x + blockIdx.x * blockDim.x;    if (idx < N) {        sharedArray[threadIdx.x] = input[idx];        __syncthreads();        output[idx] = sharedArray[threadIdx.x];    }}

Here’s the issue: all threads in a warp access consecutive elements in sharedArray, leading to conflicts.

How to Debug and Fix It:

Use cuda-memcheck to detect conflicts

cuda-memcheck ./my_cuda_app

Fix the kernel by padding shared memory:

__global__ void sharedMemoryNoConflict(float *input, float *output, int N) {    __shared__ float sharedArray[32 + 1];  // Add padding to avoid conflicts    int idx = threadIdx.x + blockIdx.x * blockDim.x;    if (idx < N) {        sharedArray[threadIdx.x] = input[idx];        __syncthreads();        output[idx] = sharedArray[threadIdx.x];    }}

Padding prevents threads from hitting the same memory bank, resolving conflicts and boosting performance.

5. Optimization Strategies for CUDA Memory Usage

Optimizing CUDA memory usage is like fine-tuning a race car. Each tweak shaves off precious milliseconds. Let’s explore some advanced strategies.

1. Memory Coalescing: The Key to Speed

Why it matters: Coalesced memory access ensures that threads within a warp access consecutive memory addresses, minimizing latency.

Example of Uncoalesced Access:

__global__ void uncoalescedAccess(float *input, float *output, int N) {    int idx = threadIdx.x * blockDim.x + blockIdx.x;    if (idx < N) output[idx] = input[idx];  // Poor access pattern}

Fixed Coalesced Access:

__global__ void coalescedAccess(float *input, float *output, int N) {    int idx = threadIdx.x + blockIdx.x * blockDim.x;    if (idx < N) output[idx] = input[idx];  // Consecutive access}

Benchmark Both Approaches: Use Nsight Compute to compare throughput. You’ll see a stark difference in performance.

2. Shared Memory Optimization with Tiling

Tiling involves loading chunks of data into shared memory to reduce redundant global memory accesses.

Example: Matrix Multiplication

__global__ void tiledMatrixMul(float *A, float *B, float *C, int N) {    __shared__ float tileA[BLOCK_SIZE][BLOCK_SIZE];    __shared__ float tileB[BLOCK_SIZE][BLOCK_SIZE];    int tx = threadIdx.x, ty = threadIdx.y;    int row = blockIdx.y * blockDim.y + ty;    int col = blockIdx.x * blockDim.x + tx;    float sum = 0.0f;    for (int k = 0; k < N / BLOCK_SIZE; ++k) {        tileA[ty][tx] = A[row * N + k * BLOCK_SIZE + tx];        tileB[ty][tx] = B[(k * BLOCK_SIZE + ty) * N + col];        __syncthreads();        for (int i = 0; i < BLOCK_SIZE; ++i) {            sum += tileA[ty][i] * tileB[i][tx];        }        __syncthreads();    }    C[row * N + col] = sum;}

Tiling significantly improves memory efficiency by reusing loaded data within the block.

3. Dynamic Parallelism

Dynamic parallelism allows child kernels to handle localized computations, optimizing memory usage.

Example: Recursive matrix subdivision. Launch child kernels for smaller submatrices, reducing memory overhead.

6. Case Study: Optimizing ResNet Training

When training a ResNet, VRAM is your most precious resource. Let’s address common challenges and solutions.

Challenges

Solutions

torch.cuda.empty_cache() # Free unused memory

2. Pinned Memory:

Use pinned host memory for faster host-to-device transfers:

cudaHostAlloc((void**)&hostBuffer, size, cudaHostAllocDefault);cudaMemcpyAsync(deviceBuffer, hostBuffer, size, cudaMemcpyHostToDevice, stream);

7. Advanced Topics in CUDA Memory

As you master CUDA’s memory model, you’ll start craving more advanced tools and techniques to squeeze every ounce of performance from your GPU. This section explores some high-level topics that can elevate your CUDA skills from proficient to expert.

Unified Memory in Multi-GPU Systems

Real-World Scenario: Imagine you’re training a massive transformer model across multiple GPUs. Each GPU holds a subset of the model and needs to share intermediate results with others. Manually managing data movement can quickly turn into a logistical nightmare. Here’s where Unified Memory shines.

What is Unified Memory?Unified Memory provides a single memory address space accessible by both the CPU and GPUs. In multi-GPU setups, it simplifies data sharing because you don’t need to explicitly transfer data between devices.

Example:

__global__ void unifiedMemoryKernel(float *data, int size) {    int idx = threadIdx.x + blockIdx.x * blockDim.x;    if (idx < size) {        data[idx] *= 2.0f;    }}int main() {    float *data;    size_t size = 1024 * sizeof(float);    // Allocate Unified Memory    cudaMallocManaged(&data, size);    // Initialize data    for (int i = 0; i < 1024; ++i) {        data[i] = i;    }    // Launch kernel    unifiedMemoryKernel<<<32, 32>>>(data, 1024);    cudaDeviceSynchronize();    // Free Unified Memory    cudaFree(data);    return 0;}

Caveats:While Unified Memory simplifies programming, it comes with trade-offs:

Pro Tip: Use cudaMemAdvise to provide memory access hints, reducing page thrashing in multi-GPU systems.

cudaMemAdvise(data, size, cudaMemAdviseSetPreferredLocation, device_id);

Asynchronous Memory Transfers

Here’s the deal: In CUDA, memory transfers between the host and device can often block computation. But by overlapping memory transfers with computation, you can dramatically speed up your pipeline.

How It Works: CUDA streams allow you to execute memory transfers asynchronously, so data is being moved while the GPU crunches numbers.

Code Example:

cudaStream_t stream;cudaStreamCreate(&stream);float *hostData, *deviceData;cudaMallocHost(&hostData, size);  // Pinned memorycudaMalloc(&deviceData, size);// Asynchronous memory copycudaMemcpyAsync(deviceData, hostData, size, cudaMemcpyHostToDevice, stream);// Launch kernel on the same streammyKernel<<<grid, block, 0, stream>>>(deviceData);// Synchronize the streamcudaStreamSynchronize(stream);// Free resourcescudaStreamDestroy(stream);cudaFree(deviceData);cudaFreeHost(hostData);

Why It Matters: By overlapping transfers and computation, you can achieve a significant performance boost, especially in data-intensive workflows.

Persistent Memory (CUDA 11.x+)

Persistent memory buffers were introduced in CUDA 11 to address scenarios where you need long-lived data storage for kernels. Think of it as a way to keep your data “warm” across kernel launches, reducing initialization costs.

Use Case: Storing intermediate results for iterative algorithms like conjugate gradient solvers.

Code Example:

cudaDeviceEnablePeerAccess(1, 0);  // Enable peer access between devicesvoid *persistentBuffer;cudaMallocAsync(&persistentBuffer, size, 0);// Use persistentBuffer in multiple kernel launcheskernelA<<<grid, block>>>(persistentBuffer);kernelB<<<grid, block>>>(persistentBuffer);cudaFreeAsync(persistentBuffer, 0);  // Free when done

Pro Tip: Combine persistent memory with asynchronous operations for high-efficiency pipelines.

8. Best Practices for CUDA Memory Management

Managing CUDA memory is like managing your finances — it’s all about discipline and avoiding costly mistakes. Here’s a checklist to keep your memory usage efficient and error-free.

Checklist for Efficient Memory Usage

Common Pitfalls to Avoid

9. Conclusion

Let’s wrap this up.

CUDA memory isn’t just a technical detail — it’s the foundation of high-performance GPU computing. By understanding the nuances of memory types, debugging tools, and optimization strategies, you can transform your code into a lean, efficient powerhouse.

Key Insights:

Call-to-Action: Now it’s your turn. Apply these techniques to your projects and measure the difference. I’d love to hear about the performance gains you achieve — drop a comment or share your benchmarks!

--

--

Written by Hey Amit

Data Scientist & Founder at: oneiszero.com

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Decoding Attention-LLM Inference Optimization

How to optimize MHA in the decoding stage of LLM inference?

Bruce-Lee-LY

Follow

--

Listen

Share

1 Background

Thanks to flash-attention’s accelerated optimization of MHA (Multi Head Attention) calculations, the performance of LLM training and inference has been greatly improved. Tri Dao specially wrote Flash Decoding based on Flash Attention v2 to optimize inference. It mainly modified the block division method, used multiple blocks to process the same attention head, realized parallel loading of KV Cache, and finally launched a single kernel to rescale and merge the results. Therefore, in the case of small batches and very long sequences, the inference performance of Flash Decoding is greatly improved.

However, when Flash Attention is applied to LLM inference, there are still the following problems:

Since seq_q in the decoding stage of LLM inference is always 1, the calculation process of each head attention of MHA is simplified to two HGEMV and one softmax, as shown in the schematic diagram below.

It is noted that in the case of RTX3090 in FP32 multiplication and accumulation, the computing power of Tensor Core FP16 is only 2 times that of CUDA Core FP32. In this case, the handwritten CUDA Core kernel is used to calculate the decoding MHA, while ensuring the accuracy, regardless of the delay or hardware utilization will achieve certain benefits. On the other hand, the handwritten decoding MHA kernel not only facilitates the further use of KV cache quantification methods to improve the throughput performance of LLM inference, thereby reducing the cost of inference, but can also support more variants of attention (such as ALiBi, etc.).

2 Result

This article mainly refers to ppl.llm.kernel.cuda and flash-attention, and uses CUDA Core to optimize the performance of MHA in the decoding stage of LLM inference. Currently, in some inference decoding scenarios, the performance of Decoding Attention is better than Flash Decoding (Flash Attention) and FlashInfer. Decoding Attention supports GQA (Group Query Attention) / MQA (Multi Query Attention) and ALiBi (Attention with Linear Biases) inference scenarios and supports both FP16 and BF16 data types. Decoding Attention provides C++ API and Python API and the code is open source in decoding_attention.

2.1 Test Conditions

2.2 Equipment Specifications

The device specifications of RTX3090 are as follows.

2.3 RTX3090

（1）Seq Len

The performance of Decoding Attention is better when the sequence length is below 1536, while the performance of Flash Decoding (Flash Attention) and FlashInfer is better when the sequence length is above 1536.

（2）Batch Size

Regardless of bacth size, Decoding Attention has better performance than Flash Decoding (Flash Attention) and FlashInfer.

3 Decoding Attention

As mentioned above, seq_q in the decoding stage of LLM inference is always 1, so the calculation process of each head attention of MHA is simplified to two HGEMV and one softmax. Decoding Attention divides blocks according to batch and head. Kernel is mainly divided into three parts, namely HGEMV (S = Q * K^T), softmax (P = Softmax(S)) and HGEMV (O = P * V). The source code is in decoding_attention.

3.1 HGEMV（S = Q * K^T）

Each group processes one or more seqlen_k, and the calculation process is consistent with HGEMV, refer to cuda_hgemv.

// S = Q * K^T    T RQ[thread_elem_nums];#pragma unroll    for (size_t i = 0; i < thread_iters; ++i) {        *(int4 *)(&RQ[i * thread_copy_elem_nums]) =            *(int4 *)(&q_ptr[binfo.q_offset(params.q_row_stride, params.q_head_stride,                                            (i * threads_per_group + group_lane_id) * thread_copy_elem_nums)]);    }    extern __shared__ float S_smem[];    float S_max = -std::numeric_limits<float>::max();#pragma unroll    for (size_t base_seq_k = warp_id * groups_per_warp; base_seq_k < binfo.actual_seq_k;         base_seq_k += groups_per_block) {        size_t seq_k = base_seq_k + group_id;        T RK[thread_elem_nums];        float acc = 0.0;        if (seq_k < binfo.actual_seq_k) {#pragma unroll            for (size_t i = 0; i < thread_iters; ++i) {                *(int4 *)(&RK[i * thread_copy_elem_nums]) =                    *(int4 *)(&k_ptr[binfo.k_offset(seq_k, params.k_row_stride, params.k_head_stride,                                                    (i * threads_per_group + group_lane_id) * thread_copy_elem_nums)]);            }#pragma unroll            for (size_t i = 0; i < thread_elem_nums; ++i) {                if constexpr (std::is_same_v<T, half>) {                    acc += (__half2float(RQ[i]) * __half2float(RK[i]));                } else {                    acc += (__bfloat162float(RQ[i]) * __bfloat162float(RK[i]));                }            }        }#pragma unroll        for (size_t i = threads_per_group / 2; i >= 1; i /= 2) {            acc += __shfl_xor_sync(shfl_mask, acc, i);        }        if (group_lane_id == 0 && seq_k < binfo.actual_seq_k) {            acc *= params.scale_softmax;            if (IsAlibi) {                acc += (binfo.h_slope * (static_cast<int>(seq_k) - binfo.actual_seq_q - binfo.row_shift));            }            S_smem[seq_k] = acc;            S_max = fmaxf(acc, S_max);        }    }

3.2 Softmax（P = Softmax(S)）

First, reduce is performed based on the maximum value of S in each group calculated in the previous step to obtain the maximum value of S in a row, and then the softmax corresponding to each seqlen_k is calculated on S_smem.

// P = Softmax(S)    __shared__ float softmax_smem[warps_per_block];#pragma unroll    for (size_t i = warp_size / 2; i >= 1; i /= 2) {        S_max = fmaxf(S_max, __shfl_xor_sync(shfl_mask, S_max, i));    }    if (lane_id == 0) {        softmax_smem[warp_id] = S_max;    }    __syncthreads();    if (lane_id < warps_per_block) {        S_max = softmax_smem[lane_id];    } else {        S_max = -std::numeric_limits<float>::max();    }#pragma unroll    for (size_t i = warps_per_block / 2; i >= 1; i /= 2) {        S_max = fmaxf(S_max, __shfl_xor_sync(shfl_mask, S_max, i));    }    S_max = __shfl_sync(shfl_mask, S_max, 0);    float exp_sum = 0.0;#pragma unroll    for (size_t seq_k = threadIdx.x; seq_k < binfo.actual_seq_k; seq_k += threads_per_block) {        S_smem[seq_k] -= S_max;        S_smem[seq_k] = exp(S_smem[seq_k]);        exp_sum += S_smem[seq_k];    }#pragma unroll    for (size_t i = warp_size / 2; i >= 1; i /= 2) {        exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);    }    if (lane_id == 0) {        softmax_smem[warp_id] = exp_sum;    }    __syncthreads();    if (lane_id < warps_per_block) {        exp_sum = softmax_smem[lane_id];    }#pragma unroll    for (size_t i = warps_per_block / 2; i >= 1; i /= 2) {        exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);    }    exp_sum = __shfl_sync(shfl_mask, exp_sum, 0);#pragma unroll    for (size_t seq_k = threadIdx.x; seq_k < binfo.actual_seq_k; seq_k += threads_per_block) {        S_smem[seq_k] /= exp_sum;    }    __syncthreads();

3.3 HGEMV（O = P * V）

Due to the particularity of V matrix storage, each group here calculates the outer product of each row or multiple rows in V, and then Reduce Sum gets the final result.

// O = P * V    T RV[thread_elem_nums];    float RO[thread_elem_nums];    memset(RO, 0, sizeof(RO));#pragma unroll    for (size_t base_seq_k = warp_id * groups_per_warp; base_seq_k < binfo.actual_seq_k;         base_seq_k += groups_per_block) {        size_t seq_k = base_seq_k + group_id;        if (seq_k < binfo.actual_seq_k) {#pragma unroll            for (size_t i = 0; i < thread_iters; ++i) {                *(int4 *)(&RV[i * thread_copy_elem_nums]) =                    *(int4 *)(&v_ptr[binfo.k_offset(seq_k, params.v_row_stride, params.v_head_stride,                                                    (i * threads_per_group + group_lane_id) * thread_copy_elem_nums)]);            }#pragma unroll            for (size_t i = 0; i < thread_elem_nums; ++i) {                if constexpr (std::is_same_v<T, half>) {                    RO[i] += (S_smem[seq_k] * __half2float(RV[i]));                } else {                    RO[i] += (S_smem[seq_k] * __bfloat162float(RV[i]));                }            }        }    }#pragma unroll    for (size_t i = 0; i < thread_elem_nums; ++i) {#pragma unroll        for (size_t j = threads_per_group; j <= warp_size / 2; j *= 2) {            RO[i] += __shfl_xor_sync(shfl_mask, RO[i], j);        }    }    __syncthreads();#pragma unroll    for (size_t i = threadIdx.x; i < head_dim; i += threads_per_block) {        S_smem[i] = 0.0;    }    __syncthreads();    if (lane_id < threads_per_group) {#pragma unroll        for (size_t i = 0; i < thread_iters; ++i) {#pragma unroll            for (size_t j = 0; j < thread_copy_elem_nums; ++j) {                atomicAdd(S_smem + (i * threads_per_group + lane_id) * thread_copy_elem_nums + j,                          RO[i * thread_copy_elem_nums + j]);            }        }    }    __syncthreads();#pragma unroll    for (size_t i = threadIdx.x; i < head_dim; i += threads_per_block) {        if constexpr (std::is_same_v<T, half>) {            o_ptr[binfo.q_offset(params.o_row_stride, params.o_head_stride, i)] = __float2half(S_smem[i]);        } else {            o_ptr[binfo.q_offset(params.o_row_stride, params.o_head_stride, i)] = __float2bfloat16(S_smem[i]);        }    }

4 Other

4.1 Next Plan

--

--

Written by Bruce-Lee-LY

LLM Infer, AI Infra, CUDA

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Introduction to GPU Programming with Python & CUDA

Sequential programming is really hard, parallel programming is a step beyond that — Andrew S. Tanenbaum

Geminae Stellae 💫

Follow

--

2

Listen

Share

Keywords

Before diving into the topic, we would like to define some concepts related to parallel computing:

CPU: The Central Processing Unit, is the processor installed at the heart of a computer. It receives information, processes it and distributes it to the screen or to the other units connected to it such as graphics card, sound card, printer, scanner, etc.

GPU: The Graphics Processing Unit (GPU) is a specialized processing unit, mainly designed to process images and videos. GPUs process data in order to render images on an output device, such as a screen. However, modern GPUs are general purpose computing devices that can be used to perform any kind of computation.

CuPy: A GPU array library that implements a subset of the NumPy and SciPy interfaces. It is a convenient tool for those familiar with NumPy to explore the power of GPUs, without the need to write code in a GPU programming language like CUDA & OpenCL.

Host & Device: The host is often used to refer to the the CPU, while device is used to refer to the GPU.

Thread: The smallest execution unit in a CUDA program.

Block: A set of CUDA threads sharing resources.

Grid: A set of blocks launched in one kernel.

Kernel: A large parallel loop, where each thread executes one iteration.

Programming environment

The aim of this article is to learn how to write optimized code on GPU using both CUDA & CuPy. For this, we will be using either Jupyter Notebook, a programming environment that runs in a web browser. Or Google Colab, in case you do not have any GPU available on your laptop.

Note: Don’t forget to change the runtime type from None to GPU or CPU

GPU vs CPU

CPUs are known of their capability to perform any general purpose computation in a high speed, so what’s the point behind using GPUs for the same purpose?

Well, simply because GPUs are designed to process images & videos.

CPUs cannot handle parallel processing, therefore large tasks that require millions of similar operations will choke a CPU’s capacity to process data. GPUs, on the other hand, are massively parallel devices that can execute thousands of threads at the same time.

That being said, you might still wonder why anyone would prefer GPU over CPU to compute something that can be easily computed on a CPU.

Let’s answer this question with a simple example:

Sorting an array

We will be creating a large array in python, then sorting it with the sort() function of NumPy on the CPU.

import numpy as npsize = 8192 * 8192array = np.random.random(size).astype(np.float32)

Tip: To compute the time of the sorting we can run the following command

%timeit -n 1 -r 1 result = np.sort(array)

So, it took 7.84 s to sort the array.

Now let’s perform the same sorting, but this time on a GPU. We will be using the sort() function from CuPy

import cupy as cparray_gpu = cp.asarray(array)%timeit -n 7 -r 1 result_gpu = cp.sort(array_gpu)

36.6 ms, that’s faster!

Speedup

From the results, we noticed that sorting the array with CuPy, i.e. using the GPU, is faster than with NumPy, using the CPU.

To quantify the speed, we will compute the speedup of using CuPy over NumPy. The speedup is defined as the ratio between the sequential (NumPy in our case) and parallel (CuPy in our case) execution times.

Note: Execution times must be in the same unit (milliseconds or seconds)

speedup = 7.84 / 0.0366print(speedup)

We can therefore say that by only using the GPU with CuPy to sort an array of size 8192 * 8192 we achieved a performance improvement of 214 times.

Convolution in Python

We start by generating an image on the host using Python and NumPy. Basically, it’s an image filled with zeros, except for isolated pixels with value one, on a regular grid. The plan is to convolve it with a Bicubic interpolation. Then record the time it takes to execute this convolution on the host.

import numpy as np# Construct an image with repeated delta functionsdeltas = np.zeros((4096, 4096))deltas[8::16,8::16] = 1

Visualization

import pylab as pyl# Necessary command to render a matplotlib image in a Jupyter notebook.%matplotlib inline# Display the imagepyl.imshow(deltas[0:100, 0:100])pyl.show()

The computation we want to perform on this image is a convolution, first on the host then on the device to compare the results and execution times. We can think of an image as a matrix of colour values, so when we convolve that image with a filter, we generate a new matrix with different colour values.

In our example, we will convolve our image with a Bicubic interpolation shown in the illustrations:

The BICUBIC Interpolation method is known for its ideal combination of processing time and output quality. That’s why it’s often used in image editing programs like Adobe Photoshop.

Convolution on a CPU — Scipy

Let us first construct the Bicubic, and then display it.

x, y = np.meshgrid(np.linspace(-2, 2, 15), np.linspace(-2, 2, 15))dst = np.sqrt(x*x + y*y)sigma = 1muu = 0.000pyl.imshow(dst,interpolation='bicubic')pyl.show()

Now we are ready to do the convolution on the host. We do not have to write this convolution function ourselves, as it is very conveniently provided by SciPy. Let us also record the time it takes to perform this convolution and inspect the top left corner of the convolved image.

from scipy.signal import convolve2d as convolve2d_cpuconvolved_image_using_CPU = convolve2d_cpu(deltas, dst)pyl.imshow(convolved_image_using_CPU[0:100, 0:100])pyl.show()%timeit -n 1 -r 1 convolve2d_cpu(deltas, dst)

Convolution on a GPU — CuPy

Although there is a physical connection — cable — between the CPU and the GPU, they do not share the same memory space. This image depicts the different components of CPU and GPU and how they are connected:

This means that an array created using NumPy is physically located into the main memory of the host, and thus visible to the CPU but not the GPU. It is not yet in GPU memory, so we need to copy the input image and the convolving function to the GPU, before we can execute any code on it. The arrays deltas and dst are in the host’s RAM. Let’s copy them to GPU memory using CuPy.

import cupy as cpdeltas_gpu = cp.asarray(deltas)dst_gpu = cp.asarray(dst)

Execution of the convolution:

from cupyx.scipy.signal import convolve2d as convolve2d_gpuconvolved_image_using_GPU = convolve2d_gpu(deltas_gpu, dst_gpu)%timeit -n 7 -r 1 convolved_image_using_GPU = convolve2d_gpu(deltas_gpu, dst_gpu)

That’s impressive! You see, there is a far cry between the speedup we got when performing a normal computation and the one we got in image redering! That’s the where the power of GPUs comes into play!

Performing NumPy routines on the GPU

We saw above that we cannot execute routines from the cupyx library directly on NumPy arrays. In fact we need to first transfer the data from host to device memory.

convolve2d_gpu(deltas, dst)

Vice versa, if we try to execute a regular SciPy routine (i.e. designed to run the CPU) on a CuPy array, we will also encounter an error.

convolve2d_cpu(deltas_gpu, dst_gpu)

Transfer function

This function is mandatory for any CUDA program execution which is divided to 3 main steps:

def transfer_compute_transferback():    deltas_gpu = cp.asarray(deltas)    dst_gpu = cp.asarray(dst)    convolved_image_using_GPU = convolve2d_gpu(deltas_gpu, dst_gpu)    convolved_image_using_GPU_copied_to_host = cp.asnumpy(convolved_image_using_GPU)

SciPy routines cannot have CuPy arrays as input. Contrary to SciPy, NumPy accepts CuPy arrays, i.e. arrays that exist in GPU memory, as input.

Something we can do, is to execute a linear (1D) convolution with NumPy instead of SciPy. To generate input for a linear convolution, we can flatten our image from 2D to 1D (using ravel()), but we also need a 1D kernel. For the latter we will take the diagonal elements of our 2D dst kernel.

deltas_1d = deltas.ravel()dst_1d = dst.diagonal()%timeit -n 1 -r 1 np.convolve(deltas_1d, gauss_1d)

We have performed a regular linear convolution using CPU, and now, we will transfer the 1D arrays to the GPU and use the NumPy routine to do the convolution.

deltas_1d_gpu = cp.asarray(deltas_1d)gaudst_1d_gpu = cp.asarray(dst_1d)%timeit -n 7 -r 1 np.convolve(deltas_1d_gpu, dst_1d_gpu)

The linear convolution is actually performed on the GPU, which is shown by a nice speedup!

Creating a GPU Kernel

def vector_add(A, B, C, size):    for item in range(0, size):        C[item] = A[item] + B[item]        return C

In this code, each iteration of the for loop is independent from another. So, even if we reorder the iterations, or even compute each iteration in parallel or on a different device, we will still come up with the same output. This kind of programs is called naturally parallel, and they are the best candidates to be executed on a GPU.

We could use CuPy to run something similar to our vector_add function on a GPU. But our aim is to write code that can be executed by GPUs, hence we will be using CUDA.

The CUDA-C language is a GPU programming language and API developed by NVIDIA. It is mostly equivalent to C/C++, with some special keywords, built-in variables, and functions.

Let’s start coding CUDA by writing a small kernel — It’s a GPU program, that computes the same function that we just wrote in Python.

extern "C"__global__ void vector_add(const float * A, const float * B, float * C, const int size){    int item = threadIdx.x;    C[item] = A[item] + B[item];}

In order to compile the code and manage the GPU in Python, we will use the interface provided by CuPy

import cupy# size of the vectorssize = 1024# allocating and populating the vectorsa_gpu = cupy.random.rand(size, dtype=cupy.float32)b_gpu = cupy.random.rand(size, dtype=cupy.float32)c_gpu = cupy.zeros(size, dtype=cupy.float32)# CUDA vector_addvector_add_cuda_code = r'''extern "C"__global__ void vector_add(const float * A, const float * B, float * C, const int size){    int item = threadIdx.x;    C[item] = A[item] + B[item];}'''vector_add_gpu = cupy.RawKernel(vector_add_cuda_code, "vector_add")vector_add_gpu((1, 1, 1), (size, 1, 1), (a_gpu, b_gpu, c_gpu, size))

To confirm that the CUDA code does exactly what we want, let’s run the sequential Python code below which compares the results, it will return ‘Correct results!” if they are similar:

import numpya_cpu = cupy.asnumpy(a_gpu)b_cpu = cupy.asnumpy(b_gpu)c_cpu = numpy.zeros(size, dtype=numpy.float32)vector_add(a_cpu, b_cpu, c_cpu, size)# testif numpy.allclose(c_cpu, c_gpu):    print("Correct results!")

Understanding the CUDA Code

This is the definition of the CUDA vector_add function:

__global__ void vector_add(const float * A, const float * B, float * C, const int size)

Where:

__global__ is a specefic CUDA keyword that identifies the execution space of our kernel. It allows the function to both run on the GPU, and to be called from the host (in our case the Python interpreter running on the CPU). Other execution space identifiers in CUDA-C are:

__host__which identifies a function that can only be run and called from the host — CPU

__device__identifies functions that can inly be run and called from the device — GPU

Python item vs CUDA item

int item = threadIdx.x;C[item] = A[item] + B[item];

In Python, the content of item is the result of the range function. However, in CUDA we are reading a special variable which is threadIdx, it is a triplet that represents the id of a thread inside a 3D CUDA block. In this particular case we are working on a one dimensional vector, and therefore only interested in the first dimension, that is stored in the first field of this variable.

Computing Hierarchy in CUDA

In the previous example we had a small vector of size 1024, where each of the 1024 generated threads was working on one of the elements.

If we change the size of the vector to a larger number, such as 2048, we will get an Error, and the reason behind it is that most of GPUs accept a maximum block size of 1024 threads.

Let’s get back to the parameters of our function:

vector_add_gpu((1, 1, 1), (size, 1, 1), (a_gpu, b_gpu, c_gpu, size))

The first triplet represents the size of the CUDA grid (i.e., Number of Blocks).

The second represents the size of the CUDA block (i.e., Number of Threads).

The grid is a three-dimensional structure in the CUDA programming model and it represents the organization of a whole kernel execution. A grid is made of one or more independent blocks.

In our case, the grid’s size is (1,1,1) which means it is composed of a single block.

The block’s size is (size, 1, 1), where size is the number of threads. Blocks are independent, whereas threads composing a block are not, because they can share resources and communicate with each other.

To solve the issue of the limited number of threads per block, we need -in our case- to divide 2048 on 2 so that we get two blocks of 1024 threads in our grid. This means that the new grid size will change from (1,1,1) to (2,1,1), and the block size from (size,1,1) to (size // 2, 1, 1).

We already introduced the special variable threadIdx when introducing the vector_add CUDA function, it contains a triplet specifying the coordinates of a thread in a thread block.

There are other important CUDA variables that we need to understand. Having the same structure as threadIdx, these variables are:

blockDim: The size of a bloc, i.e. the number of threads per dimension

blockIdx: The ID of a bloc in the grid

gridim: The size of thegrid, i.e. the number of blocks per dimension

Before re-running the code again (after increasing the number of threads), we need to modify the kernel code, precisely, the item value.

What we need to do is to make grid and block more flexible, so that they can adapt to vectors of arbitrary size. To do that, we can replace the Python code to call vector_add_gpu with the following code:

import maththreads_per_block = 1024grid_size = (int(math.ceil(size / threads_per_block)), 1, 1)block_size = (threads_per_block, 1, 1)vector_add_gpu(grid_size, block_size, (a_gpu, b_gpu, c_gpu, size))

Conclusion

From curious hardware to dominant processors that power the world’s fastest supercomputers. In the foreseeable future, it is obvious that GPUs and parallel computing will play a huge role in the computational science, and it is important for researchers and developers to consider how to adapt to these kinds of architectures.

Link to the code

Note: This article is a summary of what we learned from the course GPU Programming

Written with ❤ By Geminae Stellae (Ihssene Brahimi & Assala benmalek)

--

--

2

Written by Geminae Stellae 💫

Quantum computing | Data science | Medical Imaging | Healthcare | Parallel computing | Personal experiences

Responses (2)

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Reproducing GPT-2 (124M): Key Insights and Techniques

Tushar Madaan

Follow

--

1

Listen

Share

In this blog, we’ll explore the key details of reproducing the GPT-2 (124M) model, focusing on its architecture, training process, and performance optimizations specifically w.r.t GPU awareness addition to covering a few innovations that add training stability. The goal is to understand how to efficiently implement and train this transformer-based language model in a stable manner while leveraging modern tools like PyTorch and CUDA for better way faster performance on a single GPU. The code is not my original work but a faithful reproduction of Andrej Karpathy’s lets reproduce GPT2 video and some of my own research to understand a few concepts better. Full code to follow soon → here.

What this blog is not: This isnt a 101 explainer on transformers. If you’re unfamiliar with the basics of transformer architecture and want a foundational understanding, I recommend watching working through Andrej Karpthy’s excellent Let’s build GPT: from scratch, in code, spelled out and Umar Jamil’s YouTube video on transformers both of which further clarified my understanding and are fantastic resources to go through first. However, If you feel you have a rough understanding of transformers and attention mechanism then continue to read on..

What’s not covered? I am also not going to cover tokenization schemes as it is an elaborate topic probably deserving a blog of its own. Tokenization also does not lend itself well to the current theme of training stability and cuda optimizations, since its a preprocessing CPU native implementation — it does indirectly impact architecture design (vocab size etc.) but still not relevant to my current focus.

Why reproduce GPT2? Our focus here is on the specific decisions made in GPT-2’s design and training, especially around ensuring faster convergence, training stability in model training and improving performance from a better understanding of how GPUs and CUDA kernels work.

GPT-2 Architecture Overview

GPT-2 is a decoder-only transformer with 124 million parameters, specifically the smallest model in the GPT-2 family. GPT2 is not very different from a typical decoder only transformer so we wont go over the architecture core components like position encodings, self attention etc. in a ton of detail but hit a few salient aspects that make GPT2 unique.

That said, from an algorithmic point of view, GPT2 does not implement any special optimizations in the attention layer itself to tackle quadratic complexity (V tokens * V token possible edges on a attention graph) and its memory bound nature. There were also later papers like reformer, longformer etc are addressing this bottleneck with sparse attention schemes. Later models like Mistral 7B have implemented sliding window attention as covered in longformer paper. However GPT does leverage Flash Attention with is a form of hardware optimization on GPUs to make the calculation run with faster with a smaller memory footprint — more on that later in the blog.

Weight Sharing: GPT-2 ties the token embeddings and the final output logits, which reduces the parameter count by about 30%. This design makes sense as an inductive bias because it enforces consistency between the input embeddings and output probabilities. The similarity between two tokens should be reflected in both how they are embedded and the logits that predict their likelihoods, ensuring that semantically similar tokens are correlated between in teh input layer and the final output logit layer. Adding this inductive bias reduces parameters (about 30% of total) therefore reducing computational cost and leading to faster convergece.

By leveraging these architectural choices, GPT-2 achieves a good balance between performance and parameter efficiency, while causal self-attention ensures that the text is generated in a meaningful, sequential manner.

The Training Process

Training GPT-2 involves careful data preparation, loss function configuration, and optimizer selection. Here’s what the process looks like:

A quick GPU architecture & memory hierarchy primer

Training a model like GPT-2 on modern GPUs involves understanding the underlying GPU architecture and how memory hierarchies play a role in performance. Here’s an overview of how these elements interact:

Performance Improvements in Reference to GPU Architecture

Understanding how GPUs manage their memory hierarchy and how Tensor Cores accelerate matrix operations is essential to efficiently training large models like GPT-2. These architectural optimizations are key to achieving better performance without excessive hardware demands.

Optimization and Performance Improvements

Given the large size of GPT-2, several performance optimizations are essential to train it efficiently, especially on GPUs. Key optimizations include:

Reproducing GPT-2 (124M) offers valuable insights not only for model pre-training but also for optimizing models in areas like inference, fine-tuning, and deployment. For AI engineers working on real-world applications, the techniques discussed — such as FlashAttention, mixed-precision training, weight sharing, and learning rate schedulers — are essential for creating efficient, scalable models.

These techniques are particularly important in knowledge distillation, where a smaller student model is trained to mimic the behavior of a larger teacher model. For instance, learning rate schedulers ensure that the student model learns smoothly and avoids destabilizing updates, whilemixed-precision training speeds up the distillation process by reducing memory usage and computation time. This allows for faster and more resource-efficient training of the student model.

FlashAttention further reduces memory overhead, allowing the student model to handle longer sequences and more data during training without increasing memory requirements. This is crucial in making both training and inference faster and more efficient, which is often the goal of distilling large models for deployment.

By leveraging these techniques, AI engineers can optimize models for fine-tuning, inference, and knowledge distillation, ensuring that they not only perform well but also scale efficiently across a variety of hardware environments. Whether you’re focused on co-training models, fine-tuning for specific tasks, or deploying large-scale systems, applying these insights will help you achieve better performance and resource efficiency.

Feel free to explore these methods in your own workflows, and reach out if you’d like to dive deeper into specific techniques. Happy optimizing!

--

--

1

Written by Tushar Madaan

Responses (1)

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

CUDA Memory Management & Use cases

Dung Le

Follow

Distributed Knowledge

--

1

Listen

Share

In my previous article, Towards Microarchitectural Design of Nvidia GPUs, I have dissected in-depth a sample GPU architectural design, as well as its memory hierarchy. In order for GPGPU applications to achieve the true performance that GPU promises, a throughout understanding and correct use of each memory space in the hierarchy is a must. In this article, let’s discuss on how to optimally utilize different types of GPU memories and cycle through some notable use cases for each memory type.

The content of this article will be organized as follows:

I. Coalesced & un-coalesced global memory access

The global memory of a CUDA device is implemented with DRAMs. Each time a DRAM location is accessed, a range of consecutive locations that includes the requested location is actually accessed. Many sensors are provided in each DRAM chip and they work in parallel. Each senses the content of a bit within these consecutive locations. Once detected by the sensors, the data from all these consecutive locations can be transferred at very high-speed to the processor. These consecutive locations accessed and delivered are referred to as DRAM bursts.

Recognizing the burst mechanism, current CUDA devices employ a technique that allows the programmers to achieve high global memory access efficiency by organizing memory access of threads into favorable patterns. This technique takes advantage of the fact that threads in a warp execute the same instruction at any given point in time (SIMT). When all threads in a warp execute a load instruction, the hardware detects whether they access consecutive global memory locations. If they do, the hardware combines, or coalesces, all these accesses into a consolidated access to consecutive DRAM locations. Such coalesced access allows the DRAMs to deliver data as a burst.

To optimally utilize global memory, it is important to improve coalescing. There are multiple strategies that can be used. One such strategy is to improve the data access pattern. Let’s take a simple kernel to perform matrix multiplication with a scalar value as an example. In CUDA, multidimensional-array elements are placed into the linearly addressed memory space according to the row-major convention as following:

Now, let’s consider two ways of writing kernel with two different access patterns:

For the above access pattern, as threads’ indexes on x-dimension increase within a warp, they access consecutive locations on the matrix array m. Therefore, when consecutive threads issue load instructions to global memory to access m, these accesses form a coalesced access pattern.

On the other hand, notice that threads within a warp are issuing load instructions in column-major order the above kernel, in which two consecutive accesses will be distant dimx locations away from each other. This results in an un-coalesced access pattern, and henceforth reduce kernel global memory bandwidth.

Another strategy is to change the data layout to improve locality. Depends on how the program accesses data, SOA (structure of arrays) and AOS (array of structs) data type can be used for consecutive threads to issue load/store instructions to consecutive memory locations. For example, computer vision algorithms that need to apply filters onto an image requires the image to be stored onto a data structure. With SOA or AOS data structure, the image can be stored as follows:

Due to the fact that memory addresses of class data members are placed linearly in memory, in case of algorithms whose kernel threads access r, g, b value of each image pixel at the same time, storing the image as SOA type will allow coalesced memory access. On the other hand, in the scenario that kernel threads need to access the values in image’s channels at the same time, AOS type will allow coalesced memory access.

II. Tiling matrix transpose using shared memory

In contrast to global memory which resides in DRAM, shared memory is a type of on-chip memory. This allows shared memory to have a significantly low memory access latency for just several instruction cycles per instruction. One can relate shared memory usage to a CPU cache; however, while CPU cache cannot be explicitly managed, shared memory can. Shared memory can be declared by the programmer by using keyword __shared__, with size hardcoded in the kernel code or passed on explicitly to the kernel call using extern keyword.

With low latency accessing, shared memory is utilized heavily in programs in which memory bound is a problem. One key usage of shared memory comes from the fact that threads within a block can share memory access. Therefore, different threads can use shared variables to hold the data that was reused many times during the computation phase. In order to maximize memory bandwidth, threads can load this data from global memory in a coalesced manner and store it into declared shared memory variables. Threads then can load or write the data in any order due to the fact that shared memory is not affected by un-coalesced read/write order (corner turning technique). For example, in the problem of optimizing matrix transpose performance, this usage of shared memory comes in handily:

In order to have coalesced reads/ writes into global memory, threads load (line 13) and write (line 21) a tile of input data to consecutive memory locations as threadIdx.x varies. At line 13, each block of threads loads a tile of input data having block-sized float numbers from global matrix t into shared memory array. It then makes sure all threads finish loading their element by using __syncthreads, a barrier that synchronizes threads within each block. Each thread then writes into consecutive locations in the global memory as variable to varies in the same amount as threadIdx.x.

III. Optimizing convolution with constant memory, shared memory, and memory padding

The convolution operation (or filtering) is another common operation in many applications, especially in image and signal processing. It consists of source data and a filter (known as mask). By applying the filter against matrix data, we can obtain the convoluted matrix. Let’s divide our use cases into using convolution operation on 1D array and 2D matrix.

1. 1D Convolution

a. Naive convolution kernel

Figure 3 is an example of a convolution operation in a 1D array. For elements that are near the array bounds such as c1 and c6, we substitute 0s to make up for the missing cells (or ghost cells) in the sum that constitute them. Hence, c1, for example, equals to 0 * w1 + i1 * w2 + i2 * w3. A naive kernel for 1D convolution may be written like this:

We can make two observations about the kernel in the naive implementation. First, there will be a control flow divergence. The threads that calculate the output elements near the left end or the right end of the output array have to handle the ghost cells while the other threads do not. Second and a more serious problem is memory bandwidth. For every two calculations in line 7, two global memory reads are issued, which makes the ratio of floating-point arithmetic calculation to global memory access is only about 1.0 in the kernel. Therefore, this simple kernel can only be expected to run at a small fraction of peak performance.

b. Using constant memory for filter

Let’s notice the filter used in convolution operation. First, it’s unchanged throughout the convolution computation. Second, the filter size tends to be small in most cases. Hence, these properties make the filter array a great candidate for constant memory, which is defined with keyword __constant__ in CUDA. Like global memory variables, constant memory variables are also located in DRAM. However, since CUDA knows that constant memory variables are not modified during kernel execution, GPU hardware caches the constant memory data aggressively in L1 or constant cache. In case of a cache hit, data will be served from the cache instead of going down to global memory. Furthermore, the design of caches is typically optimized to broadcast a value to a large number of threads. As a result, when all warp access the same constant memory variable, as in the case of convolution masks, the caches can provide a tremendous amount of bandwidth to satisfy the data needs of threads. Also, as previously stated, the size of the masks is typically small so we can assume that all mask elements are effectively always accessed from caches. Now, let’s turn to the updated kernel code using constant memory for the filter array:

c. Optimize memory bandwidth with shared memory

To further optimize memory bandwidth, notice that each input element is used for more than one output element, and, henceforth, loaded multiple times from global memory. Therefore, to reduce global memory load and improve memory bandwidth, we will pre-load all elements need for the calculation into shared memory. For example, in figure 3, assuming our thread block calculates three output elements c1, c2, c3, the block needs to load all input elements used in c1, c2, c3 calculations, which are i1, i2, i3, and i4. To generalize this idea, let’s look at a sample code snippet:

Threads within a block will collaboratively load (tile_size+ mask_width / 2) input elements used for convolution in tile_size output elements into shared memory. The last threads load first (mask_width / 2) elements, the first threads load the last (mask_width / 2), and the other threads load the rest. This helps reduce the number of global memory loads since each element needed in the input for the convolution is loaded once. To be specific, all threads only need to load (array_width + mask_width) input elements from global memory. On the other hand, in naive kernel implementation, for each idx of a thread that maps to a valid idx index in the output array, each thread needs to load mask_width input elements. In total, all threads in naive implementation issue (mask_width * array_width) global memory transactions.

2. 2D Convolution

a. Memory Padding

Now, let’s move on to 2D matrix convolution implementation. 2D matrix convolution is heavily applied in the realm of computer vision since real-world images are represented as 2D matrices and come in all sizes and shapes. These matrices are generally stored in the row-major layout when reading from files to memory. If the width of the image in terms of bytes is not a multiple of the DRAM burst size, the starting point of row 1 and beyond can be misaligned from the DRAm burst boundaries. Such misalignment can result in poor utilization of DRAM bandwidth as there might be the case where a row is spanned between two DRAM bursts, which requires two DRAM bursts to deliver the whole row data instead of one.

Therefore, in order to utilize the strength of DRAM burst to deliver data fastly, we will pad additional bytes into each row of the image matrices so that each row ends at the DRAM burst boundaries:

After padding, each row of the image matrix will have a length of pitch units. Therefore, when the image is linearized in row-major order in DRAM memory, we need to use pitch instead of width to access the element at (row, col) coordinate: row * pitch + column. However, when we iterate through a row, width is still needed to be used as the loop bound to calculate the elements that actually exist in the original matrix. That brings us to the kernel code for 2D convolution using memory padding:

Each thread of the kernel first calculates the y and x, or col_o,and row_o, indices of its output element (line 5 and 6). It then calculates the y and x indices of the input element it needs to load into the shared memory by substracting (mask_width / 2) from row_o and col_o (line 8 and 9). Applying the same idea of dealing with halo cells in 1D convolution, each thread block loads all halo cells and internal cells needed by all threads in the block for the convolution operation (line 13). Remember that when loading input elements from the padded image, we need to use “pitch” as the row length instead of width (line 14), but when doing convolution operations, we will use “width” as the loop bound (line 26).

So how memory-efficient this kernel is compared to the naive 2D kernel? In a naive kernel, each thread in a thread block will perform (mask_width)² accesses to the image array in global memory. Hence, each thread block issues a total (mask_width)² * (tile_size)² accesses to global memory. In the tiled kernel, all threads in a thread block collectively load one input tile. Therefore, the total number of access by a thread block to the image array is (tile_size + mask_width-1)². Hence, memory access speedup is (mask_width)² * (tile_size)² / (tile_size + mask_width-1)². The larger the ratio, the more effective the tiled kernel in reducing the number of memory accesses as compared to naive kernel.

IV. Takeaways

A throughout understanding and correct use of different memory space in the GPU device will deliver a large speedup for your applications. For data that is unchanged throughout the kernel execution, consider using constant memory to utilize hardware caching for optimized load performance. For data that is accessed repeatedly during the kernel execution, consider using shared memory to reduce global memory loads, and utilize the corner turning technique. When you have to go down to global memory, remember that a coalesced memory read/ write access pattern leads to a much more optimized memory bandwidth than an un-coalesced read/write pattern, as the former can improve its performance based on DRAM bursts mechanism.

V. What’s next?

For the next article, let’s take a look at another two types of CUDA memory, texture memory and unified memory, and ways that they can help to boost GPGPU applications.

Stay tuned for the next one!

--

--

1

Published in Distributed Knowledge

Exploring the realm of Parallel & Distributed Systems — by Dung Le

Written by Dung Le

Previously at @Citadel, @Facebook, @Tesla, @Amazon

Responses (1)

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

FlashAttention: Understanding GPU Architecture-Part 1

Sachin Kalsi

Follow

--

Listen

Share

Introduction

In this three-part blog series, we will delve into the intricate world of FlashAttention, a technology that has been making waves in the field of large language models (LLMs). FlashAttention is a novel approach to optimizing the attention mechanism, making it faster and more memory-efficient.

Before we explore FlashAttention in detail, it’s crucial to have a solid grasp of the underlying hardware, specifically GPU (Graphics Processing Unit) architecture, as FlashAttention leverages the GPU for efficient execution. Let’s break down the essential concepts:

How Data Moves Through a System

To understand how data moves within a system, let’s begin with a simple breakdown. Data typically starts on the hard disk (HDD drive). However, for processing, we need this data in the Random Access Memory (RAM), also called main memory. Depending on the system, there may be multiple layers of memory, each with varying speeds and sizes.

It is indeed essential to minimize data movement between these storage levels to optimize computation. This is because the time it takes to access data increases as you move further from the CPU, and data transfer can be a bottleneck in computing tasks.

Understanding GPU Memory Hierarchy

In the case of NVIDIA A100 80GB GPUs, for instance, HBM may be implemented using 6 vertical stacks, where each stack, except one, contains 16GB of memory, summing up to the total GPU memory.

Remember, to fully utilize GPU processing power, you need to ensure that data can be moved efficiently between GPU memory, cache, and processing units.

FlashAttention and GPU Memory

One of the challenges in training large language models is efficient memory usage. As the demand for GPU memory increases, optimizing the memory hierarchy becomes crucial. FlashAttention is designed to maximize the utilization of GPU memory and leverage its high-speed components, such as tensor cores.

The process of optimizing involves increasing the data throughput between different levels of memory, such as from GPU memory to L2 cache, and ensuring that data fits into on-chip memory like L1 cache or SRAM. This optimization allows FlashAttention to unlock the full potential of the GPU’s parallel processing power. By fusing kernels, FlashAttention can achieve faster execution and improved memory efficiency, making it a crucial optimization technique.

In Part 2 of this series, we will delve deeper into the FlashAttention algorithm and how it leverages GPU architecture to enhance the performance of language models.

References

--

--

Written by Sachin Kalsi

Problem Solver | NLP

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

GPU Architecture and Programming — An Introduction

Najeeb Khan

Follow

--

Listen

Share

Explore Kernel Grids, Blocks, Warps, and Threads to Accelerate Your Code

A general purpose graphics processing unit (GPU) provide parallel cores designed to process data simultaneously. Similar to the single instruction multiple data (SIMD) in CPUs, GPUs use several threads to execute instructions in parallel, the paradigm is known as single-instruction multiple threads (SIMT). Modern GPUs can theoretically run one instruction across thousands of data points in a single cycle. However, practical applications often require data exchange between operations, which can consume hundreds of clock cycles. To address this, GPUs use a hierarchical structure to manage communication latency.

Architecture

In this section, we explore the architecture of a typical GPU. While each GPU generation introduces unique optimizations, we focus on the core concepts that are common across most GPUs.

Warps

At the core level, each thread operates on individual scalar values with private registers. While running thousands of threads simultaneously is impractical, individual threads are not efficient on their own either. Instead, threads are organized into small groups called warps or wavefronts, typically consisting of 32 threads. Each warp executes a single instruction across 32 data points. For example, in matrix multiplication, a warp might process a row and column from two matrices, performing multiplication and accumulation to generate results as shown in the figure below.

Thread Blocks

When operations exceed the warp size of 32 threads, GPUs use tiling to manage larger dimensions. This involves dividing the input into chunks or tiles that fit the warp size, processing these chunks, and then combining the results from all warps. To accumulate partial results, a placeholder is needed, which is where thread blocks come in. A thread block groups multiple warps, allowing them to share memory and synchronize their execution, as illustrated in the figure below.

Grid

The hierarchy from warps to blocks is repeated one more level: if the matrix is larger than what a single thread block can handle, we use a grid of thread blocks that share global memory. The grid enables the GPU to process large datasets by distributing the workload across multiple thread blocks.

All GPU programs, known as kernels, are executed within this grid structure. When you launch a kernel, you specify both the grid size (the number of thread blocks) and the block size (the number of threads per block). This hierarchical approach ensures efficient computation and data management, allowing the GPU to handle extensive and complex tasks effectively.

Memory Hierarchy

Following the structure of the computations, memory is organized into a hierarchy starting from the small and fast registers with ultra low latency and a few kilobytes in size. Registers are private to threads. Next, warps within a thread block share state using shared memory comprising several hundred kilobytes. Finally, global memory is accessible across the device and provides large capacity on the order of tens of gigabytes with high throughput approaching a terabyte per second. Global memory has higher latency and thus caching is used to reduce latency. The figure below shows the relative scope of each memory type.

Programming

Programming GPUs are supported by dedicated software libraries in C/C++ depending on the make of the GPU: NVIDIA GPUs can be programmed using Compute Unified Device Architecture (CUDA) interface whereas AMD GPUs offer a similar SDK known as HIP.

In this section we will briefly show how to run a hello world program on multiple threads using CUDA and how to multiply two matrices.

Hello World!

The entry point of a GPU program is called a kernel. The global ID of a thread can be calculated using three compiler intrinsics — blockIdx, blockDim, and threadIdx, representing the id of the block, the total number of threads in a block, and the thread id within the thread block, respectively. A kernel is defined by the __global__ qualifier as shown in the listing below. To launch a kernel the <<<numBlocks, blockSize>>> is used. The kernel is executed asynchronously, i.e., the host code will continue to run right after making the kernel call. To sync memory between the host and the GPU device the cudaDeviceSynchronize function is called, which blocks the execution on the host until the kernel finishes its work.

#include <cuda_runtime.h>#include <iostream>__global__ void helloFromGPU() {    printf("Hello World from Thread %d, Block %d, BlockDim %d\n",             threadIdx.x, blockIdx.x, blockDim.x);}int main() {    // Launch the kernel with 2 blocks of 4 threads each    helloFromGPU<<<2, 4>>>();    cudaDeviceSynchronize();  // Wait for the GPU to finish    return 0;}

The above code can be compiled using the NVIDIA compiler and run as follows:

> nvcc hello_gpu.cu -o hello_gpu> ./hello_gpuHello World from Thread 0, Block 0, BlockDim 4Hello World from Thread 1, Block 0, BlockDim 4Hello World from Thread 2, Block 0, BlockDim 4Hello World from Thread 3, Block 0, BlockDim 4Hello World from Thread 0, Block 1, BlockDim 4Hello World from Thread 1, Block 1, BlockDim 4Hello World from Thread 2, Block 1, BlockDim 4Hello World from Thread 3, Block 1, BlockDim 4

Matrix Multiplication

Now that we know the basic structure of a CUDA program, let’s look at a more involved example of matrix multiplication. The CUDA kernel for matrix multiplication is given in the listing below. CUDA provides block IDs and thread IDs in three dimensions. In our case, since we’re dealing with matrices, we use only two dimensions: x and y for the row and column indices.

The kernel calculates the global row and column indices of each thread by combining the block index and thread index. Each thread then performs the dot product of the corresponding row from matrix A and the column from matrix B, storing the result in matrix C. This approach ensures that each element of the output matrix is computed in parallel, leveraging the GPU’s ability to handle many threads simultaneously.

__global__ void matrixMul(const float* A, const float* B, float* C, int n){    int row = blockIdx.y * blockDim.y + threadIdx.y;    int col = blockIdx.x * blockDim.x + threadIdx.x;    if (row < n && col < n) {        float value = 0.0f;        for (int k = 0; k < n; ++k) {            value += A[row * n + k] * B[k * n + col];        }        C[row * n + col] = value;    }}

While this example provides a straightforward implementation of matrix multiplication, it is not optimized for performance. In real-world applications, achieving efficient computation requires careful consideration of memory access patterns and cache utilization. Techniques such as tiling and shared memory usage can significantly enhance performance by reducing memory access latency and improving data locality. Proper cache planning and optimization strategies are essential for scaling these algorithms to handle larger datasets and more complex computations efficiently.

--

--

Written by Najeeb Khan

Life-long Learner, PhD in ML, Applied Scientist at Amazon

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams