讲师中心 微信公众号
AI工具推荐 视频效率加速

C++ 怎么编写一个支持大模型(如 LLaMA)推理的高性能 KV Cache 显存内存池

落强君_5812

落强君_5812

发布时间:2026-10-23 08:53:21

|

324人浏览过

|

来源于php中文网

原创

不能直接 new/delete 管理 KV Cache,因自回归解码中频繁 malloc/free 会引发 CUDA 同步、显存碎片化及分配失败;llama.cpp 中 std::vector 反复 resize 导致单次更新卡顿 2~5ms。

c++ 怎么编写一个支持大模型(如 llama)推理的高性能 kv cache 显存内存池

为什么不能直接 new/delete 管理 KV Cache

LLM 自回归解码时,KV Cache 按 token 步进动态增长,每层每头都要存 K 和 V 张量。一次请求从 100 token 扩到 4096 token,KV Cache 显存占用就涨 40 倍——频繁 malloc/free 会触发 CUDA 上下文同步、显存碎片化、甚至 cudaMalloc 失败。llama.cpp 里用 std::vector 或裸指针反复 resize,实测在长上下文场景下,单次 kv_cache_update 调用可能卡住 2~5ms,全是内存管理开销。

用 pinned pool + slab allocator 管理 GPU 显存块

核心思路是:预分配大块连续显存(比如 2GB),按固定大小切片(如 64KB/片),用位图标记空闲状态,避免 runtime 分配。C++20 的 std::pmr::monotonic_buffer_resource 不适合 GPU,得自己写。关键点:

  • 用 cudaMallocPitch 分配对齐内存(非 cudaMalloc),保证每个 KV slice 起始地址满足 128-byte 对齐,适配 Tensor Core 加载
  • 每个 slice 存一个完整 token 的 K+V(含 layer/num_kv_heads/head_dim/dtype),不跨 slice 存,避免 cache line 跨界
  • 用 std::atomic<uint64_t></uint64_t> 做位图,64 位掩码管 64 个 slice,bit=1 表示已分配;分配时用 __builtin_ctzll 找首个 0 位,O(1) 定位
  • 释放不立即归还位图,而是延迟到 batch 结束后统一回收(避免高频原子操作拖慢 decode loop)

如何对接 llama.cpp 的 kv_cache 结构

llama.cpp 的 struct llama_kv_cache 是 host 端控制结构,真正数据在 ctx->kv_self.k/v 两个 llama_tensor 里。你要改的是这两个 tensor 的 data 指针和 backend 字段:

  • 把 data 指向你 pool 里分配的 device ptr,不是 cudaMalloc 出来的零散地址
  • 把 backend 设为 LLAMA_BACKEND_GPU,并确保 ctx->cparams.n_gpu_layers > 0
  • 重写 llama_kv_cache_seq_rm:不调 cudaFree,只清位图对应 bit,并把该 slice 加入 freelist 链表(用 device ptr 做 next 指针,存在 slice 开头 8 字节)

注意:llama_kv_cache_update 默认假设 k/v 是连续 layout,你的 pool 必须严格按 [layer][kv_head][head_dim] 排列,否则 llama_get_kv_cache 读出来就是乱码。

C++14
C++14

C++14 对 C++11 的修正与增强版本,适合旧系统维护和较老工具链兼容。

下载

立即学习“C++免费学习笔记(深入)”;

容易被忽略的边界:context length 变化时的重映射

用户发来新请求,prompt 长度不定(128/2048/8192),但你的 pool 是固定大小 slice。不能让一个 8192-token 请求占满全部 128 个 slice——得支持“逻辑序列”到“物理 slice”的映射。做法是:

  • 每个请求维护一个 std::vector<int32_t></int32_t>,存它实际占用的 slice index(不是地址)
  • decode 第 n 个 token 时,查这个 vector 得到第 n 个 slice 的 offset,再算出 device ptr = pool_base + offset * slice_size
  • 当 context 扩展超过当前分配数,才触发新 slice 分配;缩小时不清内存,只 truncate vector

这个映射表本身很小(8192 个 int 是 32KB),放 host memory 即可,别往 GPU 上搬。真正的性能杀手从来不是计算,而是显存地址跳转——少一次 indirect access,就少一次 L2 cache miss。

热门AI工具

更多
豆包大模型

豆包大模型是一款由字节跳动推出的企业级大语言模型服务平台。

SkildArt
SkildArt Hot

SkildArt是一款AI文本写作工具,一站式 AI 视觉创作平台。

Seko
Seko Hot

一款AI视频创作工具,主要用于商汤科技推出的创编一体的AI短视频创作Agent,适合需要提升相关任务效率的用户。

VibeKnow
VibeKnow Hot

一款AI视频创作工具,主要用于全球首个AI知识视频创作平台,文档、文章、网页,一键生成视频,适合需要提升相关任务效率的用户。

Laper
Laper Hot

Laper是专为编剧、导演和制片人推出的 AI 原生剧本创作工具。

WorkBuddy

一款AI办公效率工具,主要用于腾讯云推出的AI原生桌面智能体工作台,适合需要提升相关任务效率的用户。

Atoms
Atoms Hot

Atoms是一款AI智能体工具,第一支自动构建真实业务的 AI 团队。

DeepSeek

DeepSeek是一款面向对话、写作、编程和推理场景的AI大模型工具。

UpDream
UpDream Hot

一款AI视频创作工具,主要用于哔哩哔哩推出的自研AI视频创作工具,适合需要提升相关任务效率的用户。

相关专题

更多
c++和c语言的区别有哪些
c++和c语言的区别有哪些

c++和c语言的区别:1、面向对象编程(OOP)支持不同;2、新增特性不同;3、标准库不同;4、编译方式不同;5、命名空间不同等等。想了解更多c++的相关内容,可以阅读本专题下面的文章。

2248

2024.03.14

c++和python学习顺序推荐
c++和python学习顺序推荐

一般建议先学习C++,再学习Python,因为这样可以逐步从较为底层的编程语言向更高级的语言过渡。想了解更多python的相关内容,可以阅读本专题下面的文章。

1019

2024.03.14

python和c++学习性价比分析
python和c++学习性价比分析

Python易于学习,广泛应用于Web开发、数据科学和人工智能等领域,但性能较低。C语言性能高,适用于对性能要求较高的场景,如游戏开发和系统编程,但学习曲线陡峭,错误处理复杂。想了解更多python的相关内容,可以阅读本专题下面的文章。

447

2024.03.14

c语言和c++一样吗
c语言和c++一样吗

c语言和c++是两种不同的编程语言,虽然有相似之处,但存在显著差异。c语言专注于过程式编程和系统级开发,以简洁、高效著称。c++作为c语言的超集,引入了面向对象编程,增强了代码组织和管理能力,但学习曲线也更陡峭。本专题为大家提供相关的文章、下载、课程内容,供大家免费下载体验。

327

2024.03.14

c语言和c++先学哪个好
c语言和c++先学哪个好

初学者选择学习c语言还是c++语言,需要根据个人学习目标、背景以及编程兴趣和预期应用方向来决定。本专题为大家提供相关的文章、下载、课程内容,供大家免费下载体验。

426

2024.03.14

c语言和c++的区别和联系
c语言和c++的区别和联系

c语言和c++是计算机科学领域应用广泛的编程语言。虽然它们有着相似的基础,但它们在语言类型、语法功能和内存管理方面存在着显著差异。本专题为大家提供相关的文章、下载、课程内容,供大家免费下载体验。

580

2024.03.14

c++软件中文更改教程
c++软件中文更改教程

对于 ide,可通过打开设置,找到语言设置,选择中文,并保存更改。对于非 ide 应用程序,可查找设置或选项,选择语言设置,更改为中文,并保存更改。想了解更多c++的相关内容,可以阅读本专题下面的文章。

1409

2024.03.21

python和java和c++学习性价比分析
python和java和c++学习性价比分析

Python以其易学性、丰富的库和活跃的社区而著称,适合数据科学、人工智能和Web开发。Java以其跨平台性、企业级应用开发和Android应用开发而闻名。C++以其底层控制能力、高效性能和游戏开发而著称。选择哪种语言取决于个人兴趣、职业方向和特定需求。想了解更多python和java和c++的相关内容,可以阅读本专题下面的文章。

1217

2024.03.22

Gradle插件开发入门教程
Gradle插件开发入门教程

本专题介绍Gradle插件开发的基本流程,涵盖插件类编写、自定义任务注册、扩展参数设计和插件应用方法,帮助你理解插件如何组织构建逻辑,减少多个项目之间的重复配置。

0

2026.10.22

热门下载

更多
网站特效
/
网站源码
/
网站素材
/
前端模板

精品课程

更多
相关推荐
/
热门推荐
/
最新课程
C++ Core Guidelines
C++ Core Guidelines

共0课时 | 0人学习

C++ Reference
C++ Reference

共0课时 | 0人学习

C++ 官方标准说明
C++ 官方标准说明

共0课时 | 0人学习

关于我们 免责申明 举报中心 意见反馈 讲师合作 广告合作 最新更新
php中文网:公益在线php培训,帮助PHP学习者快速成长!
关注服务号
PHP中文网订阅号
每天精选资源文章推送

Copyright 2014-2026 https://www.php.cn/ All Rights Reserved | php.cn