1.Review of GPUs

• 每个 Thread 独享一组 Register
• 每个 Block 占有一块 shared_memory ( shared_memory 位于 SM 内部，但即使一个 SM 上被分配了多个 Block，这些 Block 之间也不能共享 shared_memory，所以一个 SM 内的多个 Block 之间的通信也比较昂贵)
• Block 内部共享 shared_memory，通信代价低；跨 Block 代价高。

所有需要交互的数据都应尽量保持在同一个 Block 或计算单元内部，确保运算速度达到极致。

2. Benchmarking and Profiling

2.1 Benchmarking

做一次 warm_up；torch.cuda.synchronize() 保持 GPU 和 CPU 状态同步（ Assignment 1 里写过）

2.1.1 sleep()

mean_time = benchmark("sleep", lambda: time.sleep(50/1000))
print("mean_time = ", mean_time)

结果为（RTX 3070）：

2.1.2 矩阵乘法

def run_operation2(dim, operation):
    a = torch.randn(dim, dim)
    b = torch.randn(dim, dim)

    def run():
        return operation(a, b)

    return run


if torch.cuda.is_available():
    dims = (1024, 2048, 4096, 8192, 16384)  # @inspect dims
else:
    dims = (1024, 2048)  # @inspect dims

matmul_results = []
for dim in dims:
    # @ inspect dim
    result = benchmark(f"matmul(dim={dim})", run_operation2(dim=dim, operation=lambda a, b: a @ b))
    matmul_results.append((dim, result))  # @inspect matmul_results

for dim, time_ms in matmul_results:
    print(f"{dim:<10} {time_ms:<15.4f}")
    # :<10 ———— 左对齐，占用10个字符宽度
    # :<15.4f ———— 左对齐，总宽度15格，保留4位小数

结果为：

发现了一个比较有意思的点：对于自己的 3070 小破卡，基本上时间差都是 8 倍：因为对于维度为 n 的方阵来说，矩阵乘法的复杂度为O(n^3)。所以维度为 2 倍，所需时间为 8 倍。但对于 H100 来说，前期不遵循 8 倍关系，只有最后 8192 → 16384 大致遵循。猜测是因为 H100 计算力太强，直到 8192 维，Roofline 还停留在左侧的内存带宽瓶颈，没有完全释放计算性能。

2.1.3 MLP

class MLP(nn.Module):
    """Simple MLP: linear -> GeLU -> linear -> GeLU -> ... -> linear -> GeLU"""
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(dim, dim) for _ in range(num_layers)])

    def forward(self, x: torch.Tensor):
        for layer in self.layers:
            x = layer(x)
            x = torch.nn.functional.gelu(x)
        return x


def run_mlp(dim: int, num_layers: int, batch_size: int, num_steps: int) -> Callable:
    # Callable 表示返回的是一个可调用的函数
    # Define a model (with random weights)
    model = MLP(dim, num_layers).to(get_device())

    # Define an input (random)
    x = torch.randn(batch_size, dim, device=get_device())

    def run():
        # Run the model `num_steps` times (note: no optimizer updates)
        for step in range(num_steps):
            # Forward
            y = model(x).mean()

            # Backward
            y.backward()

    return run

一个简单的 MLP 模型，甚至没计算损失函数，只是用了 model(x) 的平均值 .mean() 来计算梯度。

Test 1: 

Test 2：在 Test 1 的基础上，引入新变量 scale，分别与 run_mlp 的四个参数相乘。发现运行时间与 num_layers 以及 num_steps 呈现线性相关。

2.2 Profiling

benchmark 过于粗粒度，只能表示代码所用时长。而 profiling 不仅能精确显示时间消耗在哪些函数，还能追溯到代码与 PyTorch 接口的交互，追踪从高层到底层的调用，看到底层实际执行的命令，可以更直观地理解程序如何在硬件上实际执行。

def profile(description: str, run: Callable, num_warmups: int = 1, with_stack: bool = False):
    # Warmup
    for _ in range(num_warmups):
        run()
    if torch.cuda.is_available():
        torch.cuda.synchronize()  # Wait for CUDA threads to finish (important!)

    # Run the code with the profiler
    with torch.profiler.profile(
            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
            # 同时监测 CPU 端和 GPU 端的所有操作
            with_stack=with_stack,
            # with_stack 如果为 True，会记录每个操作的 Python 调用栈（用于后面生成堆栈跟踪可视化）
            experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)
            # 启用更详细的实验性配置，让 profiler 输出更多底层信息
    ) as prof:
        run()
        if torch.cuda.is_available():
            torch.cuda.synchronize()  # Wait for CUDA threads to finish (important!)

    # Print out table
    table = prof.key_averages().table(sort_by="cuda_time_total",
                                      max_name_column_width=80,
                                      row_limit=10)
    # 按 CUDA 总耗时从高到低排序（cuda_time_total 是该操作在所有调用中花费的 GPU 总时间）
    # max_name_column_width=80：限制操作名称列的最大宽度，避免过长。
    # row_limit=10：只显示耗时最长的前 10 个操作

    # Write stack trace visualization
    # 用于生成火焰图
    if with_stack:
        text_path = f"var/stacks_{description}.txt"
        svg_path = f"var/stacks_{description}.svg"
        prof.export_stacks(text_path, "self_cuda_time_total")

    return table

2.2.1 sleep()

sleep_function = lambda: time.sleep(50 / 1000)
sleep_profile = profile("sleep", sleep_function)

print(sleep_profile)

2.2.2 矩阵加法

def run_operation2(dim: int, operation: Callable) -> Callable:
    # Setup: create two random dim x dim matrices
    x = torch.randn(dim, dim, device=device)
    y = torch.randn(dim, dim, device=device)
    # Return a function to perform the operation
    return lambda: operation(x, y)

add_function = lambda a, b: a + b
add_profile = profile("add", run_operation2(dim=2048, operation=add_function))
print(add_profile)

表头中 Self 指不包括其子调用，Total 指包括子调用

①. aten::add 是 Pytorch 底层 C++ 核心库中的加法操作函数。（aten：A Tensor ENgine）

②. Unrecognized

③. 第三项一大串：被 C++名称修饰（Name Mangling）后的符号名，还原后为：at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<float, float>, ...>。  这是一个 CUDA 向量化逐元素内核，负责 float 型加法操作。                                                        外函数 vertorized_elementwise_kernel 进行逐元素向量化，内部第一个参数 4 代表向量化宽度，即每个 CUDA 线程一次处理 4 个元素；第二个参数 CUDAFunctor_add 指明执行的操作为加法。这是执行加法操作的核心部分。

④. cudaLaunchKernel：CPU 接受指令并发送给 GPU 的过程。

注意到这四项的 Self CPU 时间和恰好等于 aten::add 的 CPU total，说明它们为 aten::add 的子函数

⑤. cudaDeviceSynchronize：等待 GPU 完成计算并传回数据

2.2.3矩阵乘法

matmul_function = lambda a, b: a @ b
matmul_profile = profile("matmul", run_operation2(dim=2048, operation=matmul_function))
print(matmul_profile)

分析方法差不多，前六项的 Self CPU 之和为 aten::matmul 的 CPU total。

①. aten:matmul：矩阵乘法入口，只负责调度，实际运算位于子函数；

②. aten:mm：二维矩阵乘法底层函数

③. Unrecognized

④. ampere_sgemm_128x64_nn：NVDIA的线性代数库 cuBLAS 根据矩阵大小以及 GPU 硬件选出的高性能计算方案

⑤. cudaOccupancyMaxActiveBlocksPerMultiprocessor：计算给定内核函数在 GPU 的一个 SM 上最多能同时有多少个线程块

⑥. cudaLaunchKernel：同上

⑦. cudaDeviceSynchronize：同上

将矩阵维度从 2048改为 128，结果如下：

可以观察到第四项执行了不同的指令，调用了不同的计算内核。在高抽象层，矩阵乘法被视为一个整体操作。但在底层实现时，根据矩阵维度以及硬件配置的差异，系统实际调用的矩阵乘法运算内核可能完全不同，这会导致相当大的性能差异。

Torch Compile 工具：内置了一个能够对硬件上的矩阵乘法性能进行微基准测试（micro benchmark），然后为模型选择性能最高的矩阵乘法子程序（subroutines）——得到 10% 左右的效率优化。（具体见下 3.4）

2.2.4 torch.cdist()

无论是加法还是乘法， CPU 和 GPU 之间为一对一的关系：一个 CPU 操作对应一个 GPU 操作。

torch.cdist 计算的是两组矩阵之间的欧式距离，即两组词向量的逐对距离度量。

cdist_function = lambda a, b: torch.cdist(a,b)
cdist_profile = profile("matmul", run_operation2(dim=2048, operation=cdist_function))
print(cdist_profile)

2.2.5 gelu()

gelu_function = lambda a, b: torch.nn.functional.gelu(a+b)
gelu_profile = profile("matmul", run_operation2(dim=2048, operation=gelu_function))
print(gelu_profile)

2.2.6 softmax()

softmax_function = lambda a, b: torch.nn.functional.softmax(a+b, dim=-1)
softmax_profile = profile("matmul", run_operation2(dim=2048, operation=softmax_function))
print(softmax_profile)

2.2.7 MLP

if torch.cuda.is_available():
    mlp_profile = profile("mlp", run_mlp(dim=2048, num_layers=64, batch_size=1024, num_steps=2), with_stack=True)
else:
    mlp_profile = profile("mlp", run_mlp(dim=128, num_layers=16, batch_size=128, num_steps=2), with_stack=True)
    
print(mlp_profile)

2.3 NVIDIA —— Nsight System

粗略看了一下这东西不简单，工程实践上的东西，先略过

没自己用过听老哥讲根本听不懂

在代码运行过程中， CPU 进度要远远快于 GPU 。例如在迭代过程中打印损失值 loss，会影响 CPU 与 GPU 的运行状态。由于打印操作是发生在 CPU 上的，所以 CPU 必须等待 GPU 计算成损失结果，才能继续往下进行。故这种情况下，两者进度同步， CPU 有大量的空转时间。

3. CUDA Kernels

3.1 pytorch & manual（gelu）

Pytorch 内部的 GeLu 实现方式如下：

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

x = torch.tensor([1.])
y1 = pytorch_gelu(x)

利用 tanh 来近似计算 GeLu 以加快计算速度，没有使用 GeLu 的精确定义用标准高斯分布的累计分布函数 CDF 来计算。

原始方法计算：

def manual_gelu(x: torch.Tensor):
    return 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x * x * x)))

y2 = manual_gelu(x)

计算两者结果，并分别进行 benchmark 以及 profiling 。 

def run_operation1(dim:int,operation:Callable)->Callable:
    x = torch.randn(dim, dim, device=device)
    return lambda:operation(x)

# benchmark
pytorch_time = benchmark("pytorch_gelu", run_operation1(dim=16384, operation=pytorch_gelu))
manual_time = benchmark("manual_gelu",run_operation1(dim=16384, operation=manual_gelu))

# profiling
pytorch_table = profile("pytorch_gelu",run_operation1(dim=16384, operation=pytorch_gelu))
manual_table = profile("manual_gelu",run_operation1(dim=16384, operation=manual_gelu))

pytorch_table：

manual_table：

manual_gelu 执行了大量运算，触发了多个 CUDA kernel（没有实现 Computation Fesion，数据被搬运太多次）。 而 pytorch_gelu 只采用了一个 kernel 就完成了计算，所以出现了 8 倍的 benchmark 时间差异。

3.2 Write a Kernel （gelu）

CUDA is an extension of C/C++ with APIs for managing GPUs and a programming model for expressing parallelism.

语言层面： C/C++ 扩展

语言层面：提供操作 GPU 的 API

编程模型：实现并行

编写 CUDA kernel 函数并启动它时，它会自动在 GPU 的数千个线程上并行执行，对向量或矩阵的所有元素同时进行计算。在 CUDA 的编程模型中，Grid 是顶层容器，包含若干个 Block；每个 Block 包含若干个 Thread。例如在二维 Grid 中，每个 Block 通过 (blockIdx.x, blockIdx.y) 来定位；二维 Block 中，每个 Thread 通过 (threadIdx.x, threadIdx.y) 定位。（这些参数不是通过参数传给 Kernel 的，Kernel 可以直接访问 CUDA 提供的内置变量 threadIdx、blockIdx、blockDim、gridDim，通过这些变量，每个线程可以计算出自己在全局数据中的唯一索引，从而处理对应的数据元素）

调试 CUDA 时，需要将环境变量设置为 os.environ["CUDA_LAUNCH_BLOCKING"] = "1"  
只有这样才能正确调试 CUDA，不过这牺牲了运行时性能，系统能够返回详细的错误信息。

#include <math.h>
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>

__global__ void gelu_kernel(float* in, float* out, int num_elements) {
    // Get the index into the tensor
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    // 算出当前线程要处理张量中的哪个位置

    if (i < num_elements) {  // 防止超出张量范围（因为总线程数 ≥ 元素数）
        // Do the actual computation
        out[i] = 0.5 * in[i] * (1.0 + tanh(0.79788456 * (in[i] + 0.044715 * in[i] * in[i] * in[i])));
    }
}
// __global__ 不能有返回值，用 out 指针写回结果

inline unsigned int cdiv(unsigned int a, unsigned int b) {
    // inline: 建议编译器消除函数调用开销，直接展开代码
    // Compute ceil(a / b)
    return (a + b - 1) / b;    // -1 防止原本就能整除   (a  = 9, b = 3)
}

// Host函数：发生在CPU侧
torch::Tensor gelu(torch::Tensor x) {
    // 检查张量在 GPU 上且连续
    TORCH_CHECK(x.device().is_cuda());
    TORCH_CHECK(x.is_contiguous());

    // 在 GPU 上分配一块和 x 同样大小（形状）的内存，但不初始化
    torch::Tensor y = torch::empty_like(x);

    // Determine grid (elements divided into blocks)
    int num_elements = x.numel();   // 获取元素总数   例如 [3, 4] 的张量 → num_elements = 12
    int block_size = 1024;  // Number of threads in a block
    int num_blocks = cdiv(num_elements, block_size);   //用向上取整除法算出需要启动多少个 Block 才能覆盖所有元素

    // Launch the kernel
    gelu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), num_elements);
    // 在 GPU 上启动 num_blocks 个 Block，每个 Block 有 1024 个线程，总共 num_blocks × 1024 个线程并行执行
    // 由于 gelu_kernel 由核心修饰符 __global__ 修饰，说明 gelu_kernel 由 CPU 调用， GPU 个线程并行执行
    // gelu_kernel 需要 <<<>>> 传参
    // 所有被 __global__ 修饰的核函数都需要传这两个参数
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    //检查内核启动是否成功。如果内核里有 bug ，这里会立即捕获并报错

    return y;
}

（搞了一整天没编译成功 .cu 和 .cpp 文件，先搁置吧。心累，这种实践上的东西 AI 也是张口就来，解决不了）贴一个完整代码在下面，benchmark 和 profiling 也只有略过。

import torch
import os
from torch.utils.cpp_extension import load_inline


# 1. 确保目录存在
def ensure_directory_exists(path):
    os.makedirs(path, exist_ok=True)


# 2. CUDA 源代码（核函数实现）
cuda_gelu_src = """
#include <math.h>

__global__ void gelu_kernel(float* in, float* out, int num_elements) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < num_elements) {
        float x = in[i];
        out[i] = 0.5f * x * (1.0f + tanhf(0.79788456f * (x + 0.044715f * x * x * x)));
    }
}

// 这个函数被 C++ 封装代码调用
void launch_gelu_kernel(float* in, float* out, int num_elements) {
    int block_size = 1024;
    int num_blocks = (num_elements + block_size - 1) / block_size;
    gelu_kernel<<<num_blocks, block_size>>>(in, out, num_elements);
}
"""

# 3. C++ 封装代码（PyTorch 绑定）
cpp_gelu_src = """
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>

// 声明 CUDA 函数（在 cuda_sources 中实现）
void launch_gelu_kernel(float* in, float* out, int num_elements);

// 被 Python 调用的函数
torch::Tensor gelu(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be on GPU");
    TORCH_CHECK(x.dtype() == torch::kFloat32, "x must be float32");

    auto y = torch::empty_like(x);
    int num_elements = x.numel();

    launch_gelu_kernel(x.data_ptr<float>(), y.data_ptr<float>(), num_elements);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

# 4. 安全检查
if not torch.cuda.is_available():
    print("CUDA 不可用，跳过编译")
    exit(1)

# 5. 确保编译目录存在
ensure_directory_exists("var/cuda_gelu")

# 6. 使用 load_inline 编译
try:
    module = load_inline(
        name="inline_gelu",
        cuda_sources=[cuda_gelu_src],
        cpp_sources=[cpp_gelu_src],
        functions=["gelu"],  # 暴露给 Python 的函数名
        extra_cflags=["-O2"],
        extra_cuda_cflags=["-O2"],
        verbose=True,
        build_directory="var/cuda_gelu",
    )
    print("✅ 编译成功！")

    # 7. 测试
    x = torch.randn(10, device='cuda', dtype=torch.float32)
    y = module.gelu(x)  # 调用编译好的函数
    print(f"输入: {x}")
    print(f"输出: {y}")

except Exception as e:
    print(f"❌ 编译失败: {e}")

1. pip install ninja； 

2.解决报错：subprocess.CalledProcessError: Command '['where', 'cl']' returned non-zero exit status 1      Pytorch 找不到 Visual Studio 的 C++ 编译器 cl.exe  解决方法链接

3.解决报错：subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 2.     将 torch.utils.pp_extension 中的

command = ['ninja', '-v']  ----->     command = ['ninja', '--version']

4.：新问题：搞不定了，摆烂了

文件也成功生成了

3.3 Triton Kernels （gelu）

Triton 的优点：可以用纯 Python 写 GPU Kernel，无需关心线程管理，只需专注于线程块的设计；能自动处理许多底层细节，比如自动调整内存访问模式。代码是以 Tile 为中心写的，编译器负责将这些 Tile 调度到 SM 上。跨 SM 的数据共享需要手动处理，跨 SM 的并行 Triton 能够自动实现。

首先需要安装 triton（由于自己是 Windows 环境，需要进行一些设置）

import torch
import triton
import triton.language as tl


# 原本用 C++ 实现的 CPU 侧的 Host 函数
# torch::Tensor gelu(torch::Tensor x)
def triton_gelu(x: torch.Tensor):
    assert x.is_cuda
    assert x.is_contiguous()

    # Allocate output tensor
    y = torch.empty_like(x)

    # Determine grid (elements divided into blocks)
    num_elements = x.numel()
    block_size = 1024  # Number of threads
    num_blocks = triton.cdiv(num_elements, block_size)

    # 这一步的实现与 CUDA 稍有区别
    # gelu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), num_elements);
    triton_gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=block_size)
    # Triton 要求 grid 参数是一个元组，用来支持多维网格
    # (M,) 一维：启动 M 个 block，索引为 0, 1, 2, ..., M-1
    # (M,N) 二维：启动 M*N 个 block，索引为 (0,0), (0,1), ..., (M-1, N-1)
    # (M,N,K) 三维

    return y


@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    # BLOCK_SIZE: tl.constexpr：编译时常量（在编译时确定，不能运行时改变）

    pid = tl.program_id(axis=0)    # 可以理解成 blockId
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # 计算索引
    # CUDA 需要计算出每个 thread 的索引
    # Triton kernel 计算的offsets 是一个向量，长度为 block_size

    # Handle boundary
    mask = offsets < num_elements

    # Read
    x = tl.load(x_ptr + offsets, mask=mask)

    # Approx gelu is 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
    # Compute (tl.tanh doesn't exist, use tanh(a) = (exp(2a) - 1) / (exp(2a) + 1)
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    exp = tl.exp(2 * a)
    tanh = (exp - 1) / (exp + 1)
    y = 0.5 * x * (1 + tanh)

    # Store
    tl.store(y_ptr + offsets, y, mask=mask)

    # 类似于 SIMT，许多 threads 同时拿到这个命令，然后计算不同的数据

Triton 编译后得到 PTX，再经过一步变成 GPU 最终执行的机器码。

代码还是跑不通，没招了

3.4 Torch Compile（gelu）

能将未优化的 PyTorch 代码自动转换为更高效的版本，它自动尝试进行内核融合等优化操作。

compiled_gelu = torch.compile(manual_gelu)

compiled_time = benchmark("compiled_gelu", run_operation1(dim=16384, operation=compiled_gelu))
compiled_table = profile("compiled_gelu", run_operation1(dim=16384, operation=compiled_gelu))

（还是跑不了，这章听得我难受，配环境变量搞了一整天，啥也没成，还不知道原因）

3.5 Triton Kernel (softmax)

def triton_softmax(x: torch.Tensor):
    # Allocate output tensor
    y = torch.empty_like(x)

    # Determine grid
    M, N = x.shape                          # Number of rows x number of columns
    block_size = triton.next_power_of_2(N)  # 将 N 向上取整到 2 的幂
    num_blocks = M                          # Each block is a row 每行独立计算 Softmax，一行对应一个 block

    # Launch kernel
    triton_softmax_kernel[(M,)](
        x_ptr=x, y_ptr=y,
        x_row_stride=x.stride(0), y_row_stride=y.stride(0),
        num_cols=N, BLOCK_SIZE=block_size
    )
    # 传递了行跨度 (stride) 以便在内存中正确跳转到不同行

    return y


@triton.jit
def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride, num_cols, BLOCK_SIZE: tl.constexpr):
    assert num_cols <= BLOCK_SIZE

    # Process each row independently
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)

    # Read from global memory
    x_start_ptr = x_ptr + row_idx * x_row_stride
    x_ptrs = x_start_ptr + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))
    # other=float("-inf")：被 mask 掉的元素用负无穷填充，这样在后续 max 操作中不会影响结果

    # Compute
    x_row = x_row - tl.max(x_row, axis=0)
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator
    # Block 内部会有一个比较 max 以及求和的操作
    # 这两步 triton 会自行处理
    # 虽然单个线程看起来需要等待其他线程的结果，但 Triton 在 Block 层面实现了同步

    # Write back to global memory
    y_start_ptr = y_ptr + row_idx * y_row_stride
    y_ptrs = y_start_ptr + col_offsets
    tl.store(y_ptrs, y_row, mask=col_offsets < num_cols)