简介:一份面向Armv8/Armv9底层开发者的《深度学习cache系列》PDF文档,系统讲解高速缓存工作原理与实际工程应用。内容从“为什么要用cache”切入,依次介绍L1/L2/L3多级缓存结构、索引/路/集合的组织形式,以及VIVT、PIPT、VIPT等缓存种类、分配与替换策略;在此基础上,结合MESI协议和CCI、CMN、DSU等总线/接口机制,说明多核与多cluster场景下缓存一致性的维护方法,并梳理CLIDR_EL1、CTR_EL0、CCSIDR_EL1等系统寄存器的用途,给出软件使用flush/invalidate指令维护一致性的示例。同时讨论不可缓存内存区域、MMU关闭时的缓存行为,以及页表属性对缓存策略的影响。资源为单个PDF文件,大小约5.71MB,已有351人学习。作者结合ARM官方手册与一线支持经验整理,配有架构图、查询流程、动图和常见问题思考,适合嵌入式、内核或安全固件工程师系统梳理cache知识,也便于在实际项目中排查内存属性配置与性能瓶颈问题。
1. 深度学习里那些躲不开的 cache:从数据管道到 KV Cache,都在烧你的时间与显存
做过几年深度学习训练的人都有这种体感:明明显存没涨、算力没掉,可 GPU 利用率就是上不去,训练时间不稳定,推理首字延迟突然多出几百毫秒。查到最后,八成不是模型代码的问题,而是 cache 在背后捣乱。cache 不是一个单一东西,在深度学习里它至少横跨三层:数据读取时的文件系统页缓存、GPU 显存分配器的缓存复用、以及大模型推理里用来存历史 Key/Value 的 KV Cache。这三层任何一层没调好,都会让模型空转。这个系列想做的就是把这层窗户纸捅开:把 cache 的机制讲透,把参数给到能直接用,把你踩过和没踩过的坑一并排掉。适合正在折腾训练提速、推理显存优化、或者刚把模型搬到新环境的人。
2. 数据管道的 cache:为什么 Worker 加满,训练还是慢
2.1 先搞清三层缓存:磁盘页缓存、进程内预读、数据张量缓存
训练样本在进入 GPU 之前,要经过磁盘 → 内存 → 进程 → 副本 → GPU 这条链路。每一步都有 "缓存" 参与,但它们作用完全不同。
第一层是操作系统页缓存(Page Cache)。Linux 读文件时,文件内容会被读进空闲内存,下次再读同一文件直接命中内存,不必碰磁盘。这一层对训练集这种高频重复读取的文件特别友好:如果你的机器内存大,整个数据集可以被 OS "吞" 进页缓存,后面每个 epoch 读起来都快如闪电。
第二层是 DataLoader 的预读缓冲。PyTorch 的 DataLoader 里有两个参数直接控制这部分:num_workers决定有几个子进程去读和预处理数据,prefetch_factor决定每个 worker 提前预取多少批数据。这两个值不只是“多开几个线程”那么简单。num_workers太小,GPU 会等数据;太大,子进程间 IPC 和设备内存拷贝开销会反噬。我一般用的经验值:本地 NVMe 磁盘时num_workers=4到8,网络文件系统时提高到8到16,同时prefetch_factor=2或4。
第三层是 GPUDirect / 固定内存(Pinned Memory)的副本缓存。DataLoader 默认先拷到普通内存,再从普通内存拷到 GPU 页表管理的固定内存,最后才拷贝到显存。pin_memory=True会省掉普通内存到固定内存这一步拷贝,但代价是固定内存不能被 OS 交换,占用的是不可回收的 RAM。如果你的内存吃紧,千万别开。
2.2 最小可复现的数据管道调优脚本
下面这段是训练脚本开头的标准配置,能同时调节上面三层。
from torch.utils.data import DataLoader, Dataset class SimpleDataset(Dataset): def __init__(self, n=100000): self.data = list(range(n)) def __len__(self): return len(self.data) def __getitem__(self, idx): # 这里模拟读取图片/文本并做预处理 return self.data[idx] dataset = SimpleDataset() loader = DataLoader( dataset, batch_size=128, shuffle=True, num_workers=8, # 子进程数,按 CPU 核数和磁盘类型调整 prefetch_factor=4, # 每个 worker 预取的 batch 数 pin_memory=True, # 用固定内存加速到 GPU 的拷贝 persistent_workers=True # 避免每个 epoch 重新拉起 worker )num_workers=8配合prefetch_factor=4,等于每个 worker 在内存里提前准备 4 个 batch,绝大多数时候 GPU 不用干等。persistent_workers=True是很多新手容易漏的:如果为 False,每个 epoch 结束后 worker 会退出销毁,下个 epoch 再重新创建,白白浪费几十秒。注意prefetch_factor在num_workers=0时无效,而且 PyTorch 1.13 之后要求num_workers>0才能同时设persistent_workers=True,否则直接报错。
2.3 实战中怎么确认数据加载真的命中了 cache
跑训练时开着htop看 CPU 负载、用iostat看磁盘%util,能粗判数据管道是否卡顿。但如果想确认 OS 页缓存是否覆盖了数据文件,Linux 下用fincore或读/proc里的mincore信息比较直接。一般我用的命令是vmtouch这种小工具扫一遍数据集目录。
# 检查数据集目录被页缓存命中的比例 vmtouch /data/train/images/ # 输出里会显示:Pages: 45123 / 54321 (83.1%)如果命中率低,比如每次 epoch 都要重新读磁盘,就要看内存是否被其他进程挤占了。常见做法是在训练前置入数据:
# 把数据集读入页缓存,避免训练过程中反复读盘 cat /data/train/images/*.jpg > /dev/null这只对首次冷启动有效。训练期间如果有别的进程把内存吃掉,页缓存会自动回收,命中率又会掉下去。另一个更可控的方案是用 LMDB 或 TFRecord 这类把数据打包成单文件,再配合mmap模式读取——文件系统对这单个大文件的缓存命中率要远高于成千上万个零碎小文件。
提示:数据管道的调优,永远先看 GPU 利用率,再决定要不要动 cache。如果 GPU 利用率已经 95% 以上,再去堆 prefetch 没有任何收益,只会让 CPU 更忙。
3. GPU 显存里的 cache 玄学:Caching Allocator 与空显存陷阱
3.1 PyTorch 的 Caching Allocator 是怎么“假报”显存的
用nvidia-smi看显存占用时会发现:程序只占 3G,可显存却显示用了 9G;或者反过来,Python 进程快退场了显存还没释放。这是 PyTorch 显存分配器的缓存策略在起作用。PyTorch 内部用 Caching Allocator 管理显存:当 Tensor 被释放时,显存块不会立刻还给驱动,而是留在一个缓存池里供后续分配复用。这么做的原因是 CUDA 的cudaMalloc/cudaFree是慢操作,频繁调用会让训练速度下降一个数量级。缓存池本质上是一个“二次利用”的机制。
但这个缓存池有两个副作用。第一,显存占用看起来居高不下。第二,当训练中出现新的显存需求,比如torch.cuda.empty_cache()只释放空闲块,如果缓存池里的块全部被占用(哪怕里面存的是已经释放但尚未回收的 Tensor),OOM 照样发生。很多人的“翻车”现场是这么来的:加了empty_cache()以为清干净了,结果每个 step 都慢一倍,而且 OOM 一般会在估值计算或梯度累积时才爆炸。
3.2 用一组参数控制显存缓存块的切分:PYTORCH_CUDA_ALLOC_CONF
显存缓存池的行为在 PyTorch 1.10 之后可以用环境变量PYTORCH_CUDA_ALLOC_CONF调节。最常用的两个键是max_split_size_mb和expandable_segments:True。前者控制缓存块的最大切分粒度,后者让显存段按需扩展。
# round_robin 让多卡分配更均匀,max_split_size_mb 控制大块缓存切割 export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,round_robin:true"这个配置解决的是“显存碎片化”。当你的模型同时有大量大小不均的 Tensor 时,默认的分配策略会形成很多不可用的小空洞。把max_split_size_mb调成你模型中最大 Tensor 的稍大值,比如最大中间张量是 200MB,就设 256,能显著减少碎片。round_robin:true是对多卡场景下让缓存池在多个设备上轮流分配,避免某张卡缓存过多、另一张卡却 OOM。注意expandable_segments:True是和 CUDA Graph 搭配的;如果模型里有动态 shape,开它反而容易出错。
3.3 用一个可见的最小脚本观察缓存池的“假占用”
下面这段脚本直观展示了缓存池的分配和回收行为。
import torch def show_mem(name): allocated = torch.cuda.memory_allocated() / 1024**2 cached = torch.cuda.memory_reserved() / 1024**2 print(f"{name}: allocated={allocated:.1f}MB cached={cached:.1f}MB") show_mem("init") a = torch.randn(1024, 1024, device="cuda") # 分配 4MB 显存 show_mem("after alloc") del a # 释放张量 show_mem("after del") torch.cuda.empty_cache() # 强制将空闲块还给出驱动 show_mem("after empty_cache")del a之后allocated会立刻下降,但cached仍然维持原高位;直到empty_cache()才真正把显存归还。所以如果你的代码在第 n 个 step 报 OOM,但 n 之前显存一直平稳,大概率是缓存池里的空闲块不够连续,而不是真的显存用满了。此时不要急着调批量大小,先检查是不是某个临时 Tensor 在循环里不断改变 shape 导致缓存块反复拆分。
3.4 分布式训练里的缓存“黑匣子”:NCCL 固定缓冲区
用torch.distributed跑多卡时,显存里还有一块你看不见的缓存:NCCL 的通信缓冲区。PyTorch 会在初始化torch.distributed.init_process_group时给每个通信算子预留一块固定大小的显存,默认是NCCL_BUFFSIZE=4194304字节(4MB)。你在nvidia-smi里看到每张卡多出的几百 MB,往往就是它。
更隐蔽的是 AllReduce 梯度时如果模型太大,设置NCCL_BUFFSIZE过小会导致通信性能断崖。我调过的经验值是 256MB 对 8 卡 BERT Base 足够;更大模型建议直接用环境变量NCCL_BUFFSIZE=268435456显式设置,别信默认值。这个东西一旦设错,现象极其诡异:速度忽快忽慢,带宽跑不满,换卡也没用。检查命令是export NCCL_DEBUG=INFO看通信日志里的 buffer 使用情况。
提示:显存缓存池是给“单进程内反复分配释放”设计的。如果你的服务是常驻推理进程,显存一直不降是正常现象,不要用empty_cache()强行回收,那会让下一次推理多花几十毫秒重新cudaMalloc。
4. KV Cache:大模型推理里最贵的一笔显存生意
4.1 为什么每推理一个 token 都要重新计算前面所有人的 Key/Value
自回归生成模型(GPT、LLaMA、Qwen 这一族)在解码时,每生成一个 token,都要拿当前 token 的 Query 去和前面所有 token 的 Key、Value 做注意力计算。如果不做任何缓存,第 100 个 token 需要重新算第 1~99 个 token 的 K、V,成本按序列长度平方增长。KV Cache 的方案是:每生成一个 token,就把它的 Key 和 Value 追加到缓存里;后续 token 的注意力计算直接复用缓存中的 K、V,只需要算新 token 的 K、V。这让生成阶段的计算量从平方降到线性。
但代价是显存。KV Cache 的显存占用与 batch size、序列长度、层数、注意力头数、head 维度、精度直接挂钩。计算公式是:
KV Cache 字节数 = batch_size × 生成长度 × num_layers × 2 × num_key_value_heads × head_dim × bytes_per_element
注意这里的系数是 2,因为要同时存 Key 和 Value 两份。很多没有接触过多头潜在注意力的新手容易按 num_attention_heads 算,结果翻好几倍。现在的模型如果用 GQA(分组查询注意力)或 MQA(多查询注意力),KV Cache 里的 head 数要按num_key_value_heads算,不是 query 的 head 数。
4.2 用一段 Python 代码估算你的推理该留多少显存
写一个小函数在部署前估算 KV Cache 上限,比跑到 OOM 再回退稳妥得多。
def kv_cache_bytes(batch_size, max_gen_len, num_layers, num_kv_heads, head_dim, dtype_bytes=2): """ 估算 KV Cache 最大占用字节数 dtype_bytes: FP16/FP8 为2或1,FP32为4 """ per_token_bytes = (2 * num_layers * num_kv_heads * head_dim * dtype_bytes) return batch_size * max_gen_len * per_token_bytes # 以 7B 规模、GQA 的模型为例:32层, 4个KV头, 头维度128, FP16 bytes_ = kv_cache_bytes(batch_size=8, max_gen_len=2048, num_layers=32, num_kv_heads=4, head_dim=128) print(f"KV Cache 需求: {bytes_ / 1024**3:.2f} GB")这算出来的是生成到 2048 个 token 时的峰值。如果只有 24GB 显存,模型权重已经吃了 14GB,那 KV Cache 最多只剩 10GB,上面例子的 8 并发明显超了。实际调参时要同时考虑“权重显存 + 激活显存 + KV Cache 上限”三块,我给的行内经验是按峰值预留 20% 余量。代码里dtype_bytes用 2 对应 FP16/BF16,用 1 对应 FP8 或 INT8 量化。别小看这个预设,很多推理框架只做到 FP16,你写成 INT8 会高估容量。
4.3 两个让 KV Cache 物尽其用的常驻技巧:PagedAttention 和 reuse
PagedAttention 是 vLLM 的核心思想:把 KV Cache 切成固定大小的块,用类似操作系统虚拟内存分页的机制按需分配,避免每请求都用“最大可能的显存”去预留连续空间。这句话落地到使用上只有一个动作——别自己手写 KV Cache 管理,直接用 vLLM / SGLang 等推理框架。它们内部已经做了 block 分配和换出。
第二个技巧是前缀复用(Prefix Caching)。当多个请求拥有相同的前缀 prompt 时,比如系统提示词和几十轮历史消息,这部分 KV Cache 不需要重复计算。DeepSeek 的 DeepSeek-V2 提出的 MLA(Multi-head Latent Attention)更是把 KV Cache 大幅压缩,本质是缓存低秩投影后的隐向量,而不是每层的完整 K、V。如果你在自己写推理服务,可以在内存里维护一个“已计算前缀 hash → KV Cache 索引”的字典,命中直接拼接,省掉首字等待。这在多轮对话里收益尤其明显。
4.4 什么情况下别开 KV Cache
KV Cache 不是没有边界。当你的业务场景是“一次问完,不续写多轮”,并且生成长度很短(比如十几 token 的标签分类),KV Cache 的显存开销可能比加速收益更扎眼。更极端的情况是流式场景里用户中断了生成,但已经产生的 KV Cache 不丢掉的话,会让下一个完全不相关的请求误以为共享前缀,产生语义污染。
所以推理服务里要留一个开关:如果模型只需要单轮短应答,就设置max_kv_cache_len=0强制关闭缓存;多轮对话场景才开启,并且要按用户会话维度做隔离。我踩过这个坑:客服机器人会话切换时没清缓存,下个用户问了“今天天气”,模型答成了“上一单退款进度”。
提示:KV Cache 的显存不是静态的,它会随生成长度线性增长。长对话出现 OOM 不一定是权重太大,更多是 KV Cache 撑满了。线上服务要对max_generated_tokens做硬限制,否则聊到 20 轮以后必定翻车。
5. 深度学习 cache 连环坑:现象、原因、解决一条龙排查
5.1 坑一:torch.cuda.empty_cache()变成每步必调的“后悔药”
现象:训练时显存看起来一直满,代码里每次迭代后调用empty_cache(),结果训练速度掉了一半以上,GPU 利用率从 90% 滑到 50%。
原因:empty_cache()把显存缓存池里所有空闲块全部归还给 CUDA 驱动,下一次迭代重新分配时每次都要走cudaMalloc,这个系统调用既是同步的又是慢的,把缓存池的复用价值彻底废掉。
解决:删掉循环里的empty_cache(),只在 OOM 异常处理器里为“降 batch 后重试”这种场景使用。如果是推理服务,一次都别调。改用前面说的PYTORCH_CUDA_ALLOC_CONF去控制碎片,而不是清池子。
5.2 坑二:DataLoader 开了num_workers=32结果内存被吃光
现象:机器 64GB 内存,数据集是 200GB 的高清图片,num_workers调大后内存直接触顶,训练反而变慢,甚至 OOM 被杀。
原因:每个 worker 都会把当前解码的图像和增强后的副本放在自己进程的堆里。prefetch_factor=4时,32 个 worker × 4 个 batch × 每个 batch 100 张图,数百 GB 的内存需求根本扛不住。这个“多开缓存”的直觉在数据管道里恰恰是最常见的翻车点。
解决:先看单个 batch 的内存占用:batch_size × 单样本解码后字节数 × prefetch_factor × num_workers。把num_workers降到 CPU 核数的一半以内,同时改用解码后更紧凑的格式(JPEG 解码成 uint8 再转 tensor 而不是保留 PNG 原始)。如果数据很大,考虑用 LMDB 或 WebDataset 流式读取,避免整目录缓存。
5.3 坑三:pin_memory=True导致 CPU 内存耗尽
现象:模型不大,显存没满,但系统内存持续飙升,最终触发Cannot allocate memory,训练进程直接退出。
原因:固定内存是锁页内存,不可被换出,也不能被页缓存回收。当 DataLoader 的pin_memory=True配合大num_workers时,每个 worker 产生的 batch 都要在固定内存里排队等拷贝,内存占用远超普通内存。
解决:pin_memory=True只在你确认内存充足时开。检查方法是看/proc/meminfo的MemAvailable和SwapTotal;如果内存只有模型显存的 2 倍多,就别开。更细的控制是让 DataLoader 的batch_sampler保持固定 batch shape,避免固定内存反复分配大块,allocator的碎块问题也会缓解。
5.4 坑四:KV Cache 量化后效果崩掉
现象:用 INT8 量化 KV Cache 后显存减半,但长文本生成质量明显变差,开始重复、答非所问。
原因:KV Cache 的数值分布在不同 layer 差异极大,无脑把全部 Cache 量化到 INT8,某些 head 的敏感位置误差被放大,注意力分数漂移,代表性 token 被错误匹配。
解决:KV Cache 量化要做带精度感知的混合量化。NVIDIA 的 KV Cache quantization 方案里,某些层保留 FP16,只量化敏感度低的层。如果你不想碰这块,直接用 FP8 也是折中,因为 FP8 的指数位比 INT8 多,动态范围更接近 FP16。量化前至少用torch.profiler打点看每层 Cache 的数值分布,而不是闭着眼睛转格式。
5.5 坑五:系统页缓存被训练日志和 checkpoints 冲掉
现象:数据读取速度一直稳定,突然到某个 epoch 开始飙高,IO 等待时间上升。
原因:每轮保存checkpoint.pt时,写入几 GB 文件到磁盘,write-back cache(写回缓存)会迅速占满空闲页缓存,把原本属于数据集的热点页缓存挤出去。再次读数据集就变成冷加载。
解决:把 checkpoint 写到独立磁盘分区,或用posix_fadvise对数据集文件设置POSIX_FADV_DONTNEED标识避免被其他读操作干扰。简单粗暴的做法是:训练开始时用vmtouch把数据集锁进页面,并在保存 checkpoint 时把写入缓冲区直接回写后 drop,用sync && sysctl -w vm.drop_caches=3手动释放写缓存(仅限单机测试环境,生产环境谨慎)。
6. 进阶:用 Profiler 把 cache 命中率和显存碎片打回原形
到了这个阶段,你已经知道 cache 不是玄学,而是可以用工具量化的指标。我推荐你养成一个习惯:调优前先torch.profiler打点 100 步,把 CPU 到 GPU 的每个阶段拆开看,而不是凭感觉改参数。下面是一个标准用法。
from torch.profiler import profile, ProfilerActivity, record_function with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True, profile_memory=True) as prof: for step, data in loader: with record_function("train_step"): # 前向、反向、优化器更新 pass if step >= 50: break print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))profile_memory=True会给出每步的显存分配和释放信息;sort_by="cuda_time_total"能直接看到最耗时的算子是不是涉及数据拷贝。如果copy_类算子排第一,说明数据管道的缓存没对接好,优先调 DataLoader;如果是cudaMemsetAsync频繁出现,多半是显存缓存池在反复扩缩,去调PYTORCH_CUDA_ALLOC_CONF。
再看缓存碎片率,命令行工具nvidia-smi -q -d MEMORY会显示每个 GPU 显存段的Total FB Free和其他字段。如果空闲显存被分割成大量碎片,说明分配器策略不适合当前模型的 tensor 尺寸分布。此时我会用torch.cuda.memory_snapshot()导出一个 JSON,分析每个 segment 的大小和用途,找出谁在制造大洞。
prof = torch.cuda.memory_snapshot() # 每个 segment 结构里的 "size" 与 "allocated_size" 可以算碎片率 for seg in prof: if seg["allocated_size"] / seg["size"] < 0.7: print(f"segment {seg['addr']} 有 {seg['size'] - seg['allocated_size']} 字节空洞")碎片率高于 30% 时,优先考虑max_split_size_mb调大,其次是开启expandable_segments:True(只兼容 CUDA Graph 场景)。还有一个更冷门的招:把torch.cuda.set_per_process_memory_fraction(0.95)放宽到 0.98,让分配器有更多连续空间可用,但只建议在显存大于模型需求 1.3 倍时用。
我自己的兜底习惯是:每次换新机器、换新 PyTorch 版本、换新 CUDA 驱动,都会跑一遍上面这段 profiler 脚本,把 cache 行为当成上传数据一样存进 notes。因为 PyTorch 版本升级后缓存分配策略会改,驱动更新也可能改变显存管理,这些都不写进 release notes,只能自己用打点看出区别。希望你也能把这套“量化 cache、再动参数”的流程沉淀到自己的部署环境里,至少能让线上服务的显存和延迟表现稳定下来。希望帮到你。
本文还有配套的精品资源,点击获取