teaching_llm_applications

GPUs and Flash Attention

🤔 ❓Question

image

Motivation

image

What does a GPU look like?

image

Resources

🤔 Question for class

🤔 Question for class

Animation of how GPU/TPU works

image

image

image

Introduction

image

image

Software model

TPU

Back to GPUs

Summary

Recap

recap

Control divergence

Low precision computation

image

image

Operator fusion

image

image

image

image

Burst mode of DRAM/global memory

Tiling

Circling back to motivation

image

image

Recap

image

Flash Attention

image

image

Key takeaways

Benchmarking and profiling

add_profile = profile(run_operation2(dim=2048, operation = lambda a, b: a + b))

import torch

def builtin_gelu(x: torch.Tensor):
    return torch.nn.functional.gelu(x, approximate="tanh")


import os
import time
from typing import Callable
import torch
from torch.profiler import ProfilerActivity
import triton
import triton.language as tl
from edtrace import text, link, image
from lecture_util import get_local_url
from gpu_util import cuda_if_available

def main():
    # Last lecture: high-level overview of GPUs and performance
    # This lecture: benchmarking/profiling + writing kernels
    # review_of_gpus()
    benchmarking_and_profiling()           # Where are the bottlenecks?
    naive_vs_builtin_vs_compiled_gelu()    # Apply it to the GeLU example
    # Write Triton kernels

def benchmarking_and_profiling():
    # Recipe for success:

    # Benchmark and profile your code
    
    # Make changes
    
    # Benchmark and profile your code again
    
    benchmarking()   # How long does it take?
    
    profiling()      # Where time is being spent?
    
    # Benchmark and profile your code!

def benchmarking():
    # Benchmarking measures the wall-clock time of performing some operation.
    # It only gives you end-to-end time, not where time is spent (profiling).
    # It is still useful for:
    
    # comparing different implementations (which is faster?), and
    
    # understanding how performance scales (e.g., with dimension).
    
    # You can use torch.utils.benchmark.
    # We will roll our own to make benchmarking more transparent.
    
    # Benchmark matrix multiplication
    
    matmul = run_operation2(dim=1024, operation=lambda a, b: a @ b)
    result = benchmark(matmul)  
    
    # See how timing scales with dimension
    results = {}
    for dim in [256, 512, 1024, 2048, 4096, 8192]:
         results[dim] = benchmark(run_operation2(dim=dim, operation=lambda a, b: a @ b))  
    # Note: time is roughly constant when dimension is small, then cubic scaling.

def benchmark(run: Callable, num_warmups: int = 1, num_trials: int = 3) -> float:
    """Benchmark `func` by running it `num_trials`.  Return the average time."""
    # Warmup: first times might be slower due to compilation, etc.
    # Since we will run the kernel multiple times, the timing that matters is steady state.
    for _ in range(num_warmups):
        run()
    torch.cuda.synchronize()  # Wait for CUDA threads to finish (important!)
    # Time it for real now!
    times: list[float] = [] 
    for trial in range(num_trials):  # Do it multiple times to capture variance
        # Use CUDA events for accurate GPU timing (avoid capturing CPU overhead)
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()  # Start timing
        run()  # Actually perform computation
        end_event.record()  # End timing
        torch.cuda.synchronize()  # Wait for CUDA threads to finish
        times.append((start_event.elapsed_time(end_event))) 

    mean_time = mean(times)   
    return mean_time


def profiling():
    # While benchmarking looks at end-to-end time, profiling looks at where time is spent.
    # Independent of time, profiling also helps you understand what's going under the hood.
    # PyTorch has a built-in profiler.
    # In your assignment, you will use nsight to get more details.
    add(dim=2048)
    add_profile = profile(run_operation2(dim=2048, operation=lambda a, b: a + b))

    matmul_profile = profile(run_operation2(dim=2048, operation=lambda a, b: a @ b)) 

    matmul_profile = profile(run_operation2(dim=128, operation=lambda a, b: a @ b)) 


def naive_vs_builtin_vs_compiled_gelu():
    # Let's benchmark and profile the GeLU activation function.
    x = torch.tensor([1.])  
    # 1. Implementation naively from scratch in PyTorch (non-fused)
    y1 = naive_gelu(x)  
    # 2. Built-in PyTorch implementation (fused)
    y2 = builtin_gelu(x) 

    check_equal_1d(naive_gelu, builtin_gelu)  # Check it works
    # 3. Use PyTorch compiler on the naive implementation
    compiled_gelu = torch.compile(naive_gelu)  

    y3 = compiled_gelu(x)  

    check_equal_1d(naive_gelu, compiled_gelu)  # Check it works (compilation shouldn't change semantics) 
    # Benchmarking
    naive_time = benchmark(run_operation1(dim=16384, operation=naive_gelu)) 

    builtin_time = benchmark(run_operation1(dim=16384, operation=builtin_gelu)) 

    compiled_time = benchmark(run_operation1(dim=16384, operation=compiled_gelu)) 
    # The builtin and compiled versions are significantly faster!
    # To understand why, let's look at the profiler to see where time is being spent.
    # naive_gelu
    naive_gelu_profile = profile(run_operation1(dim=16384, operation=naive_gelu))

TODO 📚: Questions Assignment

🎮 TODO Practical

This site is open source. Improve this page.