LLM Systems 03: GPU Programming-2

This article is the learning note of CMU 11868 LLM Systems course Lecture 3. For more details, please refer to the original slides.

0. Recap

• GPU is composed of
  • streaming processing units (SMs)
    • each with four partitions of 32 cores
    • shared L1 cache
  • memory
  • L2 cache: share with all SMs
• Threads organized in
  • grid of thread blocks
  • each block is divided into warps running in parallel on one SM.

1. Basic GPU CUDA operations

1.1 memory management

Typical CUDA programming steps:

• CPU allocates GPU memory: cudaMalloc
• CPU copies data to GPU memory (host to device): cudaMemcpy
• CPU launches GPU kernels
• CPU copies results from GPU (device to host): cudaMemcpy
• Freeing GPU memory: cudaFree

GPU memory hierarchy:

• Each thread has private registers (fastest to access)
• Each thread block has shared memory
  • Visible to all threads in a block
  • __shared__
• All threads can access global gpu memory
  • Persistent across kernel launches in the same app

TMA(Tensor Memory Accelerator) in H100:

• void cuda::memcpy_async(void* destination, void const* source, Shape size, cuda::barrier<Scope, CompletionFunction>& barrier);
• call barrier.wait() to wait TMA copy completes.

• 2D Tensor layout :
  • base addr: start of tensor in memory.
  • Tensor width / height: dimensions of the tensor.
  • Tensor stride: bytes between consecutive rows (includes padding).
  • block width / height: the sub-region (pink) to copy.
  • padding: extra memory at row ends; TMA handles stride and skips padding automatically.

1.2 defining functions to be executed on GPU

Declaration of Host/Device function

keyword	call on	execute on
__global__	host (cpu)	device (gpu)
__device__	device (gpu)	device (gpu)
__host__	host	host

Define kernel function:

__global__ void VecAddKernel(int* A, int* B, int* C, int n) {
	int i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < n) {
		C[i] = A[i] + B[i];
	}
}

int main() {
	VecAddKernel<<<1, N>>>(A, B, C, N);
}

1.3 launch kernel

Calling Kernel at Runtime

• Host program specifies grid-block-threads configurations for kernel at run time
  • Dg and Db are either dim3 or int

dim3 Dg(4, 2, 1);
dim3 Db(8, 8, 1);
kernelFuncName<<<Dg, Db>>>(args)

• Dg: size of grid (num. of blocks)
  • Dg.x * Dg.y * Dg.z is num. of blocks
• Db: size of block
  • Db.x * Db.y * Db.z is num. of threads per block, <=1024)

Device Runtime Variables

• Host launches kernels on a gpu device
• Each kernel thread needs to know which thread it is running
• Compiler generates build-in variables, with x, y, z fields

Name	Type	Explanation
gridDim	dim3	dimensions of grid
blockIdx	uint3	index of block within grid
blockDim	dim3	dimensions of block
threadIdx	uint3	index of thread within block

Calling CUDA Kernel from CPU

// n: the size of the vector
int n = 1024;
int threads_per_block = 256;
int num_blocks = (n + threads_per_block - 1) / threads_per_block;
VecAddKernel<<<num_blocks, threads_per_block>>>(dA, dB, dC, n);

2. Matrix/Tensor Computation on GPU

Matrix Add

__global__ void MatAddKernel(float* A, float* B, float* C, int N) {
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	int j = blockIdx.y * blockDim.y + threadIdx.y;
	C[i * N + j] = A[i * N + j] + B[i * N + j];
}

int main() {
	int N = 32;
	dim3 threads_per_block(N, N);
	int num_blocks = 1;
	MatAddKernel<<<num_blocks, threads_per_block>>>(dA, dB, dC, N);
}

3-D Grids & Blocks Demo

__global__ void fullKernel(float* din, float* dout) {
	// global tid computation process for 3-D grids & blocks.
	int block_id = blockIdx.x +
				   blockIdx.y * gridDim.x +
				   blockIdx.z * gridDim.x * gridDim.y;
	int block_offset = block_id * blockDim.x * blockDim.y * blockDim.z;

	int thread_offset = threadIdx.x +
				        threadIdx.y * blockDim.x +
				        threadIdx.z * blockDim.x * blockDim.y;
	int tid = block_offset + thread_offset;

	dout[tid] = func(din[tid]);
}

int main() {
	dim3 threads_per_block(2, 4, 8);
	dim3 blocks_per_grid(2, 3, 4);
	fullKernel<<<blocks_per_grid, threads_per_block>>>(some_input, some_output);
}

3. Summary

• Basic GPU CUDA operations
  • memory allocation
  • data movement
  • creating threads and running on SMs
    • specifying number of threads and number of blocks in a grid
  • referring to data in GPU memory within a thread
    • using building index variables to refer to the data