Hotdry.

Article

tiny-vLLM 的 CUDA 内存池分层与 KV 缓存碎片整理机制

深入解析 tiny-vLLM 中 PagedAttention 的分页显存管理策略、CUDA Memory Pool 的分层设计，以及 KV 缓存碎片预防与整理的具体实现参数。

2026-05-30ai-systems

在 LLM 推理服务中，显存管理是制约吞吐量和延迟的核心瓶颈。传统的连续内存分配策略在应对变长序列时会产生严重的外部碎片，导致显存利用率低下。tiny-vLLM 作为 vLLM 的精简教学实现，通过引入 PagedAttention 分页机制和 CUDA Memory Pool 分层分配策略，实现了细粒度的显存复用与接近零拷贝的推理路径。

PagedAttention 的分页显存架构

tiny-vLLM 采用 PagedAttention 机制管理 KV 缓存，其核心思想借鉴操作系统虚拟内存的分页设计。KV 缓存被划分为固定大小的块（block），每个块存储固定数量的 token 的 Key 和 Value 向量。与连续分配不同，序列的逻辑位置通过 block table 映射到物理块，允许非连续存储。

这种设计的优势在于彻底消除了外部碎片。传统连续分配中，不同长度的序列会留下无法被利用的内存间隙；而分页机制下，只要存在空闲块即可满足分配请求，无论这些块在物理上是否连续。内部碎片被限制在每个序列最后一个未填满的块中，当块大小设置合理（如 16-64 tokens）时，这部分开销通常低于 5%。

在 tiny-vLLM 的实现中，block table 维护每个序列的逻辑块到物理块的映射关系。当序列生成新 token 时，引擎从全局空闲块池中分配新块；当序列完成或被抢占时，其占用的块立即归还到空闲池供复用。这种动态分配策略避免了为每个序列预留最大长度的显存空间，显著提升了显存利用率。

CUDA Memory Pool 的分层分配策略

tiny-vLLM 在 CUDA 层面采用分层内存池设计，通过 cudaMallocAsync 和 cudaMemPool API 实现高效的显存管理。与传统 cudaMalloc/cudaFree 相比，内存池通过预分配大块后备内存并从中服务小块请求，将分配延迟从微秒级降至纳秒级，同时减少设备端碎片。

内存池的分层结构包含三个层级：

设备全局内存池：管理 HBM/VRAM 中的大块显存，通过 cudaMemPoolCreate 创建。池的大小应根据模型规模和并发请求数预设，通常设置为可用显存的 70-80%，预留空间给权重加载和临时缓冲区。

流序分配层：利用 cudaMallocFromPoolAsync 实现流序分配语义，允许跨流的内存复用。当一条 CUDA 流释放内存后，同一设备上其他流在满足同步依赖的前提下可立即复用该内存，无需等待设备级同步。

共享内存层：在 kernel 执行期间使用 __shared__ 修饰符分配的 SRAM，用于 RMSNorm 的并行归约、Softmax 的数值稳定计算等场景。tiny-vLLM 的 RMSNorm kernel 中，共享内存用于存储中间归约结果，避免全局内存的频繁访问。

内存池的关键参数包括 cudaMemPoolAttrReleaseThreshold，控制何时将内存归还操作系统。在推理场景中，建议将该阈值设为较高值（如 0.9），优先在池内复用内存而非频繁与 OS 交互，降低分配延迟。

KV 缓存碎片整理机制

PagedAttention 的碎片管理策略以预防为主，而非传统的碎片整理（defragmentation）。由于物理块无需连续，系统不需要执行内存压缩操作来合并碎片。当显存压力升高时，tiny-vLLM 采用抢占（preemption）和块交换（swapping）策略：将低优先级请求的 KV 块从显存换出到主机内存，或选择重新计算其 KV 缓存，而非尝试整理碎片。

这种设计的工程合理性在于：GPU 内存的随机访问特性使得原地整理成本高昂，而分页机制从根本上消除了整理的需求。唯一需要关注的碎片问题是内部碎片，可通过以下参数优化：

  • 块大小选择：较大的块（如 64 tokens）减少 block table 开销但增加内部碎片；较小的块（如 16 tokens）提高显存利用率但增加管理开销。推荐在 16-32 tokens 范围内根据序列长度分布调优。

  • 请求调度策略：通过 continuous batching 将短序列优先填充到现有批次的空闲槽位，减少新块的分配频率。

Buffer Reuse 与生命周期分析

tiny-vLLM 实现了显式的 Buffer Reuse 机制。在推理过程中，多个中间缓冲区（如 buf_2048_1、buf_2048_2）通过生命周期分析确定复用时机。当某个缓冲区的数据不再被后续计算依赖时，其内存区域可被后续操作复用，避免重复分配。

生命周期分析遵循以下原则：

  1. 追踪每个缓冲区的最后使用位置
  2. 确认后续计算流程中无读取依赖
  3. 将已释放的缓冲区标记为可复用
  4. 新缓冲区请求优先匹配尺寸相近的可复用区域

这种策略在 16 层 Transformer 的推理中尤为有效。例如，RMSNorm 的输出缓冲区在注意力计算完成后即可被 Feed Forward 网络复用，实现零额外显存开销的层间数据流转。

零拷贝推理路径

tiny-vLLM 通过减少 CPU-GPU 数据传输实现接近零拷贝的推理。模型权重在初始化阶段通过 cudaMemcpy 一次性加载到设备内存，后续推理仅涉及 token ID 和生成结果的传输。输入 token 序列通过 cudaMemcpyHostToDevice 复制到 GPU，但采用预分配的固定缓冲区（pinned buffer）减少传输开销。

在 decode 阶段，每次迭代仅传输新生成的单个 token ID，而非整个序列。KV 缓存的更新完全在设备端完成，通过 PagedAttention kernel 直接写入分配的块中，无需额外的数据搬运。

可落地参数与监控清单

基于 tiny-vLLM 的实现经验，以下是可直接应用的配置参数：

内存池配置：

  • 初始池大小：可用显存 × 0.75
  • Release Threshold：0.9（优先池内复用）
  • 块大小：16-32 tokens（根据平均序列长度调整）

运行时监控：

  • cudaMemPoolAttrReservedMemCurrent：监控池当前预留内存
  • cudaMemPoolAttrUsedMemCurrent：监控实际使用内存
  • 空闲块池水位：低于总块数 10% 时触发告警
  • 内部碎片率：计算未填满块的 token 浪费比例

性能调优检查点：

  • 批处理大小与显存占用的平衡点
  • 抢占策略的阈值设置（显存使用率 > 85% 时触发）
  • 主机内存交换带宽是否成为瓶颈

总结

tiny-vLLM 的 CUDA 内存池分层策略与 PagedAttention 分页机制展示了如何在资源受限环境下实现高效的 LLM 推理。通过固定大小的块分配、block table 映射、以及显式的 buffer 生命周期管理，系统在消除外部碎片的同时保持了低延迟的显存分配。这种设计哲学 —— 以预防替代治理、以复用替代分配 —— 为构建高性能推理引擎提供了可复用的工程范式。

资料来源：

  • tiny-vLLM GitHub 仓库（jmaczan/tiny-vllm）：Paged KV Cache 与 Buffer Reuse 实现
  • vLLM PagedAttention 论文：Efficient Memory Management for LLM Serving with PagedAttention
  • NVIDIA CUDA Runtime API 文档：Memory Pools 与 Stream-Ordered Allocator

ai-systems

内容声明：本文无广告投放、无付费植入。

如有事实性问题，欢迎发送勘误至 i@hotdrydog.com。