@maxxfuu: 推理工程第6/90天 我写了一个用于一维卷积的CUDA内核，只是为了练习写未优化的……

推理工程第6/90天

我写了一个一维卷积的CUDA内核，只是为了练习编写未经优化、平淡但正确的CUDA代码！我还阅读了PagedAttention论文（2312.07104），观看了一个关于GPU内存的短视频讲座，并且作为最高组，深蹲了455磅×6次和495磅×3次！

以下是我对PagedAttention的理解：

PagedAttention减少了内部碎片，并解决了外部碎片。

在解码阶段，KV缓存管理器实际上一次只分配一个缓存块来存储输出token。如果输出token超出了缓存块的容量，KV缓存管理器会分配一个新的缓存块来容纳之前缓存块装不下的剩余token。最终的效果是VRAM被有效利用，每个缓存块都被填满后才在VRAM上分配新的缓存块。

在朴素的KV缓存实现中，内存会预先为最大序列长度预留，而这些内存可能永远不会被用到，意味着有大量VRAM未被触及。

PagedAttention通过确保每个解码序列之间的内存永远不会被浪费，从而减少了内部碎片。

为了更高效地使用内存，对于共享相同前缀token的不同提示，KV缓存管理器允许它们共享同一个缓存块，但从两个提示中token开始不同的位置起，写入到一个新的缓存块。

尽管每个解码阶段实际上会写入VRAM的随机位置，但块表提供了一种抽象，使KV缓存对模型而言看起来像是连续的内存。这使得分配过程变得简单，只需检查缓存块是否已满。

这种持续使用缓存块来分配恰好足够的内存，并使用块表来分配VRAM上非连续的、分散的空闲内存的做法，减少了内部碎片，解决了外部碎片。

观看了一个关于GPU内存的视频：https://youtube.com/watch?v=Zrbw0zajhJM&list=PL5XwKDZZlwaY7t0M5OLprpkJUIrF8Lc9j&index=7…

我阅读了这篇博客，它完全巩固了所有知识，由（@hamzaelshafie）撰写：https://hamzaelshafie.bearblog.dev/paged-attention-from-first-principles-a-view-inside-vllm/…

TL;DR: 本集简要介绍CUDA中的五种内存类型——全局、寄存器、本地、常量、共享——及其片内/片外位置、读写权限、作用域和性能特点,为后续深入优化做铺垫。

GPU内存层次概述

GPU芯片(片上)通过PCB连接到VRAM(片外),但部分内存直接集成在芯片内部,访问速度更快。理解片内与片外内存的差异是编写高性能CUDA代码的基础。

全局内存

• 位置:片外(GPU的VRAM)
• 权限:所有线程均可读写
• 分配方式:cudaMalloc或创建全局变量
• 特点:容量最大,速度最慢
• 作用域:全局可见,生命周期由主机控制

寄存器

• 位置:片内
• 权限:每个线程私有
• 分配方式:内核中声明的局部变量默认存于寄存器
• 特点:速度极快,数量有限
• 占用率影响:使用过多寄存器可能导致占用率下降(后续集数详解)
• 溢出问题:当寄存器不足时,变量会被自动溢出到本地内存(片外,速度慢),应避免。可通过编译标志(如增加PTX详细输出)或cuobjdump查看寄存器及本地内存使用情况。

示例:使用255个寄存器后,变量开始溢出到本地内存,导致2040字节的本地内存读写操作。

本地内存

• 位置:片外(名称易混淆,仅因对线程私有而叫“本地”)
• 权限:线程私有
• 特点:访问速度慢,是寄存器溢出的结果,应尽量避免
• 检测:通过编译标志可查看本地内存访问量

常量内存

• 位置:片外(但配有缓存,只读)
• 容量:仅64KB
• 权限:所有线程只读
• 访问特性:线程束内不同线程访问不同地址时会串行化;若多个线程访问同一地址,性能可能优于全局内存
• 声明方式:__constant__关键字,用cudaMemcpyToSymbol从CPU复制数据

共享内存

• 位置:片内
• 权限:同一线程块内的所有线程可读写
• 分配方式:内核中数组声明加__shared__关键字
• 特点:访问速度远快于全局内存,适合多线程频繁访问同一地址的场景

五种内存类型总结

后续集数将深入每种内存的优化技巧。欢迎订阅以避免错过更新,点赞、评论、分享你的反馈。

Source: https://www.youtube.com/watch?v=Zrbw0zajhJM