torch_memory_saver高性能CUDA内存管理工具实现
torch_memory_saver是一个开源的高性能CUDA内存管理工具,主要功能是允许暂停和恢复PyTorch张量的CUDA内存占用。保持用户使用的虚拟地址不变,暂停后释放显存,恢复重新分配显存,绑定到虚拟地址上。
本文会介绍核心原理,以及拦截CUDA runtime API的实现。你还可以看到如何实现一个python c++扩展。在sglang大模型推理库中也有使用到这个torch_memory_saver库。
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 以极低的成本获得。
这篇文章探讨了 DeepSeek 模型中与专家混合 (MoE) 相关的另一个关键架构突破:辅助无损负载均衡策略。在本文中,我们将深入探讨 DeepSeek 如何解决 MoE 的隐藏瓶颈——负载均衡——同时消除梯度干扰并保留因果关系,为基于专家的模型的效率设定新标准。
首先介绍专家混合 (Mixture-of-Experts,MoE) 的基础知识,解释为什么负载平衡很重要,并回顾以前的工作,包括辅助损失方法(auxiliary loss methods)和专家选择(Expert Choice)。
CUTLASS库中的尾声融合(Epilogue Fusion)和尾声访问树(Epilogue Visitor Trees)
GEMM 在 NVIDIA GPU 上的高性能实现分为两个阶段:mainloop和epilogue。
mainloop:负责实际 GEMM 计算的部分
epilogue:其中进行了后处理(例如,元素激活、缩放)和数据存储
这篇文章中,我们将研究 CUTLASS 的尾声融合(epilogue fusion)方案。EVT在论文Epilogue Visitor Tree (EVT)中。
今天我们一起来阅读一篇文章deepseek-v3-explained-2-deepseekmoe,文章通过巧妙的例子生动讲解了DeepSeekMoE机制的原理。DeepSeekMoE是DeepSeek模型中的另一个关键架构创新。
将解释 Mixture-of-Experts (MoE) 的工作原理,是什么让它在 LLM 中如此受欢迎以及它面临的挑战。我们还将讨论专家专业化与知识共享之间的权衡,以及 DeepSeekMoE 如何设计以取得更好的权衡。为了使这些概念更直观,文章通过餐厅做菜选择厨师的例子,来类比分解它们,通过厨师在厨房中的角色来说明 MoE 中的每个元素。