Sign up

Sign in

Sign up

Sign in

Flash Attention with CUDA

Damien J

Follow

--

Listen

Share

Introduction

Attention mechanisms have revolutionized the field of natural language processing and deep learning. They allow models to focus on relevant parts of input data while performing tasks like machine translation, language generation, and more. Attention mechanisms enable models to weigh different parts of input data differently, focusing on the most relevant information while performing a task. This mimics the human ability to selectively pay attention to certain aspects of our surroundings while filtering out distractions. Attention mechanisms have been instrumental in improving the performance of various AI models, particularly in sequence-to-sequence tasks.

Flash Attention, as the name suggests, brings a fast and memory-efficient solution to attention mechanisms. It addresses some of the inefficiencies present in traditional attention mechanisms, making them more suitable for large-scale tasks and complex models. It is particularly crucial for training large-scale deep learning models, such as BERT-large and GPT-2, where the attention step often becomes a bottleneck. By introducing memory-efficient and exact computations, Flash Attention accelerates the training process while maintaining high accuracy.

Note: This article focuses on Flash Attention 1. Subsequent work has introduced newer variants — such as Flash Attention 2 and even Flash Attention 3 — that build upon these techniques to further improve speed, memory efficiency, and numerical stability.

1. Flash Attention Intuition & Algorithm:

Simplified Intuition - Imagine you have a huge list of words or tokens (like all the words in a paragraph) and your goal is to figure out how each word relates to all the other words. Normally, you’d compare every word with every other word in a big grid (a big table of numbers). This works, but it can use a lot of memory — like trying to keep the entire grid in your computer’s brain at once. Flash Attention is a trick to handle that comparison in smaller pieces. Instead of storing the entire giant grid, it looks at only a small part of the information at a time, does the math, and then moves on to the next piece. By processing the data in “bite-sized chunks,” Flash Attention avoids filling up memory with a huge table. The result is that you still get the same overall understanding of how each word relates to every other word, but you use much less memory, and you often run faster too. Essentially, Flash Attention is like taking the same big job and chopping it into smaller tasks that are more manageable and efficient.

In traditional mechanisms, the quadratic memory complexity (O(N²), where N is the sequence length) significantly limits the scalability of models. Flash Attention reduces this to linear complexity (O(N)) through innovative techniques such as memory tiling, thread coarsening, and recomputation, making it a more feasible option for large-scale tasks.

Core Components of Flash Attention

Flash Attention’s effectiveness lies in its understanding of the hardware it runs on. It exploits the fact that different types of memory in GPUs have varying capacities and speeds. For instance, SRAM is faster but smaller, while HBM (high bandwidth memory, the global memory) is larger but slower. By minimising the communication between these memory types, Flash Attention significantly speeds up computations

Flash Attention Algorithm: Tiling and Recomputation

Flash Attention’s algorithm can be summarised in two main ideas: tiling and recomputation.

Tiling: During both forward and backward passes, Flash Attention divides the attention matrices into smaller blocks, optimizing memory usage and improving computation efficiency.

Recomputation:In the backward pass, Flash Attention recomputes attention matrices using stored outputs and softmax normalization statistics, eliminating the need for excessive memory storage.

Complexity and Real-World Challenges

Flash Attention’s space complexity scales linearly with the sequence length and attention head dimension. This makes it suitable for handling large-scale models and tasks.

However, implementing Flash Attention comes with challenges, particularly in writing optimized CUDA kernels. The need for lower-level language coding can hinder adoption, but projects like Triton offer potential solutions to this issue.

Impact Modern AI models, particularly those in natural language processing (NLP) and computer vision, rely heavily on attention mechanisms to handle dependencies across long sequences or image patches. Flash Attention addresses the pressing need for scalable, efficient solutions in these domains. It enables researchers to train larger models faster and with fewer resources, directly impacting the pace of advancements in AI.

Benefits of Using GPUs and CUDA

These aspects make Flash Attention a pivotal step in advancing AI workloads, unlocking new possibilities for large-scale models in research and production

2. Parallelization: Concept, Semantics, and Implementation

1. Phasing-and-Tiling

Concept/Semantics:

Implementation:

2. Thread Coarsening

Concept/Semantics:

Implementation:

3. Shared Memory

Concept/Semantics:

Implementation:

4. Memory Coalescing

Concept/Semantics:

Implementation:

3. Parallel Patterns

1. Map

2. Reduce

3. Tiling

4. Experimental Results

Here are some output samples that has smaller sizes:

One can mimic the real embeddings, query, key, and value matrices with randomly generated matrices. Also implemented is the CPU code as the benchmark and expected output and designed two kernels.

5. Source Code

All the source code can be found in the https://github.com/damienjose/cuda-flashattention Github link.

// CPU Implementation of Attentionvoid computeAttentionCPU(float* query, float* key, float* value,                          float* attentionScores, float* output) {    float* transposedKey = (float*)malloc(FEATURE_DIMENSION * NUM_SAMPLES * sizeof(float));    transposeMatrix(key, transposedKey, NUM_SAMPLES, FEATURE_DIMENSION);        float scalingFactor = 1.0f / sqrtf((float)FEATURE_DIMENSION);    // Compute attention scores    for (int i = 0; i < NUM_SAMPLES; i++) {        for (int j = 0; j < NUM_SAMPLES; j++) {            for (int k = 0; k < FEATURE_DIMENSION; k++) {                attentionScores[i * NUM_SAMPLES + j] +=                     query[i * FEATURE_DIMENSION + k] * transposedKey[k * NUM_SAMPLES + j];            }            attentionScores[i * NUM_SAMPLES + j] *= scalingFactor;        }    }    // Softmax row-wise    for (int row = 0; row < NUM_SAMPLES; row++) {        float maxScore = attentionScores[row * NUM_SAMPLES];        for (int col = 1; col < NUM_SAMPLES; col++) {            if (attentionScores[row * NUM_SAMPLES + col] > maxScore) {                maxScore = attentionScores[row * NUM_SAMPLES + col];            }        }        float sumExp = 0.0f;        for (int col = 0; col < NUM_SAMPLES; col++) {            attentionScores[row * NUM_SAMPLES + col] =                 exp(attentionScores[row * NUM_SAMPLES + col] - maxScore);            sumExp += attentionScores[row * NUM_SAMPLES + col];        }        for (int col = 0; col < NUM_SAMPLES; col++) {            attentionScores[row * NUM_SAMPLES + col] /= sumExp;        }    }    // Multiply by Value matrix    for (int i = 0; i < NUM_SAMPLES; i++) {        for (int j = 0; j < FEATURE_DIMENSION; j++) {            for (int k = 0; k < NUM_SAMPLES; k++) {                output[i * FEATURE_DIMENSION + j] +=                     attentionScores[i * NUM_SAMPLES + k] * value[k * FEATURE_DIMENSION + j];            }        }    }    free(transposedKey);}

2. Naive GPU Implementation (Global Memory)

// Kernel: QK^T__global__ void computeScoresKernel(float* queryMatrix, float* keyMatrix, float* scoreMatrix) {    int row = blockIdx.y * blockDim.y + threadIdx.y;    int col = blockIdx.x * blockDim.x + threadIdx.x;    if (row < NUM_SAMPLES && col < NUM_SAMPLES) {        float score = 0.0f;        for (int d = 0; d < FEATURE_DIMENSION; ++d) {            score += queryMatrix[row * FEATURE_DIMENSION + d] *                      keyMatrix[col * FEATURE_DIMENSION + d];        }        scoreMatrix[row * NUM_SAMPLES + col] =             score / sqrtf(static_cast<float>(FEATURE_DIMENSION));    }}// Kernel: Softmax__global__ void applySoftmaxKernel(float* scoreMatrix, float* softmaxMatrix) {    int row = blockIdx.y * blockDim.y + threadIdx.y;    if (row < NUM_SAMPLES) {        float maxScore = -1e30f;        for (int col = 0; col < NUM_SAMPLES; ++col) {            maxScore = fmaxf(maxScore, scoreMatrix[row * NUM_SAMPLES + col]);        }        float sumExp = 0.0f;        for (int col = 0; col < NUM_SAMPLES; ++col) {            softmaxMatrix[row * NUM_SAMPLES + col] =                 expf(scoreMatrix[row * NUM_SAMPLES + col] - maxScore);            sumExp += softmaxMatrix[row * NUM_SAMPLES + col];        }        for (int col = 0; col < NUM_SAMPLES; ++col) {            softmaxMatrix[row * NUM_SAMPLES + col] /= sumExp;        }    }}// Kernel: Output = Softmax(QK^T) * V__global__ void computeOutputKernel(float* softmaxMatrix, float* valueMatrix, float* outputMatrix) {    int row = blockIdx.y * blockDim.y + threadIdx.y;    int col = blockIdx.x * blockDim.x + threadIdx.x;    if (row < NUM_SAMPLES && col < FEATURE_DIMENSION) {        float result = 0.0f;        for (int k = 0; k < NUM_SAMPLES; ++k) {            result += softmaxMatrix[row * NUM_SAMPLES + k] *                       valueMatrix[k * FEATURE_DIMENSION + col];        }        outputMatrix[row * FEATURE_DIMENSION + col] = result;    }}// Complete naive GPU routinevoid computeAttentionGPUGlobal(float* queryMatrix, float* keyMatrix, float* valueMatrix,                               float* attentionMatrix, float* outputMatrix) {    // Device pointers + memory allocations + memcpys omitted for brevity...    // Launch kernels    computeScoresKernel<<<gridDim, blockDim>>>(d_queryMatrix, d_keyMatrix, d_scoreMatrix);    cudaDeviceSynchronize();        applySoftmaxKernel<<<softmaxGridDim, softmaxBlockDim>>>(d_scoreMatrix, d_softmaxMatrix);    cudaDeviceSynchronize();        computeOutputKernel<<<outputGrid, outputBlock>>>(d_softmaxMatrix, d_valueMatrix, d_outputMatrix);    cudaDeviceSynchronize();        // Copy results back + free memory ...}

3. Optimized GPU Implementation (Shared Memory)

// Kernel: QK^T using Shared Memory__global__ void shared_compute_scores(float* queryMatrix, float* keyTransposeMatrix,                                       float* attentionScores) {    __shared__ float sharedQuery[MEM_WIDTH][MEM_WIDTH];    __shared__ float sharedKeyTranspose[MEM_WIDTH][MEM_WIDTH];    int threadX = threadIdx.x;    int threadY = threadIdx.y;    int blockX = blockIdx.x;    int blockY = blockIdx.y;    int scoreColumnIndex = blockX * TILE_WIDTH + threadX;    int scoreRowIndex   = blockY * TILE_WIDTH + threadY;    float scoreValue = 0.0f;    int numPhases = (FEATURE_DIMENSION + TILE_WIDTH - 1) / TILE_WIDTH;    for (int phase = 0; phase < numPhases; phase++) {        // Load query tile        if (phase * TILE_WIDTH + threadX < FEATURE_DIMENSION &&             blockY * TILE_WIDTH + threadY < NUM_SAMPLES) {            sharedQuery[threadY][threadX] =                 queryMatrix[(blockY * TILE_WIDTH + threadY) * FEATURE_DIMENSION +                             phase * TILE_WIDTH + threadX];        } else {            sharedQuery[threadY][threadX] = 0.0f;        }        // Load keyTranspose tile        if (phase * TILE_WIDTH + threadY < FEATURE_DIMENSION &&             blockX * TILE_WIDTH + threadX < NUM_SAMPLES) {            sharedKeyTranspose[threadY][threadX] =                 keyTransposeMatrix[(phase * TILE_WIDTH + threadY) * NUM_SAMPLES +                                   (blockX * TILE_WIDTH + threadX)];        } else {            sharedKeyTranspose[threadY][threadX] = 0.0f;        }        __syncthreads();        // Dot product        if (scoreColumnIndex < NUM_SAMPLES && scoreRowIndex < NUM_SAMPLES) {            for (int i = 0; i < TILE_WIDTH; i++) {                scoreValue += sharedQuery[threadY][i] * sharedKeyTranspose[i][threadX];            }        }        __syncthreads();    }    if (scoreColumnIndex < NUM_SAMPLES && scoreRowIndex < NUM_SAMPLES) {        attentionScores[scoreRowIndex * NUM_SAMPLES + scoreColumnIndex] =             scoreValue / sqrtf(static_cast<float>(FEATURE_DIMENSION));    }}// Kernel: Softmax__global__ void shared_softmax(float* attentionScores, float* softmaxScores) {    int rowIndex = blockIdx.y * blockDim.y + threadIdx.y;    if (rowIndex < NUM_SAMPLES) {        float maxScore = -1e30f;        // Find max for numerical stability        for (int colIndex = 0; colIndex < NUM_SAMPLES; ++colIndex) {            maxScore = fmaxf(maxScore, attentionScores[rowIndex * NUM_SAMPLES + colIndex]);        }        float sumExp = 0.0f;        for (int colIndex = 0; colIndex < NUM_SAMPLES; ++colIndex) {            softmaxScores[rowIndex * NUM_SAMPLES + colIndex] =                 expf(attentionScores[rowIndex * NUM_SAMPLES + colIndex] - maxScore);            sumExp += softmaxScores[rowIndex * NUM_SAMPLES + colIndex];        }        for (int colIndex = 0; colIndex < NUM_SAMPLES; ++colIndex) {            softmaxScores[rowIndex * NUM_SAMPLES + colIndex] /= sumExp;        }    }}// Kernel: Output = SoftmaxScores * V__global__ void shared_compute_output(float* softmaxScores, float* valueMatrix,                                       float* outputMatrix) {    __shared__ float sharedSoftmaxScores[TILE_WIDTH][TILE_WIDTH];    __shared__ float sharedValueMatrix[TILE_WIDTH][TILE_WIDTH];    int threadX = threadIdx.x;    int threadY = threadIdx.y;    int blockX = blockIdx.x;    int blockY = blockIdx.y;    int outputColumnIndex = blockX * TILE_WIDTH + threadX;    int outputRowIndex    = blockY * TILE_WIDTH + threadY;    float outputValue = 0.0f;    int numPhases = (NUM_SAMPLES + TILE_WIDTH - 1) / TILE_WIDTH;    for (int phase = 0; phase < numPhases; phase++) {        if (phase * TILE_WIDTH + threadX < NUM_SAMPLES &&             blockY * TILE_WIDTH + threadY < NUM_SAMPLES) {            sharedSoftmaxScores[threadY][threadX] =                 softmaxScores[(blockY * TILE_WIDTH + threadY) * NUM_SAMPLES +                               (phase * TILE_WIDTH + threadX)];        } else {            sharedSoftmaxScores[threadY][threadX] = 0.0f;        }        if (phase * TILE_WIDTH + threadY < NUM_SAMPLES &&             blockX * TILE_WIDTH + threadX < FEATURE_DIMENSION) {            sharedValueMatrix[threadY][threadX] =                 valueMatrix[(phase * TILE_WIDTH + threadY) * FEATURE_DIMENSION +                             (blockX * TILE_WIDTH + threadX)];        } else {            sharedValueMatrix[threadY][threadX] = 0.0f;        }        __syncthreads();        // Dot product        if (outputColumnIndex < FEATURE_DIMENSION && outputRowIndex < NUM_SAMPLES) {            for (int i = 0; i < TILE_WIDTH; i++) {                outputValue += sharedSoftmaxScores[threadY][i] * sharedValueMatrix[i][threadX];            }        }        __syncthreads();    }    if (outputColumnIndex < FEATURE_DIMENSION && outputRowIndex < NUM_SAMPLES) {        outputMatrix[outputRowIndex * FEATURE_DIMENSION + outputColumnIndex] = outputValue;    }}// Complete shared memory GPU routinevoid computeAttentionGPUShared(float* queryMatrix, float* keyTransposeMatrix, float* valueMatrix,                               float* attentionScores, float* outputMatrix) {    // Device pointers + memory allocations + memcpys omitted for brevity...        // Launch kernels    shared_compute_scores<<<gridDimension, blockDimension>>>(        deviceQuery, deviceKeyTranspose, deviceAttentionScores);    cudaDeviceSynchronize();    shared_softmax<<<softmaxGridDimension, softmaxBlockDimension>>>(        deviceAttentionScores, deviceSoftmaxScores);    cudaDeviceSynchronize();    shared_compute_output<<<outputGrid, outputBlock>>>(        deviceSoftmaxScores, deviceValue, deviceOutput);    cudaDeviceSynchronize();    // Copy results back + free memory ...}

Key Differences Between Implementations

CPU:

Global Memory CUDA:

Shared Memory CUDA:

6. Comparative Analysis

Below are the CPU, GPU Global Memory and GPU Shared Memory Latency measurements for different TILE_WIDTH, BLOCK_SIZE, NUM_SAMPLES and FEATURE_DIMENSION as measured on a NVIDIA GeForce RTX 4050 Laptop GPU

a. Performance Breakdown:

CPU Implementation:

Naive GPU Implementation (Global Memory):

Optimized GPU Implementation (Shared Memory):

b. Computational Cost/Complexity

CPU Implementation:

Naive GPU Implementation (Global Memory):

Optimized GPU Implementation (Shared Memory):

c. Work Efficiency

CPU Implementation:

Naive GPU Implementation (Global Memory):

Optimized GPU Implementation (Shared Memory):

7. Conclusion

Key Optimizations:

Performance Insights:

Scalability:

Limitations:

Experimental Outcomes:

Latency Measurements (1024 x 1024 matrices):

CPU: 9,427 ms

Naive GPU (Global Memory): 38.78 ms

Optimized GPU (Shared Memory): 14.12 ms

For large matrix sizes (e.g., 1024x1024), the optimized shared memory implementation is ~3x faster than the naive GPU implementation 600x faster than the CPU implementation and achieves close-to-ideal theoretical occupancy with minimal control divergence.

8. References

9. Appendix

Change the SM and compute in Visual Studio Project Properties, before you run the code for both Debug and Release builds

Profiling and Performance Analysis

Two main GPU implementations were profiled — Naive (Global Memory) vs. Optimized (Shared Memory) — on an NVIDIA GeForce RTX 4050 Laptop GPU with N = D = 1024 and BLOCK_SIZE = TILE_WIDTH = MEM_WIDTH = 32.

Naive (Global Memory) Highlights

Compute Scores Kernel:

Softmax Kernel:

Output Kernel:

Optimized (Shared Memory) Highlights

Shared Compute Scores:

Shared Softmax:

Shared Compute Output:

--

--

Written by Damien J

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 Parallel Programming with CUDA and C++

Avin

Follow

--

Listen

Share

Parallel programming on GPUs is one of the best ways to speed up processing of compute intensive workloads. Programming for CUDA enabled GPUs can be as complex or simple as you want it to be. Some developers elect to take the simpler route of using libraries such as Thrust and OpenCL. While they facilitate faster serial to parallel conversion times, they do not reveal the GPU architecture that facilitate speedups in the magnitude of 10x to 100x in comparison to CPUs. In this article, I will explore some of the key concepts that facilitate GPU acceleration and walk through parallelizing a simple algorithm.

CUDA Architecture

A modern CPU consists of several cores that are each capable of executing a single thread. Each of these threads can contain their own sequence of unique instructions. Each core executes a single thread

CUDA Architecture utilizes a different approach where a collection of “streaming multiprocessors” (SM) execute the same set of instructions, including branch conditions on multiple threads on different regions of data. Nvidia has coined the term SIMT, single instruction, multiple threads, to describe this architecture.

Each corresponding pair of arrows entering and leaving the streaming multiprocessor represent a single thread. 21 threads are working in parallel in this theoretical GPU. Therefore it is able to map 21 elements in a single step. A CPU calculating in serial will require 21 steps one after the other to map this array.

Our programs do not always map one to one to the hardware implementation of the above architecture, for example, we may have more than 21 elements in the array or this program may need to run on a lesser GPU with only 10 parallel threads. The Nvidia CUDA runtime API provides a layer of abstraction between thread execution in hardware and thread representation in software.

In the simplest of terms, developers can group threads into blocks and a collection of blocks form a grid. Maximum number of threads per block is 1024. The number of blocks per grid can go up to 2³¹-1, i.e. 2,147,483,647. i.e. the maximum value of a signed int.

Let’s apply this in practice to an array of one million elements that are to be processed in parallel. If we are to launch 1024 threads per block, this would require us to launch at least 977 blocks to process the entire array in parallel. (1,000,000 / 1024 = 976.6). Each of these threads need to know which element of the array to process. Since we have one thread for each element of the array, we can use an id based on the thread count. The CUDA runtime exposes threadIdx.x to reveal the thread Id within a block and blockIdx.x to reveal the block Id within the grid. It also exposes blockDim.x to reveal the dimensions of the block. Putting it all together we can calculate a global id for each thread within the entire grid as follows

int globalIdx = blockIdx.x * blockDim.x + threadIdx.x

Knowing the global thread id, we also need to ensure that threads do not attempt illegal memory access, as there are more threads than there are elements in an array due to rounding up. Combining these concepts, a single unit of computation, i.e. a CUDA kernel, looks as follows in pseudo C++ code.

kernel function parallel_scalar_multiply  int globalIdx = blockIdx.x * blockDim.x + threadIdx.x  if globalIdx < N then // where N is number of elements in array    array[globalIdx] = array[globalIdx] * 10  end ifend kernel

As you might have noticed some threads in this case will not do any work as they are out of bounds of the array, This is natural in massively parallel architectures such as CUDA. To launch this kernel to process all elements in an array the kernel invocation code has to specify the launch configuration that was discussed above. In pseudo code:

launch parallel_scalar_multiply with 1024 threads in 977 blocks

Before we launch

Programming on CUDA requires the CUDA Toolkit and a CUDA GPU. While the CUDA Toolkit extends the C language, C++ provides a richer syntax on top of C and will be the language of choice. The CUDA compiler and profiler is installed with the toolkit, the debugger and visual profiler may require separate installation.

Compiler — nvcc automatically adds the CUDA source headers and links the CUDA runtime libraries.

Profiler — nvprof command line profiler that profiles kernel execution times and runtime API calls

Debugger — cuda-gdb compiler with nvcc -g to enable verbose debugging output

Visual Profiler — nvvp the swiss army knife of CUDA application optimization. Includes detailed analysis of core utilization, register usage, global memory access and timing profiles.

Our first CUDA kernel launch

As with any program, execution has to start at the CPU. We will create test data, process it on the CPU (host), time it and transfer the data to the GPU (device), process it again and copy the results back to the host machine. The code below is compiled with nvcc -std=c++11 cuda_multiply.cu and profiled with nvprof ./a.out

We first generate 60 million integers as the test data:

int main() {  int count = 60000000;    // 60 million elements  int* test_data = new int[count];for(int i = 0; i < count; i++)    test_data[i] = i;}

The CPU multiplication, timed with C++ chrono, #include <chrono>

auto t1 = std::chrono::high_resolution_clock::now();for(int i = 0; i < count; i++)  test_data[i] = test_data[i] * 5;auto t2 = std::chrono::high_resolution_clock::now();

In order to perform the calculation on the GPU, the data has to be copied first, the CUDA run-time exposes cudaMalloc to allocate memory on the GPU and cudaMemcpy to transfer the data.

int* d_test_data;cudaMalloc(&d_test_data, count * sizeof(int));cudaMemcpy(d_test_data, test_data, count * sizeof(int), cudaMemcpyHostToDevice);

The kernel launch configuration is specified using the following syntax

<<<num of blocks, num of threads per block>>>

It is similar to an ordinary method call except for the above snippet follows immediately after the method name before specifying the parameters:

int block_count = ceil((double)count / 1024);_cuda_parallel_multiplication<<<block_count, 1024>>>(count, d_test_data, 5);

The kernel code is as follows

__global__ void _cuda_parallel_multiplication(int count, int* test_data, int magnitude) {  int globalIdx = blockIdx.x * blockDim.x + threadIdx.x;  if (globalIdx < count)    test_data[globalIdx] = test_data[globalIdx] * magnitude;                                               }

The __global__ specifier informs the compiler that the function is a CUDA kernel and has to be compiled as such. The if statement checks whether the current global thread id is within bounds of the array before accessing it and performing the multiplication. This set of instructions is executed identically across all 60 million threads launched.

The following code copies the processed data back to the host machine from the CUDA device.

cudaDeviceSynchronize();cudaMemcpy(test_data, d_test_data, count * sizeof(int), cudaMemcpyDeviceToHost);

CUDA kernel launches are asynchronous in relation to the host computer; cudaDeviceSynchronize() blocks the CPU code till all GPU functions complete. cudaMemcpy synchronizes by default as well, but I have been explicit here to demonstrate the need to synchronize before any calculations can be done with data from the GPU. The processed data is copied back to the initial test data set.

The program finally samples a portion of the results, prints the CPU execution time and exits.

for(int i = 0; i < 10; i++)  std::cout << i << ": " << test_data[i] << std::endl;std::cout  << "CPU time: "   << std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1).count()   << "ms"    << std::endl;

At this point it has gone through two multiplication rounds. i.e. into 5 on CPU and into 5 on the GPU and the sample output will reflect this as: i x 5 x 5

0: 01: 252: 503: 754: 1005: 1256: 1507: 1758: 2009: 225CPU time: 97ms

Running the application with nvprof ./a.out results in:

Time(%) Time       Name46.85%  26.438ms   [CUDA memcpy DtoH]42.90%  24.207ms   [CUDA memcpy HtoD]10.25%  5.7829ms   _cuda_parallel_multiplication(int, int*, int)

The 97ms computation on the CPU takes a mere 5.8ms on the GPU. It is out shadowed by memcpy times here, but only because this kernel is trivial and performs no meaningful work. Memory latency can be hidden by the use of streams to overlap memcpy and kernel execution, which I will explore in later posts.

A simple optimization

You must have wondered why we launched 60 million threads for each element in the array when physical GPUs do not possess the capacity to process such monstrous amounts of threads in parallel. It was merely a technique to easily map the sequential workload to a parallel one. However, this introduces the overhead of switching through 60million threads. While the optimal thread count varies by workload, generally the block count should be a multiple of the number of SMs on the GPU, while the threads per block should be a maximum that does not overload register usage per block. With a trivial multiplication such as above, a thread count of 1024 does not overflow the registers, leaving us to pick a block count. On my GTX1050, a single streaming multiprocessor(SM) processes two blocks simultaneously with a total of 5 SMs. (A Telsa P100 has a whopping 56 SMs !!!) Therefore, to achieve max occupancy, a total of 10(5 x 2) blocks each with 1024 threads should be launched. But how do 10240 threads process 60 million elements? Loop.

__global__ void _cuda_parallel_multiplication(int count, int* test_data, int magnitude) {  int globalIdx = blockIdx.x * blockDim.x + threadIdx.x;    while (globalIdx < count) {    test_data[globalIdx] = test_data[globalIdx] * magnitude;        globalIdx += blockDim.x * gridDim.x;  }}

The if statement from the previous kernel has been turned into a loop. Each iteration increases the id by the block dimension multiplied by the grid dimension. i.e. the entire grid of threads is “striding” to the next group of elements in the array. Switching to a grid stride resulted in a 0.6ms decrease in execution time on my GTX1050.

Another optimization

__global__ void _cuda_parallel_multiplication(int count, int* test_data, int magnitude) {int globalIdx = blockIdx.x * blockDim.x + threadIdx.x;    while (globalIdx < count) {    test_data[globalIdx] = test_data[globalIdx] * magnitude;        globalIdx += blockDim.x * gridDim.x;    __syncthreads();  }}

__syncthreads() is a block level function that synchronizes execution of all threads across the block. Using it here, ensures that all threads reach this point of execution before the next iteration begins. In the context of this kernel, it ensures that global memory loads happen in unison for the entire block. This along with the grid striding ensures maximal use of caching as entire cache lines are loaded into memory, shared across all threads in the block, read and written to before the next segment of global memory is fetched. On my GTX1050 GPU a difference in execution time of 0.2ms was observed with and without __syncthreads()

Results

CPU: Core i3–7100 @ 3.9GHZ

Average Time: 95ms

GPU: Nvidia GTX 1050, 5 SMs

Unoptimized Average Time: 5.8ms Speedup: x16

Optimized Average Time: 5.0ms Speedup: x19

GPU: Nvidia Tesla P100, 56 SMs

Optimized Average Time: 920us=0.92ms Speedup: x103

Conclusion

The results above speak to the efficacy of CUDA as a platform for accelerated computing. A simple multiplication was chosen for parallelization to allow the reader to focus on the CUDA architecture as opposed to understanding a complex calculation. The speed up seen in this trivial calculation is not an exaggeration of speed up seen in more complex useful calculations. However, overall speedup in a real application might not be as high due to memory limitations and copy latency.

Credits

Mark Harris — https://devblogs.nvidia.com/author/mharris/

Wayne Kelly — http://staff.qut.edu.au/staff/kellyw/

--

--

Written by Avin

Computer Science. Philosophy. Object Oriented.

No responses yet

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Sign up

Sign in

Sign up

Sign in

Optimizing any TensorFlow model using TensorFlow Transform Tools and using TensorRT

Model Optimization and reducing precision from FP32 to FP 16 for speedup and reducing graph size.

Alex Punnen

Follow

Better Software

--

3

Listen

Share

What is this article about? Whom is this meant for?

More than an article, this is basically how to, on optimizing a Tensorflow model, using TF Graph transformation tools and NVIDIA Tensor RT. This is a bit of a Heavy Reading and meant for Data Science Engineers than Data Scientist. Also, more than a how-to, it also illustrates the bugs and unexpected results I have encountered in doing this.

Also, these results may be unexpected to me in the role of a DS Engineer but may have a rational explanation if I understand or go deeper into the abstractions of the model or the TF framework that implements this model. I don’t have the skill yet and leave it to the reader to point out the mistakes I have made if any. I have raised these two bugs, one with TF Graph Transform and one with NVIDIA TensorRT transform tool, and it gets answered, then I will update this post

https://github.com/tensorflow/tensorflow/issues/28276

https://devtalk.nvidia.com/default/topic/1048485/tensorrt/no-speed-up-with-tensorrt-fp16-or-int8-on-nvidia-v100/post/5320989/#5320989

What was my expectation — More CUDA cores, TensorCores == ( I thought) Faster inference.

From NVIDIA TensorCore product description, I got the idea/ hope that if I can convert the FP 32 based Object detection models that we use in production to FP 16 or INT8 models and weights, then I will be able to run twice or four times as fast inference speeds; as advertised.

We had one of the best GPU’s NVIDIA V100 32 GB for our experiments which supported all these modes- Tensor Cores, not to mention a set of other GPU’s server-grade as well as gaming grade, P40, GTX 1080, 1070 etc.

Before we started with these expensive GPU’s we used to run the model in GTX 1080 HP desktops. The first expectation was that the higher number of CUDA cores (~5k) in V100 will make our models run faster. I read blogs like below and knew that 1080 is pretty good and that NVIDIA prices and markets the server side GPU’s higher. ( In 1080 they throttle the FP16 which I guess is removed in 2080).

Deep Learning GPU Benchmarks - Tesla V100 vs RTX 2080 Ti vs GTX 1080 Ti vs Titan V

At Lambda, we're often asked "what's the best GPU for deep learning?" In this post and accompanying white paper, we…

lambdalabs.com

Here is an excerpt from the above

“2080 Ti vs V100 — is the 2080 Ti really that fast?If you absolutely need 32 GB of memory because your model size won’t fit into 11 GB of memory with a batch size of 1. If you are creating your own model architecture and it simply can’t fit even when you bring the batch size lower, the V100 could make sense. However, this is a pretty rare edge case. Fewer than 5% of our customers are using custom models. Most use something like ResNet, VGG, Inception, SSD, or Yolo.

So. You’re still wondering. Why would anybody buy the V100? It comes down to marketing.”

Also, I was not expecting a drastic speed up anyway with more CUDA cores as even with HD image frames, we found that the GPU utilisation could not touch 100 per cent meaning that the processing alone was not the bottleneck. Again things are grey here. There is a lot more scope of Engineering optimisations here.

Why we choose V100 was not because of ‘marketing’; it was the only GPU with that much memory 32 GB, which would enable us to batch more image frames in parallel, basically do real-time analytics of more HD video cameras on a single edge.

Reality regarding more CUDA Cores

The truth was that other than the advantage of processing more frames in parallel due to the higher memory, there was no speedup from 1080 GPU (~2.5k CUDA core). We have also tested this in Jetson TX2 which has much fewer (~256) CUDA cores and one older gaming GPU where it was very slow. So higher CUDA cores help to run faster but beyond some threshold, there is not much difference. Maybe this fact is known already and that’s why the newer models from NVIDIA like T4 have around 2.5k CUDA cores, and these are available in GCP and other cloud providers for inference at much cheaper rates. The V100 seems to be used only for training models. It is not practical at least now to train models in lower precision. Maths is easier — gradient descent, back propagation with higher precision.

You can gain more insights from the post by Tom Dettmers regarding GPUs https://timdettmers.com/2019/04/03/which-gpu-for-deep-learning/, but do be wary of things like a number of frames you need to process in parallel etc than just raw inference speeds.

Reality regarding TensorCores, half precision/lower precision FP16, INT8

The field of Data Science Engineering is still nascent in that there is no clear distinction where Data Science ends and Engineering takes over. Frameworks like Tensorflow Serving and tools helps the Dev or DS Operations team to work on the model, develop generic clients and build on top useful applications on the model. But they treat the model without knowing too much in-depth. So when they take a model and do an Optimisation and get an error like below, they don’t know what to do, than get a real inferiority complex and make a mental note on understanding things better

details = "input_max_range must be larger than input_min_range.	 [[{{node Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/ClipToWindow_87/Area/mul_eightbit/Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/ClipToWindow_87/Area/sub_1/quantize}}]]	 [[{{node Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/zeros_like_83}}]]"

With that prelude, and based on the optimisations that worked (partially) on converting an Object detection model — Single Shot Detector from TF model zoo, here are the results I got running in V100. Basically, not much speed up, as many of the layers did not get converted**. I had initially done the experiment on a Keras converted model and got similar results, but then I thought that if I used a TF written model, it may be better, and hence the experiments on SSD model.

**tensorflow/contrib/tensorrt/segment/segment.cc:443] There are 3962 ops of 51 different types in the graph that are not converted to TensorRT: TopKV2, NonMaxSuppressionV2, TensorArrayWriteV3, Const, Squeeze, ResizeBilinear, Maximum, Where, Add, Placeholder, Switch, TensorArrayGatherV3, NextIteration, Greater, TensorArraySizeV3, NoOp, TensorArrayV3, LoopCond, Less, StridedSlice, TensorArrayScatterV3, ExpandDims, Exit, Cast, Identity, Shape, RealDiv, TensorArrayReadV3, Reshape, Merge, Enter, Range, Conv2D, Mul, Equal, Sub, Minimum, Tile, Pack, Split, ZerosLike, ConcatV2, Size, Unpack, Assert, DataFormatVecPermute, Transpose, Gather, Exp, Slice, Fill, (For more information see https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html#support-ops).

The rest of the article is more details on how I did this and which you can also follow step by step.

The post that has really helped me was these from Google team -

[1] https://medium.com/google-cloud/optimizing-tensorflow-models-for-serving-959080e9ddbf

[2] https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/graph_transforms#optimizing-for-deployment

I am writing this post as a more detailed explanation to [1], as some parts were not clear when I started following the steps.

My Colab/Jupyter notebook for Optimization is given here; you can skip the article and follow the Notebook also as I have documented it in the notebook. The TF Serving and the Client parts are however in this article.

https://colab.research.google.com/drive/1wQpWoc40kf__WSjfTqDaReMx6fFjUn48

If you have a frozen TF graph you can use the following methods to optimize it before using it for inferences.

There are two types of optimization. One to make it faster or smaller in size to run inferences. And the other to change the weights from higher precision to lower precision. Usually from FP32 to FP16 or INT8. For the latter, the GPU should have the ability to run mixed precision operations (Tensor Cores). Usually, NVIDIA’s desktop or laptop class GTX 1080 or similar are restricted from running lower precision operations. NVIDIA’s server-class GPUs support this. Especially the newer GPUs V100, T4, etc. Not all server GPU’s support it.

The GPU I l use is NVIDIA V100 32 GB GPU which has support for mixed precision operations. Also, you need to run the optimization in the GPU that you are optimizing for. Especially if you are using TensorRT.

Step 0. The model, and the Docker Containers

The first thing that has to be done is to convert the TensorFlow graph to a Frozen Graph. If the graph is Kearns based it is the HD5 format and has to be converted to the TF model and then to the frozen graph. A frozen graph has the value of variables embedded in the graph itself. It is a GrpahDef/protocol buffer (pb) format like a Saved Model only it cannot be retrained.

What is difference frozen_inference_graph.pb and saved_model.pb?

frozen_inference_graph.pb, is a frozen graph that cannot be trained anymore, it defines the graphdef and is actually a…

stackoverflow.com

The model that we are using is the SSD model ssd_resnet_50_fpn_coco form TF model zoo -https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md

Docker container used for the optimization is tensorflow/tensorflow:1.13.0rc1-gpu-jupyter

docker run --entrypoint=/bin/bash --runtime=nvidia  -it --rm -p 8900:8500 -p 8901:8501 -v /usr/alex/:/coding --net=host tensorflow/tensorflow:1.13.0rc1-gpu-jupyteronce insidecd /codingjupyter notebook --allow-root &

Note- I changed the entry point to something more convenient to me than default tf-notebook I believe.

After optimizing, to run inferences I am using the same docker image after installing on that TF serving API’s, as well as headless opencv-python version. This is because we will be converting the optimized model to a TF serving compatible model for inference.

docker run --entrypoint=/bin/bash --env http_proxy=<my proxy> --env https_proxy=<my proxy>  --runtime=nvidia  -it --rm -p 8900:8500 -p 8901:8501 -v /usr/alex/:/coding --net=host tensorflow/tensorflow:1.13.0rc1-gpu-jupyterpip install tensorflow-serving-apipip install opencv-python==3.3.0.9cd codingpython ssd_client_1.py -num_tests=1 -server=127.0.0.1:8500 -batch_size=1 -img_path='../examples/google1.jpg/'

Step 1. Get the output node names in the Tensorflow Graph

Why is this important? We need to find the output node names of the frozen graph as it is needed to optimize the graph. Note Tensorflow version that is used in TF 1.13

# To Freeze the Saved Model# We need to freeze the model to do further optimisation on itfrom tensorflow.python.saved_model import tag_constantsfrom tensorflow.python.tools import freeze_graphfrom tensorflow.python import opsfrom tensorflow.tools.graph_transforms import TransformGraphdef freeze_model(saved_model_dir, output_node_names, output_filename):  output_graph_filename = os.path.join(saved_model_dir, output_filename)  initializer_nodes = ''  freeze_graph.freeze_graph(      input_saved_model_dir=saved_model_dir,      output_graph=output_graph_filename,      saved_model_tags = tag_constants.SERVING,      output_node_names=output_node_names,      initializer_nodes=initializer_nodes,      input_graph=None,      input_saver=False,      input_binary=False,      input_checkpoint=None,      restore_op_name=None,      filename_tensor_name=None,      clear_devices=True,      input_meta_graph=False,  )

For this, we can plot the model in TF Board and see the output nodes, or print the nodes and grep on some keywords.

# Source https://medium.com/google-cloud/optimizing-tensorflow-models-for-serving-959080e9ddbfdef get_graph_def_from_file(graph_filepath):  tf.reset_default_graph()  with ops.Graph().as_default():    with tf.gfile.GFile(graph_filepath, 'rb') as f:      graph_def = tf.GraphDef()      graph_def.ParseFromString(f.read())      return graph_def

let us use the above helper to print the input and output nodes, input nodes via the for loop -

graph_def =get_graph_def_from_file('/coding/ssd_inception_v2_coco_2018_01_28/frozen_inference_graph.pb')for node in graph_def.node:    if node.op=='Placeholder':        print node # this will be the input node

and output nodes by plotting it in a format readable by Tensor Board.

with tf.Session(graph=tf.Graph()) as session:    mygraph = tf.import_graph_def(graph_def, name='')    writer = tf.summary.FileWriter(logdir='/coding/log_tb/1', graph=session.graph)    writer.flush()

Let us invoke Tensor board.

#ssh -L 6006:127.0.0.1:6006 root@<remoteip> # for tensor board - in your local machine type 127.0.0.1tensorboard --logdir '/coding/log_tb/1'

From this, I could make out the output nodes. Note that if you are building the graph yourself you don’t need to do this circus. Since I am using a model that is opensourced and with less documentation I am using this. Sometimes for auto converted/TF imported graphs, the names will be pretty long. You can then print the nodes in a for a loop as I did for Placeholder and from the output, shape make out ( for detections class, score, rectangle coordinates)

# These are the output names. Add a index usually 0 for graph nodes. # You can print the node details by nodenamesoutput_node_names = ['detection_boxes','detection_scores','detection_classes','num_detections']outputs =  ['detection_boxes:0','detection_scores:0','detection_classes:0','num_detections:0']

Step 3 Optimise using TF Graph Transform Tools

The snippet below illustrates how you can optimize a graph after reading it from disk.

# Source https://medium.com/google-cloud/optimizing-tensorflow-models-for-serving-959080e9ddbf#https://gist.github.com/lukmanr# Optimizing the graph via TensorFlow libraryfrom tensorflow.tools.graph_transforms import TransformGraphdef optimize_graph(model_dir, graph_filename, transforms, output_names, outname='optimized_model.pb'):  input_names = ['input_image',] # change this as per how you have saved the model  graph_def = get_graph_def_from_file(os.path.join(model_dir, graph_filename))  optimized_graph_def = TransformGraph(      graph_def,      input_names,        output_names,      transforms)  tf.train.write_graph(optimized_graph_def,                      logdir=model_dir,                      as_text=False,                      name=outname)  print('Graph optimized!')

Let us use the above helper to optimize the graph first quantize_weights

# Optimization without Qunatization - Reduce the size of the model# speed may actually be slower# see https://medium.com/google-cloud/optimizing-tensorflow-models-for-serving-959080e9ddbftransforms = ['remove_nodes(op=Identity)', \ 'merge_duplicate_nodes', \ 'strip_unused_nodes', 'fold_constants(ignore_errors=true)', 'fold_batch_norms', 'quantize_weights'] #this reduces the size, but there is no speed up , actaully slows down, see belowoptimize_graph('/coding/ssd_inception_v2_coco_2018_01_28', 'frozen_inference_graph.pb' ,               transforms, output_node_names,outname='optimized_model_small.pb')

Let’s then convert the optimized model to TF serving compatible format.

#lets convert this to a s TF Serving compatible mode;convert_graph_def_to_saved_model('/coding/ssd_inception_v2_coco_2018_01_28/2',                                 '/coding/ssd_inception_v2_coco_2018_01_28/optimized_model_small.pb',outputs)

The helper that does this is given below

# Source https://medium.com/google-cloud/optimizing-tensorflow-models-for-serving-959080e9ddbf#https://gist.github.com/lukmanrdef convert_graph_def_to_saved_model(export_dir, graph_filepath,outputs):graph_def = get_graph_def_from_file(graph_filepath)  with tf.Session(graph=tf.Graph()) as session:    tf.import_graph_def(graph_def, name='')    tf.saved_model.simple_save(        session,        export_dir,# change input_image to node.name if you know the name        inputs={'input_image': session.graph.get_tensor_by_name('{}:0'.format(node.name))            for node in graph_def.node if node.op=='Placeholder'},        outputs={t:session.graph.get_tensor_by_name(t) for t in outputs}                                            )    print('Optimized graph converted to SavedModel!')

And then ‘quantize_weights’ and ‘quantize_nodes’.

This should really covert also the calculation to lower precision - but does not work as of now.

"This process converts all the operations in the graph that have eight-bit quantized equivalents and leaves the rest in floating point. Only a subset of ops are supported and on many platforms, the quantized code may actually be slower than the float equivalents, but this is a way of increasing performance substantially when all the circumstances are right.”https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/graph_transforms#optimizing-for-deployment

transforms = ['add_default_attributes', \  'strip_unused_nodes', \  'remove_nodes(op=Identity, op=CheckNumerics)',\  'fold_constants(ignore_errors=true)',  'fold_batch_norms',  'fold_old_batch_norms',  'quantize_weights',  'quantize_nodes',  'strip_unused_nodes',  'sort_by_execution_order']optimize_graph('/coding/ssd_inception_v2_coco_2018_01_28', 'frozen_inference_graph.pb' ,               transforms, output_node_names,outname='optimized_model_weight_quant.pb')

However this does not work in the sense, inference using this optimized model gives the error. I had tried with a Keras model earlier and got another error message. This seems to be a bug as now this model is a pure Tensorflow model and I have not changed anything here

(‘Got an error’, <_Rendezvous of RPC that terminated with: status = StatusCode.INVALID_ARGUMENT details = “input_max_range must be larger than input_min_range. [[{{node Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/ClipToWindow_87/Area/mul_eightbit/Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/ClipToWindow_87/Area/sub_1/quantize}}]] [[{{node Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/zeros_like_83}}]]” debug_error_string = “{“created”:”@1555723203.356344655",”description”:”Error received from peer”,”file”:”src/core/lib/surface/call.cc”,”file_line”:1036,”grpc_message”:”input_max_range must be larger than input_min_range.\n\t [[{{node Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/ClipToWindow_87/Area/mul_eightbit/Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/ClipToWindow_87/Area/sub_1/quantize}}]]\n\t [[{{node Postprocessor/BatchMultiClassNonMaxSuppression/map/while/MultiClassNonMaxSuppression/zeros_like_83}}]]”,”grpc_status”:3}”>)Response Received Exiting

Step 4 Optimise using NVIDIA TenosrRT

Base reference for this is these two posts

https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html

https://developers.googleblog.com/2018/03/tensorrt-integration-with-tensorflow.html

Inference with TF-TRT `SavedModel` workflow: we are using the TF Serving model.

import tensorflow.contrib.tensorrt as trttf.reset_default_graph()graph = tf.Graph()sess = tf.Session()# Create a TensorRT inference graph from a SavedModel:with graph.as_default():    with tf.Session() as sess:        trt_graph = trt.create_inference_graph(            input_graph_def=None,            outputs=outputs,            input_saved_model_dir='/coding/ssd_inception_v2_coco_2018_01_28/01',            input_saved_model_tags=['serve'],            max_batch_size=1,            max_workspace_size_bytes=7000000000,            precision_mode='FP16')             #precision_mode='FP32')            #precision_mode='INT8')        output_node=tf.import_graph_def(trt_graph, return_elements=outputs)        #sess.run(output_node)        tf.saved_model.simple_save(sess,        "/coding/ssd_inception_v2_coco_2018_01_28/4",        inputs={'input_image': graph.get_tensor_by_name('{}:0'.format(node.name))            for node in graph.as_graph_def().node if node.op=='Placeholder'},        outputs={t:graph.get_tensor_by_name('import/'+t) for t in outputs}        )

Inference with TF-TRT `Frozen` graph workflow:

Reference https://medium.com/tensorflow/speed-up-tensorflow-inference-on-gpus-with-tensorrt-13b49f3db3fa

#Lets load a frozen model and reset the graph and usegdef =get_graph_def_from_file(‘/coding/ssd_inception_v2_coco_2018_01_28/frozen_inference_graph.pb’)tf.reset_default_graph()graph = tf.Graph()sess = tf.Session()# Create a TensorRT inference graph from a SavedModel:with graph.as_default(): with tf.Session() as sess: trt_graph = trt.create_inference_graph( input_graph_def=gdef, outputs=outputs, max_batch_size=8, max_workspace_size_bytes=7000000000, is_dynamic_op=True, #precision_mode=’FP16')  #precision_mode=’FP32') precision_mode=’INT8')  output_node=tf.import_graph_def(trt_graph, return_elements=outputs) #sess.run(output_node) tf.saved_model.simple_save(sess, “/coding/ssd_inception_v2_coco_2018_01_28/5”, inputs={‘input_image’: graph.get_tensor_by_name(‘{}:0’.format(node.name)) for node in graph.as_graph_def().node if node.op==’Placeholder’}, outputs={t:graph.get_tensor_by_name(‘import/’+t) for t in outputs}  )

Step 5: Pause and Check the models

The outputs of the various models are given below. You can see that the model size reduces after optimizations.

Original model ('/coding/ssd_inception_v2_coco_2018_01_28/frozen_inference_graph.pb', '') Model size: 99591.409 KB Variables size: 0.0 KB Total Size: 99591.409 KB ---------Tensorflow Transform Optimised model Weights Quantised ('/coding/ssd_inception_v2_coco_2018_01_28/2/saved_model.pb', '') Model size: 26193.27 KB Variables size: 0.0 KB Total Size: 26193.27 KB ---------Tensorflow Transform Optimised model Weights and Nodes Quantised ('/coding/ssd_inception_v2_coco_2018_01_28/3/saved_model.pb', '') Model size: 29265.284 KB Variables size: 0.0 KB Total Size: 29265.284 KB ---------NVIDIA RT Optimised model FP16 ('/coding/ssd_inception_v2_coco_2018_01_28/4/saved_model.pb', '') Model size: 178564.229 KB Variables size: 0.0 KB Total Size: 178564.229 KB ---------NVIDIA RT Optimised model INT8 ('/coding/ssd_inception_v2_coco_2018_01_28/5/saved_model.pb', '') Model size: 178152.834 KB Variables size: 0.0 KB Total Size: 178152.834 KB

Step 6: Ready the TF Serving container to server these models

Note the container we are using here — Client

docker run --entrypoint=/bin/bash --env http_proxy=<my proxy> --env https_proxy=<my proxy>  --runtime=nvidia  -it --rm -p 8900:8500 -p 8901:8501 -v /usr/alex/:/coding --net=host tensorflow/tensorflow:1.13.0rc1-gpu-jupyterpip install tensorflow-serving-apipip install opencv-python==3.3.0.9cd codingpython ssd_client_1.py -num_tests=1 -server=127.0.0.1:8500 -batch_size=1 -img_path='../examples/google1.jpg/'

Server -This is pasted from Step 0. This is run in the V100 32 GB Linux/machine.

docker run  --net=host --runtime=nvidia  -it --rm -p 8900:8500 -p 8901:8501 -v /usr/alex/:/models  tensorflow/serving:1.13.0-gpu --rest_api_port=0  --enable_batching=true --model_config_file=/models/ssd_inception_v3_coco.json

where the config json is like below. Since I have placed the different models in folders under “/models/ssd_inception_v2_coco_2018_01_28/” as 01 — original model, 2-TF Graph Transform Weight Quantized, 3- TF Graph Transform Weight and Node Quantized,4-TensorRT FP16,5-TensorRT INT8; I just change the versions in the file to load different servables for each test.

model_config_list { config {    name: "ssd_inception_v2_coco",    base_path: "/models/ssd_inception_v2_coco_2018_01_28/",   model_version_policy: {      specific: {         versions:[01]      }     },    model_platform:"tensorflow",  }}

Step 7: Write a TF Serving Client for tests

I have written about this in detail in a previous post.

Writing a Generic Tensorflow Serving Client for Tensorflow Serving model

For CNN based object detection models

towardsdatascience.com

The saved model of the SSD is like below You can use the saved model CLI to view it

saved_model_cli show --dir '/coding/ssd_inception_v2_coco_2018_01_28/3' --allMetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:signature_def['serving_default']:  The given SavedModel SignatureDef contains the following input(s):    inputs['input_image'] tensor_info:        dtype: DT_UINT8        shape: (-1, -1, -1, 3)        name: image_tensor:0  The given SavedModel SignatureDef contains the following output(s):    outputs['detection_boxes:0'] tensor_info:        dtype: DT_FLOAT        shape: unknown_rank        name: detection_boxes:0    outputs['detection_classes:0'] tensor_info:        dtype: DT_FLOAT        shape: unknown_rank        name: detection_classes:0    outputs['detection_scores:0'] tensor_info:        dtype: DT_FLOAT        shape: unknown_rank        name: detection_scores:0    outputs['num_detections:0'] tensor_info:        dtype: DT_FLOAT        shape: unknown_rank        name: num_detections:0  Method name is: tensorflow/serving/predict

Note that in this, the input and output node names are slightly different from the original model- whose input is ‘inputs’ and output is ‘detection_boxes’,’detection_classes’,’detection_scores’ (without the :0 part- which is a deficiency in the conversion scripts that I have used- but can be rectified easily)

Original model

root@ndn-oe:/coding/tfclient# saved_model_cli show - dir /coding/ssd_inception_v2_coco_2018_01_28/01/ - allMetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:signature_def['serving_default']: The given SavedModel SignatureDef contains the following input(s): inputs['inputs'] tensor_info: dtype: DT_UINT8 shape: (-1, -1, -1, 3) name: image_tensor:0 The given SavedModel SignatureDef contains the following output(s): outputs['detection_boxes'] tensor_info: dtype: DT_FLOAT shape: (-1, 100, 4) name: detection_boxes:0 outputs['detection_classes'] tensor_info: dtype: DT_FLOAT shape: (-1, 100) name: detection_classes:0 outputs['detection_scores'] tensor_info: dtype: DT_FLOAT shape: (-1, 100) name: detection_scores:0 outputs['num_detections'] tensor_info: dtype: DT_FLOAT shape: (-1) name: num_detections:0 Method name is: tensorflow/serving/predict

The TF Serving client is given here -https://gist.github.com/alexcpn/d7c28230af437dafb0d2cc7f50140eed

The rest of the imports are here, the client is slightly different, the names of inputs and outputs, that’s why it is on gist https://github.com/alexcpn/tf_serving_clients

The image file used for the test is https://github.com/fizyr/keras-retinanet/blob/master/examples/000000008021.jpg

Step 8: The Output from various models

Basically, there is hardly any difference between the optimized and non-optimized model. Batch size is one here.

More details below

Original Model

Invocaiton :

coding/tfclient# python ssd_client_1.py -num_tests=1 -server=127.0.0.1:8500 -batch_size=1 -img_path=’../examples/000000008021.jpg’

(‘Image path’, ‘../examples/000000008021.jpg’)(‘original image shape=’, (480, 640, 3))(‘Input-s shape’, (1, 800, 1066, 3)) → This is the size of input tensor

Ouput (‘Label’, u’person’, ‘ at ‘, array([412, 171, 740, 624]), ‘ Score ‘, 0.9980476)(‘Label’, u’person’, ‘ at ‘, array([ 6, 423, 518, 788]), ‘ Score ‘, 0.94931936)(‘Label’, u’person’, ‘ at ‘, array([ 732, 473, 1065, 793]), ‘ Score ‘, 0.88419175)(‘Label’, u’tie’, ‘ at ‘, array([529, 337, 565, 494]), ‘ Score ‘, 0.40442815)(‘Time for ‘, 1, ‘ is ‘, 0.5993821620941162)

Tensorflow Transform Optimised model Weights Quantized

(‘Label’, u’person’, ‘ at ‘, array([409, 174, 741, 626]), ‘ Score ‘, 0.99797523)(‘Label’, u’person’, ‘ at ‘, array([ 4, 424, 524, 790]), ‘ Score ‘, 0.9549346)(‘Label’, u’person’, ‘ at ‘, array([ 725, 472, 1064, 793]), ‘ Score ‘, 0.8900732)(‘Label’, u’tie’, ‘ at ‘, array([527, 338, 566, 494]), ‘ Score ‘, 0.3943166)(‘Time for ‘, 1, ‘ is ‘, 0.6182711124420 → This is higher a model size is reduced and during inference the higher precision coversion has to be done

You should see that the size of the output graph is about a quarter of the original. The downside to this approach compared to round_weights is that extra decompression ops are inserted to convert the eight-bit values back into floating point, but optimizations in TensorFlow’s runtime should ensure these results are cached and so you shouldn’t see the graph run any more slowly.- https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/graph_transforms/README.md

TensorRT FP 16 Converted model(‘Label’, u’person’, ‘ at ‘, array([412, 171, 740, 624]), ‘ Score ‘, 0.9980476)(‘Label’, u’person’, ‘ at ‘, array([ 6, 423, 518, 788]), ‘ Score ‘, 0.9493193)(‘Label’, u’person’, ‘ at ‘, array([ 732, 473, 1065, 793]), ‘ Score ‘, 0.8841917)(‘Label’, u’tie’, ‘ at ‘, array([529, 337, 565, 494]), ‘ Score ‘, 0.40442812)(‘Time for ‘, 1, ‘ is ‘, 0.5885560512542725) →

I was hoping this would be half the original value — twice as fast. But during optimization TensorRT was telling it could convert only a few of the supported* operations - "There are 3962 ops of 51 different types in the graph that are not converted to TensorRT -Conv2D" though Convolution operation is shown as supported here →https://docs.nvidia.com/deeplearning/sdk/tensorrt-support-matrix/index.html. Bug raised for this by me here

2019-04-14 08:32:31.357592: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:2019-04-14 08:32:31.357620: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 2019-04-14 08:32:31.357645: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N 2019-04-14 08:32:31.358154: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 30480 MB memory) -> physical GPU (device: 0, name: Tesla V100-PCIE-32GB, pci bus id: 0000:b3:00.0, compute capability: 7.0)2019-04-14 08:32:34.582872: I tensorflow/core/grappler/devices.cc:51] Number of eligible GPUs (core count >= 8): 12019-04-14 08:32:34.583019: I tensorflow/core/grappler/clusters/single_machine.cc:359] Starting new session2019-04-14 08:32:34.583578: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1512] Adding visible gpu devices: 02019-04-14 08:32:34.583610: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] Device interconnect StreamExecutor with strength 1 edge matrix:2019-04-14 08:32:34.583636: I tensorflow/core/common_runtime/gpu/gpu_device.cc:990]      0 2019-04-14 08:32:34.583657: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1003] 0:   N 2019-04-14 08:32:34.583986: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 30480 MB memory) -> physical GPU (device: 0, name: Tesla V100-PCIE-32GB, pci bus id: 0000:b3:00.0, compute capability: 7.0)2019-04-14 08:32:36.713848: I tensorflow/contrib/tensorrt/segment/segment.cc:443] There are 3962 ops of 51 different types in the graph that are not converted to TensorRT: TopKV2, NonMaxSuppressionV2, TensorArrayWriteV3, Const, Squeeze, ResizeBilinear, Maximum, Where, Add, Placeholder, Switch, TensorArrayGatherV3, NextIteration, Greater, TensorArraySizeV3, NoOp, TensorArrayV3, LoopCond, Less, StridedSlice, TensorArrayScatterV3, ExpandDims, Exit, Cast, Identity, Shape, RealDiv, TensorArrayReadV3, Reshape, Merge, Enter, Range, Conv2D, Mul, Equal, Sub, Minimum, Tile, Pack, Split, ZerosLike, ConcatV2, Size, Unpack, Assert, DataFormatVecPermute, Transpose, Gather, Exp, Slice, Fill, (For more information see https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html#support-ops).2019-04-14 08:32:36.848171: I tensorflow/contrib/tensorrt/convert/convert_graph.cc:913] Number of TensorRT candidate segments: 42019-04-14 08:32:37.129266: W tensorflow/contrib/tensorrt/convert/convert_nodes.cc:3710] Validation failed for TensorRTInputPH_0 and input slot 0: Input tensor with shape [?,?,?,3] has an unknown non-batch dimension at dim 12019-04-14 08:32:37.129330: W tensorflow/contrib/tensorrt/convert/convert_graph.cc:1021] TensorRT node TRTEngineOp_0 added for segment 0 consisting of 707 nodes failed: Invalid argument: Validation failed for TensorRTInputPH_0 and input slot 0: Input tensor with shape [?,?,?,3] has an unknown non-batch dimension at dim 1. Fallback to TF...2019-04-14 08:32:37.129838: W tensorflow/contrib/tensorrt/convert/convert_nodes.cc:3710] Validation failed for TensorRTInputPH_0 and input slot 0: Input tensor with shape [?,546,?,?] has an unknown non-batch dimension at dim 22019-04-14 08:32:37.129859: W tensorflow/contrib/tensorrt/convert/convert_graph.cc:1021] TensorRT node TRTEngineOp_1 added for segment 1 consisting of 3 nodes failed: Invalid argument: Validation failed for TensorRTInputPH_0 and input slot 0: Input tensor with shape [?,546,?,?] has an unknown non-batch dimension at dim 2. Fallback to TF...2019-04-14 08:32:38.309554: I tensorflow/contrib/tensorrt/convert/convert_graph.cc:1015] TensorRT node TRTEngineOp_2 added for segment 2 consisting of 3 nodes succeeded.2019-04-14 08:32:38.420585: I tensorflow/contrib/tensorrt/convert/convert_graph.cc:1015] TensorRT node TRTEngineOp_3 added for segment 3 consisting of 4 nodes succeeded.2019-04-14 08:32:38.644767: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:581] Optimization results for grappler item: tf_graph2019-04-14 08:32:38.644837: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:583]   constant folding: Graph size after: 6411 nodes (-1212), 10503 edges (-1352), time = 848.996ms.2019-04-14 08:32:38.644858: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:583]   layout: Graph size after: 6442 nodes (31), 10535 edges (32), time = 225.361ms.2019-04-14 08:32:38.644874: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:583]   constant folding: Graph size after: 6432 nodes (-10), 10535 edges (0), time = 559.352ms.2019-04-14 08:32:38.644920: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:583]   TensorRTOptimizer: Graph size after: 6427 nodes (-5), 10530 edges (-5), time = 2087.5769ms.

TensorRT INT 8 Converted model

One can see from the V100 server logs some Tensor Core magic happening

2019–04–20 01:30:39.563827: I external/org_tensorflow/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc:574] Starting calibration thread on device 0, Calibration Resource @ 0x7f4c341ac5702019–04–20 01:30:39.563982: I external/org_tensorflow/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc:574] Starting calibration thread on device 0, Calibration Resource @ 0x7f4ce8008e60

(‘Label’, u’person’, ‘ at ‘, array([412, 171, 740, 624]), ‘ Score ‘, 0.9980476)(‘Label’, u’person’, ‘ at ‘, array([ 6, 423, 518, 788]), ‘ Score ‘, 0.9493195)(‘Label’, u’person’, ‘ at ‘, array([ 732, 473, 1065, 793]), ‘ Score ‘, 0.8841919)(‘Label’, u’tie’, ‘ at ‘, array([529, 337, 565, 494]), ‘ Score ‘, 0.40442798)(‘Time for ‘, 1, ‘ is ‘, 0.5967140197753906)

With batch size 2 there is an error/ out of memory for TensorCores

python ssd_client_1.py -num_tests=1 -server=127.0.0.1:8500 -batch_size=2 -img_path=’../examples/000000008021.jpg’2019–04–20 01:34:25.042337: F external/org_tensorflow/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc:227] Check failed: t.TotalBytes() == device_tensor->TotalBytes() (788424 vs. 394212)2019–04–20 01:34:25.042373: F external/org_tensorflow/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc:227] Check failed: t.TotalBytes() == device_tensor->TotalBytes() (34656 vs. 17328)/usr/bin/tf_serving_entrypoint.sh: line 3: 6 Aborted (core dumped)

Results from other models (and Comparison with different GPU’s)

Here are some results from other tests and models

Details here — https://docs.google.com/spreadsheets/d/1Sl7K6sa96wub1OXcneMk1txthQfh63b0H5mwygyVQlE/edit?usp=sharing

Model — Resnet_50 FP 32 and FP16

FP32 = http://download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v2_fp32_savedmodel_NCHW.tar.gz

FP16 = http://download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v2_fp16_savedmodel_NCHW.tar.gz

You can see that there is a slight difference, V100 32 GB takes slightly less time than the consumer grade GTX 1070 8GB, when the batch size increases the more memory resource of V100 stands out; but not the number of CUDA cores. It seems as is noted in other blogs, that simply having more CUDA cores does not automatically mean that an inference will run faster. It may depend on memory and the model characteristics also.

Model Retinanet

One can see here that there is not much difference. Actually, this was my first experiment, but this was a Keras model that was converted to a TF frozen model and optimised. So I thought maybe I would get better results from a pure TF written model like SSD. But did not make much difference.

Summary

One can see that there are no drastic improvements in the inference time between the models. Also, TF GraphTransform for Model Quantization has not worked for me in this nor one other model I tried. Will raise a bug for that.TensorRT is better but is only able to convert a few layers to lower precision- have raised a bug/clarification for this, and if that works, hopefully, we can see the models runs twice as fast as advertised in Tensor Cores.

Main References

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/graph_transforms/README.md

https://medium.com/google-cloud/optimizing-tensorflow-models-for-serving-959080e9ddbf

https://colab.research.google.com/drive/1wQpWoc40kf__WSjfTqDaReMx6fFjUn48

Other related posts

Writing a Generic Tensorflow Serving Client for Tensorflow Serving model

For CNN based object detection models

towardsdatascience.com

Using Tensorflow Serving GRPC

Once you have your Tensorflow or Keras based model trained, one needs to think on how to use it in production. You may…

towardsdatascience.com

Running Tensor Flow and Keras in Jupyter notebook via Docker

Assuming that you have NVIDIA GPU in your machine and NVIDIA Drivers installed. Please read the below before you install…

medium.com

--

--

3

Published in Better Software

People Process Technologies in the Software Engineering

Written by Alex Punnen

SW Architect/programmer- in various languages and technologies from 2001 to now. https://www.linkedin.com/in/alexpunnen/

Responses (3)

Help

Status

About

Careers

Press

Blog

Privacy

Terms

Text to speech

Teams


Automating GPU Kernel Generation with DeepSeek-R1 and Inference Time Scaling

As AI models extend their capabilities to solve more sophisticated challenges, a new scaling law known as test-time scaling or inference-time scaling is emerging. Also known as AI reasoning or long-thinking, this technique improves model performance by allocating additional computational resources during inference to evaluate multiple possible outcomes and then selecting the best one, neural network. This enables AI to strategize and systematically solve complex problems in a similar fashion to how humans dissect complex problems and solve them individually to arrive at a final solution.

In this post, we talk about an experiment done by NVIDIA engineers who used one of the newest open-source models, the DeepSeek-R1 model, together with additional computing power during inference to solve a complex problem. The experiment was to automatically generate GPU attention kernels that were numerically correct and optimized for different flavors of attention without any explicit programming.

The results turned out to be better than the optimized kernels developed by skilled engineers in some cases.

The need for optimized attention kernels and associated challenges

Attention is a key concept that revolutionized the development of the large language model (LLM). It’s a powerful mechanism that enables AI models to focus selectively on the most relevant parts of input when performing tasks. By focusing on important information, the attention operation helps the models make better predictions and find hidden patterns in the data.

The computational complexity of the attention operation grows quadratically in relation to the input sequence length. This motivates the need for developing an optimized lower-level implementation (that is, a GPU kernel) to prevent runtime errors arising from simple implementations (for example, out-of-memory errors) and for computational efficiency purposes.

There are multiple variants of attention (causal, relative positional embeddings, alibi, and so on) and often engineers must use a combination of these variants for a given task.

Multi-modal models (for example, vision transformers) introduce an additional layer of challenges as they require specialized attention mechanisms (Spatial Neighborhood Attention) for maintaining spatio-temporal information often encountered in computer vision, video generation models, and so on.

Creating an optimized GPU kernel for attention takes a lot of skill and time, even for experienced software engineers.

Recent LLMs like DeepSeek-R1 have shown a lot of promise in code generation tasks, but they still face challenges creating optimized code on the first try. This makes it necessary to use other strategies at inference time to generate optimized code.

The following prompt is sample user input for a relative positional embeddings attention kernel.

Please write a GPU attention kernel to support relative position encodings. Implement the relative positional encoding on the fly within the kernel. The complete code should be returned, including the necessary modifications.

Use the following function to compute the relative positional encoding:

def relative_positional(score, b, h, q_idx, kv_idx):

    return score + (q_idx - kv_idx)

When implementing the kernel, keep in mind that a constant scaling factor 1.44269504 should be applied to the relative positional encoding due to qk_scale = sm_scale * 1.44269504. The PyTorch reference does not need to scale the relative positional encoding, but in the GPU kernel, use:

qk = qk * qk_scale + rel_pos * 1.44269504

Please provide the complete updated kernel code that incorporates these changes, ensuring that the relative positional encoding is applied efficiently within the kernel operations.

LLMs can occasionally produce hallucinated code or mix syntax from different languages or frameworks, causing immediate code errors or inefficiencies. Computing the optimal GPU thread mapping is also non-trivial and a challenging task, often requiring iterative refinement to achieve a correct and efficient kernel.

Inference-time scaling for generating optimized GPU Kernels

To get the best results with optimized attention kernels, NVIDIA engineers created a new workflow that includes a special verifier along with the DeepSeek-R1 model during inference in a closed-loop fashion for a predetermined duration.

The workflow is first initialized by a manual prompt and the DeepSeek-R1 model generates the GPU code (that is, the kernel) in the first pass. The verifier runs on an NVIDIA H100 GPU. It analyzes the generated kernel and creates new prompts that are provided as input to the DeepSeek-R1 model.

This closed-loop approach makes the code generation process better by guiding it in a different way each time. The team found that by letting this process continue for 15 minutes resulted in an improved attention kernel.

This workflow produced numerically correct kernels for 100% of Level-1 problems and 96% of Level-2 problems, as tested by Stanford’s KernelBench benchmark.

The Level-1 solving rate in KernelBench refers to the numerical correct metric used to evaluate the ability of LLMs to generate efficient GPU kernels for specific computational tasks. This test is part of a series of challenges to test the latest LLMs’ abilities in GPU programming.

Figure 4 shows how the inference-time budget affects the agent’s solving rate. Allocating more than 10 minutes per problem in the Level-1 category enables the workflow to produce numerical correct code for most of the 100 problems.

Optimized GPU kernels on DeepSeek-R1

These results show how you can use the latest DeepSeek-R1 model to give better GPU kernels by using more computing power during inference time. This is still a new research area with early results on a promising approach that automatically generates effective attention kernels.

While we are off to a good start, more work is needed to generate better results consistently for a wider variety of problems. We’re excited about the recent developments in DeepSeek-R1 and its potential.

For more information or to get started, see the DeepSeek-R1 NIM microservice, now available on build.nvidia.com.

Related resources

Tags

About the Authors

Comments

Related posts

Low Latency Inference Chapter 2: Blackwell is Coming. NVIDIA GH200 NVL32 with NVLink Switch Gives Signs of Big Leap in Time to First Token Performance

Deploying Accelerated Llama 3.2 from the Edge to the Cloud

NVIDIA NeMo Accelerates LLM Innovation with Hybrid State Space Model Support

Optimizing Inference on Large Language Models with NVIDIA TensorRT-LLM, Now Publicly Available

NVIDIA Slashes BERT Training and Inference Times

Related posts

Spotlight: NAVER Place Optimizes SLM-Based Vertical Services with NVIDIA TensorRT-LLM

NVIDIA AI Enterprise Adds Support for NVIDIA H200 NVL

Optimizing Qwen2.5-Coder Throughput with NVIDIA TensorRT-LLM Lookahead Decoding

Just Released: Tripy, a Python Programming Model For TensorRT

OpenAI Triton on NVIDIA Blackwell Boosts AI Performance and Programmability


Analysis-Driven Optimization: Preparing for Analysis with NVIDIA Nsight Compute, Part 1

In this three-part series, you discover how to use NVIDIA Nsight Compute for iterative, analysis-driven optimization. Part 1 covers the background and setup needed, part 2 covers beginning the iterative optimization process, and part 3 covers finishing the analysis and optimization process and determining whether you have reached a reasonable stopping point.

Nsight Compute is the primary NVIDIA CUDA kernel-level performance analysis tool. It is part of the NVIDIA Nsight family of tools for GPU computing. For thorough introductions to the NVIDIA Nsight family profiler tools, see the following posts:

These posts point out that the GPU code performance analysis process usually begins with Nsight Systems. Eventually, the analysis may select a specific kernel to focus on, for further analysis using Nsight Compute. In this post, I discuss how Nsight Compute facilitates analysis-driven optimization (ADO) of GPU kernels.

ADO is predicated on the idea that making the most efficient use of your time involves focusing on the most important limiters to code performance, in pareto order. In a cyclical process, you want the tool to identify the current most important limiter to performance and, to the fullest extent possible, give you some clues about how to address it. The most important limiter to performance is the code characteristic that, if modified, would yield the largest improvement in performance.

Start by focusing on the fixes that yield the largest performance improvements. In a cyclical fashion, you use the tool to identify these areas, make code changes, and then use the tool again to assess the impact of these changes and identify the next area to look at. The process completes when you either run out of time or have identified, perhaps through some calculations, that further optimization is unlikely to yield significant performance improvement.

To follow along with this post, I recommend using CUDA 11.1 and Nsight Compute 2020.2 or newer. Several kinds of profiler output reviewed in this post may not be present if you use a previous version of the tool.

Code for analysis

You’re going to analyze code that is more complicated than my previous post on Nsight Compute. This code has multiple steps and phases to it. Using ADO, you go through a set of steps that successively improves the performance of the code by uncovering, in each step, a performance limiter. For more information about an additional example, see the Summary section.

Use the performance limiter to guide your attempt to improve performance. The code that you analyze has two major phases:

This code repeats these steps on different sets of incoming vectors but uses the same matrix to produce a set of result vectors. The starting point for this code was a CPU algorithm using OpenMP for parallelism that has been ported to a GPU version to achieve higher performance.

// compile with: nvcc -Xcompiler -fopenmp -o t5 t5.cu -O3 -lineinfo
#include
#include

#define cudaCheckErrors(msg) \
  do { \
    cudaError_t __err = cudaGetLastError(); \
    if (__err != cudaSuccess) { \
        fprintf(stderr, "Fatal error: %s (%s at %s:%d)\n", \
            msg, cudaGetErrorString(__err), \
            __FILE__, __LINE__); \
        fprintf(stderr, "*** FAILED - ABORTING\n"); \
        exit(1); \
    } \
  } while (0)

#include
#include <sys/time.h>
#define USECPSEC 1000000ULL

unsigned long long dtime_usec(unsigned long long start){

  timeval tv;
  gettimeofday(&tv, 0);
  return ((tv.tv_sec*USECPSEC)+tv.tv_usec)-start;
}

// perform vector averaging over M vectors of length L,  followed by matrix-vector multiply
// repeat the above N times
// input vectors are stored as a set of N column-major matrices
// for each k in N: output[k] = matrix*input[k]
template
void cpu_version1(T *input, T *output, T *matrix, int L, int M, int N){
#pragma omp parallel for
  for (int k = 0; k < N; k++){      // repeat the following, N times
    std::vector v1(L);           // vector length of L
    for (int i = 0; i < M; i++)     // compute average vector over M input vectors
      for (int j = 0; j < L; j++)
        v1[j] += input[k*M*L+j*M+i];
    for (int j = 0; j < L; j++)
      v1[j] /= M;
    for (int i = 0; i < L; i++)     // matrix-vector multiply
      for (int j = 0; j < L; j++)
        output[i*N+k] += matrix[i*L+j]*v1[j];
  }
}

const int my_L = 1024;
const int my_M = 1024;
const int my_N = 1024;

template
__global__ void gpu_version1(const T * __restrict__ input, T * __restrict__ output, const T * __restrict__ matrix, const int L, const int M, const int N){

  __shared__ T smem[my_L];
  size_t idx = ((size_t)blockIdx.x)*blockDim.x + threadIdx.x;
  for (int k = 0; k < N; k++){  // iterate over N data sets
    T v1 = 0;
    for (int i = 0; i < M; i++) // perform vector averaging
      v1 += input[k*M*L+idx*M+i];
    v1 /= M;
    for (int i = 0; i < L; i++){ // perform matrix-vector multiply
      __syncthreads();
      smem[threadIdx.x] = v1 * matrix[i*L+idx];
      for (int s = blockDim.x>>1; s > 0; s>>=1){
        __syncthreads();
        if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x+s];}
      if (!threadIdx.x) output[k+i*N] = smem[0];}
  }
}

typedef float ft;

int main(){
  ft *d_input, *h_input, *d_output, *h_outputc, *h_outputg, *d_matrix, *h_matrix;
  int L = my_L; int M = my_M; int N = my_N;
  // host allocations
  h_input   = new ft[N*L*M];
  h_matrix  = new ft[L*L];
  h_outputg = new ft[N*L];
  h_outputc = new ft[N*L];
  // data initialization
  for (int i = 0; i < N*L*M; i++) h_input[i] = (rand()&1)+1;  // 1 or 2
  for (int i = 0; i < L*L; i++) h_matrix[i]  = (rand()&1)+1;  // 1 or 2
  // create result to test for correctness
  unsigned long long dt = dtime_usec(0);
  cpu_version1(h_input, h_outputc, h_matrix, L, M, N);
  dt = dtime_usec(dt);
  std::cout << "CPU execution time: " << dt/(float)USECPSEC << "s" << std::endl;
  // device allocations
  cudaMalloc(&d_input, N*L*M*sizeof(ft));
  cudaMalloc(&d_output,  N*L*sizeof(ft));
  cudaMalloc(&d_matrix,  L*L*sizeof(ft));
  cudaCheckErrors("cudaMalloc failure");
  // copy input data from host to device
  cudaMemcpy(d_input,  h_input,  N*L*M*sizeof(ft), cudaMemcpyHostToDevice);
  cudaMemcpy(d_matrix, h_matrix,   L*L*sizeof(ft), cudaMemcpyHostToDevice);
  cudaMemset(d_output, 0, N*L*sizeof(ft));
  cudaCheckErrors("cudaMemcpy/Memset failure");
  // run on device and measure execution time
  dt = dtime_usec(0);
  gpu_version1<<<1, L>>>(d_input, d_output, d_matrix, L, M, N);
  cudaCheckErrors("kernel launch failure");
  cudaDeviceSynchronize();
  cudaCheckErrors("kernel execution failure");
  dt = dtime_usec(dt);
  cudaMemcpy(h_outputg, d_output, N*L*sizeof(ft), cudaMemcpyDeviceToHost);
  cudaCheckErrors("cudaMemcpy failure");
  for (int i = 0; i < N*L; i++) if (h_outputg[i] != h_outputc[i]) {std::cout << "Mismatch at " << i << " was: " << h_outputg[i] << " should be: " << h_outputc[i] << std::endl; return 0;}
  std::cout << "Kernel execution time: " << dt/(float)USECPSEC << "s" << std::endl;
  return 0;
}

Some highlights:

The code isn’t reflective of any specific scientific algorithm. However, to give a real-world application, the vector-averaging phase could reflect a naive form of a calculation performed by a DNN parameter server, and the matrix-vector multiply is used throughout scientific and AI codes and could be the basis of a naive non-batched DNN update.

The code provides for built-in timing measurement of the CPU and GPU operations, so that you can quickly assess the measure of speedup or benefit from the latest optimization. You can also use Nsight Compute directly to measure kernel duration. In a more complete treatment of using NVIDIA Nsight tools for performance analysis or optimization, you might include Nsight Systems in the iterative analysis loop. In addition, results checking is performed between CPU and GPU versions so that you can be sure that your GPU optimized versions are producing the correct results.

This is a simplistic comparison, testing for exact equality between results. This is normally not the recommended way to compare floating-point results. Because the scope of the problem and test data is limited, this method is acceptable. For general floating-point comparison, equality should be tested against some measure of a difference threshold, not exact equality.

Initial performance baseline

If you compile and run this code on a V100 GPU, you see the following results:

$ nvcc -Xcompiler -fopenmp -o t5 t5.cu -O3 -lineinfo
$ OMP_NUM_THREADS=1 ./t5
CPU execution time: 5.65601s
Kernel execution time: 2.922s
$ ./t5
CPU execution time: 0.52372s
Kernel execution time: 2.9219s
$

If you run with only a single CPU thread, then the initial port of CPU code to GPU code seems to give about a 2x speedup. If you don’t restrict the number of CPU threads, however, the OpenMP parallelization seems to give about a 10x speedup to the CPU code, meaning that your initial CUDA kernel realization is about 5x slower than the CPU code. See if you can improve this.

Getting started with Nsight Compute

The Nsight Compute profiler can collect a large range of data on your kernel execution. In addition, you make use of rules embedded in the analysis output from Nsight Compute. A rule in Nsight Compute is a set of instructions to the profiler that indicate what metrics are to be gathered and how they are to be displayed or interpreted.

The rule system in Nsight Compute is a powerful feature that allows extending the functionality provided. It is possible to create your own rules, but this analysis uses rules that are already available. You use the bottleneck rule to guide your steps. For most of the work that you are doing in this post, you use the user interface version of Nsight Compute. The Nsight Compute user interface can be used directly, in-situ, for installations that support it. Alternatively, you can collect Nsight Compute report data from the Nsight Compute CLI and import that data into a session running elsewhere.

To capture the information needed during this investigation, Nsight Compute must have access to profiling features that require permission at the GPU driver level. Depending on your operating system, there are various methods to enable this.

You can start Nsight Compute user interface on the target using the method given in a previous post. For example, using Nsight Compute 2020.2 with appropriate path setup, type ncu-ui. When the initial dialog box opens, choose Quick Launch, Continue.

In the next dialog box, under Target Platform, choose the appropriate option. I am working on a Linux x86_64 platform so for me, that choice is already selected. For Application Executable, enter the full path/name to the executable that you just compiled and ran. You can use the browse (…) button to navigate and find the app. In the Activity section, make sure that Interactive Profile is selected. You don’t have to make any other changes. Choose Launch.

Nsight Compute then launches the app and allows it to proceed up to the first CUDA call. This step may take about 10 seconds depending on your CPU, because the code is setting up and initializing 4 GB of memory. You haven’t profiled anything yet. Because your application has only one kernel in it, called only one time, you can quickly get to profiling by choosing Auto Profile, selecting the Full section set, and then choosing Run to Next Kernel.

What you discover at this point is that the application is taking a long time to profile, compared to the kernel duration. The Nsight Compute tool may require substantial data collection to gather all the requested metrics to generate a full section set in your report. For more information about why the tool may have longer profiling times, see Kernel replay.

You can reduce the profiling scope, while still getting useful results, by reducing the number of data sets N to process in the code. Make the following code change:

const int my_N = 32;

Recompile the code. In the profiler, choose Terminate, and repeat the earlier steps. Now, profiling after you choose Run to Next Kernel should only take about 10 seconds.

Profiling results

You now have your initial profiler results (Figure 2).

The profiler results are organized into sections, and the sections are arranged from top to bottom in roughly attention-priority order. Give your attention to the top section first, GPU Speed of Light (Figure 2). The section provides a high-level overview of the utilization for compute and memory resources of the GPU. For each unit, the Speed of Light (SOL) reports the achieved percentage of utilization with respect to the theoretical maximum. As previously mentioned, you want the tool to guide your analysis, and you have the bottleneck rule to help with that. The bottleneck rule is presented in the SOL section:

[Warning] This kernel grid is too small to fill the available resources on this device. Look at Launch Statistics for more details.

Launch Statistics is reporting a launch grid of one block (Figure 3).

You see a similar recommendation in this section. As a well-trained CUDA programmer, you know that such a small grid cannot hope to fill any GPU, and such a grid runs well below the expected GPU performance. The solution is to increase the grid size, by increasing the number of blocks launched. For more information about grids, blocks, and SMs, and how scheduling more blocks can help to achieve better performance, see Hardware model.

Summary

In this post, I introduced the code for profiling, covered the basic ideas of ADO, and got you started with the Nsight Compute profiler. In part 2, you continue the ADO process, by applying what you learned in this post to refactor the code, and profile again to discover next steps. In part 3, you finish your analysis and optimization. You also perform some measurements to give you confidence that you have reached a reasonable stopping point.

The analysis work in this post was performed on a machine with the following characteristics: Ubuntu 18.04.3, CUDA 11.1, GPU Driver version 455.23.05, GCC version 7.5.0, V100-SXM2-32 GB GPU, Intel(R) Xeon(R) Gold 6130 CPU @ 2.10GHz. The code examples presented in this post are for instructional purposes only. They are not guaranteed to be defect-free or suitable for any particular purpose.

For more information, see the following resources:

Acknowledgements

The author would like to thank the following individuals for their contributions:  Sagar Agrawal, Rajan Arora, Ronny Brendel, Max Katz, Felix Schmitt, Greg Smith, and Magnus Strengert.

Related resources

Tags

About the Authors

Comments

Related posts

CUDA Toolkit 12.3 Delivers New Features for Accelerated Computing

New Video Series: CUDA Developer Tools Tutorials

NVIDIA GTC: A Complete Overview of Nsight Developer Tools

SC20 Demos: New Nsight Systems and Nsight Compute Demos

Unleashing the Power of NVIDIA Ampere Architecture with NVIDIA Nsight Developer Tools

Related posts

CUDA Toolkit Now Available for NVIDIA Blackwell

Improving GPU Performance by Reducing Instruction Cache Misses

Optimizing llama.cpp AI Inference with CUDA Graphs

Just Released: Nsight Compute 2024.3

Just Released: CUDA Toolkit 12.6