VLLM推理框架中的sleep_mode如何实现
vllm sleep_model 简介
前文torch_memory_saver 高性能 CUDA 内存管理工具实现,介绍了 sglang 中利用了该库将保存 kv_cache 和权重的显存释放出来。
在 VLLM 中也有同样的功能实现,在 VLLM 中的直接应用是“sleep mode”。将模型权重从显存(或者 NPU 内存中)卸载,并丢弃其中的 KV 缓存。
前文torch_memory_saver 高性能 CUDA 内存管理工具实现,介绍了 sglang 中利用了该库将保存 kv_cache 和权重的显存释放出来。
在 VLLM 中也有同样的功能实现,在 VLLM 中的直接应用是“sleep mode”。将模型权重从显存(或者 NPU 内存中)卸载,并丢弃其中的 KV 缓存。
torch_memory_saver是一个开源的高性能CUDA内存管理工具,主要功能是允许暂停和恢复PyTorch张量的CUDA内存占用。保持用户使用的虚拟地址不变,暂停后释放显存,恢复重新分配显存,绑定到虚拟地址上。
本文会介绍核心原理,以及拦截CUDA runtime API的实现。你还可以看到如何实现一个python c++扩展。在sglang大模型推理库中也有使用到这个torch_memory_saver库。
实现DeepSeek NextN推测解码。 通过forward_batch.spec_info.hidden_states获取推测解码的隐藏状态 使用eh_proj合并当前和推测的隐藏状态 支持分布式并行计算(通过enable_dp_attention标志)
SpecDecodeWorker + Proposers (ngram, draft model) + Scorer (top-1 scoring) + Verifier (rejection sampling)
关于投机采样是什么这里就不赘述了,参考这篇大佬的文章大模型推理妙招—投机采样(Speculative Decoding),解释的很清晰。
vllm框架对投机采样(Speculative Decoding)进行了支持, 见spec_decode。
投机采样也可以翻译为推测解码,我觉得推测解码会更加容易理解一些,下文采用推测解码。
推测解码是一种推理优化技术,它在生成当前Token的同时,对未来的Token进行有根据的猜测,这一切都在一次前向传播中完成。它融入了一种验证机制,以确保这些推测出的Token的正确性,从而保证推测解码的整体输出与普通解码的输出相同。优化大语言模型(LLMs)的推理成本,可以说是降低生成式人工智能成本并提高其应用率的最关键因素之一。为了实现这一目标,有各种推理优化技术可用,包括自定义内核、输入请求的动态批处理以及大型模型的量化。
(本小节译自deepseek-r1-and-fp8-mixed-precision-training)
深度探索(DeepSeek)发布了其推理模型深度探索-R1(DeepSeek-R1),震惊了世界。与 OpenAI 的 o1 和Google Gemini的Flash Thinking类似,R1 模型旨在通过在响应提示之前生成一条“思维链”来提高回复质量。R1 引起的兴奋源于它在包括数学、编码以及英语和中文语言理解等几个行业标准基准测试中与 o1 达到同等水平,同时它也是开源的,并且可以通过深度探索 API 以极低的成本获得。