Skip to content

Latest commit

 

History

History
2398 lines (1768 loc) · 151 KB

File metadata and controls

2398 lines (1768 loc) · 151 KB

一、大模型推理与服务框架综述

1.1 框架概览与比较

二、专用推理与服务框架深度解析

2.1 vLLM 框架技术详解

2.2 TensorRT-LLM 框架

2.3 LightLLM 框架

2.4 SGLang 框架

2.5 LMDeploy 框架

2.6 其他推理工具与框架

三、用户友好型服务与交互工具

3.1 本地化部署与交互工具

3.2 新兴服务协议与框架

四、核心性能优化技术

4.1 注意力机制计算优化

4.2 服务端请求处理与调度

4.3 模型压缩与加速

4.4 大规模分布式训练

五、服务部署、评测与实践

5.1 服务性能评测体系

5.2 部署实践与接口调用

一、大模型推理与服务框架综述

1.1 框架概览与比较

1. 大模型推理框架概述

vLLM

vLLM全称Virtual Large Language Model,由Nvidia开源,旨在降低大模型推理的显存占用。其核心思想是将模型的一部分保存在CPU内存或硬盘上,只将当前计算所需的部分加载到GPU显存中,从而打破GPU显存限制。

vLLM支持PyTorch和FasterTransformer后端,可无缝适配现有模型。使用vLLM,在配备96GB内存+440GB A100的服务器上可运行1750亿参数模型,在配备1.5TB内存+880GB A100的服务器上可运行6万亿参数模型。

TensorRT-LLM

Tensorrt-LLM是Nvidia在TensorRT推理引擎基础上,针对Transformer类大模型推理优化的框架。主要特性包括:

  1. 支持多种优化技术,如kernel融合、矩阵乘优化、量化感知训练等,可提升推理性能
  2. 支持多GPU多节点部署,可扩展到万亿规模参数
  3. 提供Python和C++ API,易于集成和部署 在Nvidia测试中,基于OPT-30B在A100上的推理,Tensorrt-LLM可实现最高32倍加速。

DeepSpeed

DeepSpeed是微软开源的大模型训练加速库,最新的DeepSpeed-Inference也提供了推理加速能力,主要特点包括:

  1. 通过内存优化、计算优化、通信优化,降低推理延迟和提升吞吐
  2. 支持多GPU横向扩展,单卡可推理数百亿参数模型
  3. 提供Transformer、GPT、BERT等模型的推理示例
  4. 集成Hugging Face transformers库,使用简单

在GPT-NeoX测试中,基于DeepSpeed的推理相比原生PyTorch可实现7.7倍加速。

Text Generation Inference

Text Generation Inference(简称TextGen)是Hugging Face主导的开源推理框架,旨在为自然语言生成模型如GPT、OPT等提供高性能推理。主要特点包括:

  1. 高度优化的核心代码,支持FP16、int8等多种精度
  2. 支持多GPU多节点扩展,可推理万亿规模参数
  3. 良好的用户体验,提供Python高层API,简化开发
  4. 支持Hugging Face生态中的模型,如GPT2、GPT-Neo、BLOOM等

在OPT-175B基准测试中,TextGen可实现最高17倍推理加速。

Torch Dynamo

Torch Dynamo是PyTorch官方开发的一种动态图优化工具,旨在提高PyTorch模型的执行效率。它通过在运行时捕获和跟踪Python代码的执行,动态地将PyTorch的eager模式(即时模式)代码转换为更高效的图模式代码。这一转换使得 PyTorch模型在执行时能够更加接近静态图框架的性能,同时仍然保持PyTorch动态计算图的灵活性。

Torch Dynamo的核心思想是:

  1. 捕获 Python 函数的执行:在 Python 代码执行时,Torch Dynamo 会捕获函数的执行,分析其逻辑和计算图结构。
  2. 动态转换:它将捕获的代码转换为一个静态的、高效的中间表示(IR)。这种表示更容易进行进一步的优化,例如内存管理、算子融合等。
  3. 执行优化的代码:在转换完成后,优化后的代码会被重新执行,以获得更高的执行性能。

这个过程是透明的,对于用户来说,不需要修改现有的 PyTorch 代码。

FullyShardedDataParallel (FSDP)

FullyShardedDataParallel (FSDP) 是 PyTorch 提供的一种分布式数据并行训练技术,专门设计用于大规模模型的高效训练。它通过将模型的参数、梯度和优化器状态全面分片(sharding),并分布在多个 GPU 上,以优化内存使用,提升训练效率。FSDP 尤其适用于处理超大模型,这些模型通常无法在单个 GPU 内存中完全容纳。

FSDP 的工作原理:

  1. 全分片

    • 在传统的数据并行方法中,每个 GPU 上通常保存一份完整的模型副本,这样在处理超大模型时,会导致显存的巨大浪费。FSDP 的核心思想是将模型的参数、梯度和优化器状态按需分片,并分布在不同的 GPU 上。
    • 在每次前向和后向传播时,FSDP 会动态地将所需的分片参数加载到 GPU 上,并在计算完梯度后重新分片。
  2. 高效的内存管理

    • 通过分片技术,FSDP 可以极大地减少每个 GPU 所需的显存,从而能够处理更大的模型或在同一显存中容纳更多的模型参数。
    • 这种内存优化的机制使得即便是在显存资源受限的环境下,也能进行超大规模模型的训练。
  3. 与 ZeRO 的关系

    • FSDP 是一种全分片的数据并行方法,与 DeepSpeed 的 ZeRO 优化策略有一些相似之处。ZeRO 也通过将模型参数、梯度和优化器状态分散到多个设备上来减少显存使用。FSDP 可以被视为 ZeRO 的一种 PyTorch 原生实现,它与 PyTorch 的其他分布式训练功能高度集成。

Megatron-LM

Megatron-LM 是由 NVIDIA 开发的一种用于训练超大规模语言模型的深度学习框架。随着语言模型规模的不断扩大,训练这些模型变得越来越具有挑战性,特别是在处理数十亿到数万亿参数的模型时。Megatron-LM 专门设计来解决这些挑战,它通过多种并行化技术(如模型并行、数据并行和流水线并行)实现了高效的大规模模型训练。

Megatron-LM 的工作原理:

  1. 模型并行(Model Parallelism)

    • 在模型并行中,Megatron-LM 将一个巨大的模型分割成多个部分,每个部分分配给不同的 GPU。这种方法允许单个模型跨多个 GPU 进行训练,从而突破单个 GPU 显存的限制。
    • 具体实现包括将神经网络层或层内的参数矩阵划分到不同的设备上,并行计算。
  2. 数据并行(Data Parallelism)

    • Megatron-LM 也使用了数据并行技术,即将输入数据批次拆分为多个子批次,每个子批次在不同的 GPU 上独立计算梯度。然后,梯度在所有 GPU 之间进行同步,以确保模型参数的一致更新。
    • 数据并行是深度学习中常见的并行化方法,特别是在处理大型数据集时非常有效。
  3. 流水线并行(Pipeline Parallelism)

    • 为了进一步提高并行计算效率,Megatron-LM 引入了流水线并行。这种方法将模型的前向和后向传播过程划分为多个阶段,每个阶段在不同的 GPU 上执行。通过流水线并行,不同阶段可以同时进行,从而减少计算的等待时间,提高 GPU 利用率。
    • 流水线并行类似于工厂的流水线作业,不同的计算任务在不同的时刻完成,但最终达成整体的并行加速效果。
  4. 张量并行(Tensor Parallelism)

    • Megatron-LM 还支持张量并行,它进一步将模型的张量操作分解为更小的计算单元,这些单元分布在多个 GPU 上。这种方法特别适合处理超大规模的矩阵乘法等操作。

2. 主流推理框架对比分析(Ollama, vLLM, LMDeploy, TensorRT-LLM, SGLang)

在AIGC时代中,对大模型进行部署是不可缺少的一环,在DeepSeek系列火爆全球后,我们需要掌握Ollama、vLLM、LMDeploy、TensorRT-LLM、SGLang五个主流大模型部署工具的异同。

我们只有了解不同部署工具在易用性、性能和适用场景上的差异化优势,才能结合实际需求中的响应速度、硬件资源和部署复杂度进行权衡选择。

一、核心特性对比

工具名称 核心定位 优势 局限性 典型适用场景
Ollama 轻量级LLM本地部署框架 安装简单、跨平台支持、内存占用低、支持多模态模型 并发能力弱、国内下载速度慢 个人开发、本地测试、原型验证
vLLM 高性能LLM推理引擎 高吞吐量、支持多GPU并行、内存优化(PagedAttention) 配置复杂、显存占用高 企业级高并发服务、在线问答平台
LMDeploy 端到端LLM服务优化工具链 支持模型压缩(KV Cache量化)、长上下文优化、多框架兼容 生态相对较小、文档资源有限 移动端部署、边缘计算场景
TensorRT-LLM NVIDIA GPU专用推理加速框架 极致低延迟(TTFT优化)、支持混合精度、与TensorRT深度集成 模型转换复杂、依赖NVIDIA生态 自动驾驶实时决策、金融高频交易
SGLang 流式交互优化框架 支持复杂逻辑链式调用、异步流式响应、自定义语法扩展 学习成本较高、社区活跃度较低 AIGC多轮对话、游戏NPC交互系统

二、实际案例解析

1. Ollama案例:个人开发者本地调试

  • 场景:开发者需要在笔记本电脑上快速测试DeepSeek-R1模型的代码生成能力。
  • 操作:通过ollama run DeepSeek-R1一键启动模型,结合Zed AI编辑器实现代码补全。
  • 优势:5分钟内完成环境搭建,内存占用仅4GB,支持离线运行保障代码隐私。

2. vLLM案例:电商大促智能客服

  • 场景:双十一期间需处理每秒上千次的用户咨询。
  • 配置:使用4台A100 GPU部署vLLM集群,通过vllm --tensor-parallel-size 4启动分布式推理。
  • 效果:吞吐量提升3倍,响应延迟稳定在200ms内,支撑日均千万级查询。

3. TensorRT-LLM案例:自动驾驶路径规划

  • 场景:车辆需在10ms内完成障碍物轨迹预测。
  • 优化:将LSTM模型转换为TensorRT-LLM引擎,启用FP16精度和动态批处理。
  • 结果:首包延迟(TTFT)降低至5ms,较原始PyTorch实现提速8倍。

三、领域应用分析

AIGC领域

  • Ollama:本地运行Stable Diffusion文本生成图像流水线,避免云服务API调用成本。
  • vLLM:支撑AI直播带货脚本批量生成,单GPU可并行处理50个主播的个性化文案需求。
  • SGLang:构建多模态创作工作流(如「文本→分镜→配乐」链式生成),通过语法规则控制创作逻辑。

传统深度学习

  • LMDeploy:将70B参数的模型量化至4-bit后部署至手机端,实现离线翻译APP。
  • TensorRT-LLM:在医疗影像分析中,将ResNet-50推理速度提升至每秒3000帧,满足实时诊断需求。

自动驾驶

  • vLLM:用于车载语音助手,支持20路并发语音指令解析(如「导航到最近的充电站并播放新闻」)。
  • TensorRT-LLM:在百度Apollo系统中实现毫秒级车道线预测,较传统MPC控制器响应速度提升5倍。

四、选型建议

考量维度 推荐工具 理由
快速原型开发 Ollama 无需复杂配置,支持即时模型切换与本地调试
高并发生产环境 vLLM PagedAttention技术显著降低显存碎片,支撑千级QPS
边缘设备部署 LMDeploy + TensorRT-LLM 量化与硬件加速结合,实现大模型在Jetson设备上的实时运行
复杂交互逻辑 SGLang 支持状态保持与条件分支,适合多轮对话和流程控制场景

二、专用推理与服务框架深度解析

2.1 vLLM 框架技术详解

1. vLLM 整体架构与核心设计理念

vllm架构

架构组成

  1. Endpoints/Engine
  • Endpoints层是整个推理服务的入口,用于传入模型类型、推理参数、对于没一个request返回其output结果合集。 API Server是一个异步服务端,它会为每一个request关联一个output stream,用来返回推理生成的结果。 LLM多数用在同步代码和测试代码当中,他是对下层引擎的封装,并将引用层传入的参数组成engine_arges, 提供给引擎。
  • Engine层推理服务的核心控制层,一方面他会创建Scheduler,并传入request规划好需要使用的KVCache,另一方面初始化模型执行器,并将KVCache和request灌入执行器,最后得到模型的执行器返回给用户。
  1. Scheduler
  • 作为PageAttention的核心调度层,首先会把request转换为SequenceGroup,这样当我们使用例如Bean Search这些并行token生成策略的时,这一个requst就会变成一组相关token序列,我们为每个序列生成目标token,并选取其中最符合当前策略的token作为下一轮输入和结果输出。

  • Scheduler维护了3个队列,其中waitting队列主要是用户刚刚输入到scheuler中的token序列;running队列是正在进行推理的token序列;swaped队列是当显存不足或者推理优先级降低时,从GPU中换出的token序列

  • BlockManager下属两个KVCache分配器: UncachedAllocator、CachedAllocator. UncachedAllocator是正常token的KVCache分配器, 通过block_table维护分配状态,CachedAllocator略有差别,主要添加了token前缀hash计算,会将相同前缀的KVCache直接复用起来,并引入了evcitor等概念,被清除的KVCache会先暂时根据LRU原则转移到evcitor中,如果短时间内又出现相同前缀的token,可以恢复他的前缀KVCache使用。

  • 在Scheduler Output封装中,我们会将调度的KVCache和相关的sequence id传入给模型,这样就可以直接使用进行推理了。

  1. Excutor 执行器主要是通过worker的一种封装,其中既有单卡单worker的封装形式,同时也有分布式集群的worker形式。

  2. Worker 对应的每个Worker,不仅需要负责加载实际的模型,同时对需要对不同模型的推理逻辑进行捕获,使用CUDA Graph加速推理效率。 上文提到BlockManager会通过block_table规划好哪些物理Block需要使用,到了worker这一层就是对具体的物理内存进行操作。 以上动作都会直接和GPU或者其他计算平台打交道,因此还会封装一层Attention Backend, 屏蔽这些不同平台的差异细节。 最终我们可以获得模型的logits输出,通过我们的sample策略后,得出最终的Output, 返回给上一层。

  3. Backend Backend这层主要关注于推理和模型层计算过程中需要用到的各种加速算法和技巧,最后通过OPS的粘结剂层将最终的数据GPU中,执行相关的kernel函数进行并行计算。

2. vLLM 推理引擎与服务实现

引擎部分主要分为两部分:

入口: 主要分为同步和异步接口,会对输入的配置和数据进行转换成引擎args对引擎进行初始化。 引擎: 这个是顶层调度的核心模块,衔接scheduler和executor等下游核心模块。

vllm engine)

API-Server是用来对网络服务暴露的入口,主要使用了asyncio库进行了协程封装,提高了服务的并发处理能力。

  1. 服务端的启动args使用create_engine_config方法转换为EngineConfig,然后通过AsyncLLMEngine.from_engine_args方法创建出带有异步接口的引擎类AsyncLLMEngine(LLMEngine)。EngineConfig包含DeviceConfig、ModelConfig、CacheConfig、ParallelConfig、SchdulerConfig
  2. 调用generate方法将传入的prompt传输上一步创建的AsyncEngine引擎
  3. add_request这里会做两件事情:
  • start_background_loop: 开启协程主循环,监听新request事件并处理, 对应线条4。
  • 将request_tracker作为当前request的response stream追踪, 并将request放入到追踪队列,对应线条3
  1. engine_step是异步服务中较为核心的方法,主要会调用成员变量中的实际engine引擎,这个是个_AsyncLLEnginel类,对LLMEngine添加了一些异步封装方法。
  • 首先会主动通过get_new_and_aborted_requests去tracker拿到相应的request(线条5)
  • 然后调用add_request_async传入request请求(线条6、7), 这个不会调用模型中的tokenlizer做预处理和数据转换操作
  • 随后获得token_ids后, 将token按照block分组,组成sequence,这个工程主要发生在内部函数_add_processed_request中,最后再讲sequence转换成seq_group数据结构,在prefill阶段,这里只有一条sequence。
  • 最后会调用step_async,对request和此时生成的KVCache scheduler计划数据传给model(线条10、11)
  • 模型的输出,会直接传入process_request_outputs中, 主要是把output放入到tracker中,最终关联上这个request的stream,返回给用户(线条13、14)

LLM类是用于测试和同步调用时的封装,首先通过engine_args初始化LLMEgine,然后调用add_request接口,将prompt信息传入engine,此时engine会通过(路径8、9)来进行scheduler调度,最后12驱动模型,给出output。

  • 引擎是从from_engine_args方法入口创建的
  • add_request: 获取到的prompt,会经历tokenlize->sequence->sequence_group的数据处理操作,然后会作为参数传入add_seq_group
  • step/step_async方法是进行推理的核心方法,会首先通过scheduler对sequence group进行KVCache调度,拿到调度之后的metadata,同各种缓存和队列一起,传入给模型执行器,执行器返回推理结果。

3. vLLM 调度策略与显存块管理

调度器(Scheduler)决定哪些请求可以参与推理,并为这些请求做好逻辑块->物理块的映射。这个过程只是根据目前step(prefill+decode)收集到的所有信息,对目前GPU的显存进行新的内存分配规划,而未实际操作GPU内存。在模型执行器中,这些规划metadata会最终执行,正式分配相应的内存。

  1. 调度器图解 调度分为两步进行:

添加引擎提供的sequence group数据,存储至waitting队列。 对目前调度队列中的任务进行遍历和重新分配,需要修改成running状态的,需要对应分配物理block,同时根据资源情况进行swap int/out。整体图解如下

vllm scheduler

从上图图解可以看出,首先通过add_seq_group添加到waitting队列后,引擎对调用schedule方法通知调度器进行调度,调度后,scheduler还会对调度的输出进行简单的后处理,缓存相关的结果为下一步迭代做准备。

注:chunked prefill是在prefill阶段不会一次性的处理完所有prefill, 根据prefill的大小规划后面的budget,从而后面每次迭代(prefill+decode)都会带上这个prefill request id的处理,针对长prompt可以显著提高TTFT指标

  • 调度器首先从规划running队列的数据,将已完成的数据状态标记为finished,然后评估资源情况,对无法调度的sg放入swapped队列,对于sg size只有1的请求,直接放入watting队列,作为全新的request处理
  • 处理完running队列后,会再处理swap队列,处理需要swap in/out的sg。swap out时,针对prefill对放入prefill group中。针对无调度空间的sg,放弃调度,标记为finished状态。
  • 最后再观察waitting队列,将可以分配的sg规划好KVCache,标记为running状态,不满足资源的sg继续放在waitting队列里直到条件满足时再进行处理。
  1. Block分配

vllm block

在BlockManager中,利用block_table分配物理block

class BlockSpaceManagerV1(BlockSpaceManager):
    def __init__(
        self,
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
        watermark: float = 0.01,
        sliding_window: Optional[int] = None,
        enable_caching: bool = False,
    ) -> None:
        # 一个block的token个数,默认16
        self.block_size = block_size
        # gpu分配的block个数
        self.num_total_gpu_blocks = num_gpu_blocks
        # cpu分配的block个数
        self.num_total_cpu_blocks = num_cpu_blocks

        if enable_caching and sliding_window is not None:
            raise NotImplementedError(
                "Sliding window is not allowed with prefix caching enabled!")

        self.block_sliding_window = None
        if sliding_window is not None:
            # 滑动窗口的中的block大小
            self.block_sliding_window = math.ceil(sliding_window / block_size)

        # 水位90%
        self.watermark = watermark
        assert watermark >= 0.0
        self.enable_caching = enable_caching
        self.watermark_blocks = int(watermark * num_gpu_blocks)

        if self.enable_caching:
            # 开启前缀缓存的Block分配器
            self.gpu_allocator: BlockAllocatorBase = CachedBlockAllocator(
                Device.GPU, block_size, num_gpu_blocks)
            self.cpu_allocator: BlockAllocatorBase = CachedBlockAllocator(
                Device.CPU, block_size, num_cpu_blocks)
        else:
            self.gpu_allocator = UncachedBlockAllocator(
                Device.GPU, block_size, num_gpu_blocks)
            self.cpu_allocator = UncachedBlockAllocator(
                Device.CPU, block_size, num_cpu_blocks)

        # 每一个sequence对应一个BT
        # Mapping: seq_id -> BlockTable.
        self.block_tables: Dict[int, BlockTable] = {}

        # 对encode-decode模型,存在cross-attention层,我们为这个层需要引入单独的KVCache表
        # Mapping: req_id -> BlockTable
        # Note that each SequenceGroup has a unique
        # request ID
        self.cross_block_tables: Dict[str, BlockTable] = {}

UncachedBlockAllocator

在UncachedBlockAllocator中,是不带prefix缓存的块分配器,在不同调度队列流程中的图解,对running队列,这里的can_allocate只是通过目前剩余的block数量是否满足

def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:

    check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
    self_num_required_blocks = self._get_seq_num_required_blocks(
        seq_group.get_seqs(status=SequenceStatus.WAITING)[0])
    cross_num_required_blocks = self._get_seq_num_required_blocks(
        seq_group.get_encoder_seq())
    num_required_blocks = self_num_required_blocks + \
                            cross_num_required_blocks

    if self.block_sliding_window is not None:
        num_required_blocks = min(num_required_blocks,
                                    self.block_sliding_window)
    num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()

    # 所有block都不满足需求.
    if (self.num_total_gpu_blocks - num_required_blocks <
            self.watermark_blocks):
        return AllocStatus.NEVER
    # 剩余free满足,进行分配
    if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
        return AllocStatus.OK
    else:
    # 再等等
        return AllocStatus.LATER
  • 真正的分配物理块号
def _allocate_sequence(self, seq: Optional[Sequence], ref_count: int, \
                       is_encoder_decoder: bool = True) -> BlockTable:
    # 获得这个sequence需要多少个block.
    num_prompt_blocks = self._get_seq_num_required_blocks(seq)

    block_table: BlockTable = BlockTable()
    assert seq is not None
    # 每个逻辑块号从0开始
    for logical_idx in range(num_prompt_blocks):
        if (self.block_sliding_window is not None
                and logical_idx >= self.block_sliding_window):
            # 开启滑动窗口,重新映射逻辑块号id
            block = block_table[logical_idx % self.block_sliding_window]
            # Set the reference counts of the token blocks.
            block.ref_count = ref_count
        elif not is_encoder_decoder and self.enable_caching:
            # 开启了前缀匹配
            # hash值是对所有seq中的token进行整体hash,并将此值与驱逐器中进行比对
            # 最终或者创建新的物理快,或者从驱逐器中回收物理块,返回使用
            block = self.gpu_allocator.allocate(
                seq.hash_of_block(logical_idx),
                seq.num_hashed_tokens_of_block(logical_idx))
        else:
            # 直接从分配池中选取一个分配
            block = self.gpu_allocator.allocate()
            # Set the reference counts of the token blocks.
            block.ref_count = ref_count
        block_table.append(block)

    return block_table
  • waiting队列中的每个seq_group都还未经历过prefill阶段,因此每个seq_group下只有1个seq,这个seq即为prompt
  • 在使用UncachedBlockAllocator为wating队列中的某个seq_group分配物理块时,就是在对初始的这个prompt分配物理块。所以这个prompt有多少个逻辑块,就分配多少个可用的空闲物理块,同时注意更新物理块的ref_count。
  • 另外,这里给定一种“物理块的分配方案”,我们只是在制定这个seq_group可以使用哪些物理块,但并没有实际往物理块中添加数据。
  • 具体物理块分配,由CacheEngine按照这个方案,往物理块中实际添加KVCache。

以上主要waitting阶段的块分配,主要面向prefill的数据。

下面是running/swaped队列中的分配逻辑,主要面对decode阶段的数据,用到了block_manager中的can_append_slots和append_slots两个方法。

running/swaped

def can_append_slots(self, seq_group: SequenceGroup,
                     num_lookahead_slots: int = 0) -> bool:
    assert (num_lookahead_slots == 0), "lookahead allocation not supported in BlockSpaceManagerV1"

    # 就是看目前所有free的块是否满足running队列中需要的空间
    num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
    num_seqs = seq_group.num_seqs(status=SequenceStatus.RUNNING)
    return num_seqs <= num_free_gpu_blocks

def append_slots(
        self,
        seq: Sequence,
        num_lookahead_slots: int = 0,
    ) -> List[Tuple[int, int]]:
        n_blocks = seq.n_blocks
        block_table = self.block_tables[seq.seq_id]
        # 查找block_table的物理块个数,看是否还有能力分配
        # 如果物理块数量 < 逻辑块数量
        if len(block_table) < n_blocks:
            # 需要验证物理块只允许比逻辑块少1块
            assert len(block_table) == n_blocks - 1
            # 使用滑动窗口,做取整处理
            if (self.block_sliding_window
                    and len(block_table) >= self.block_sliding_window):
                # reuse a block
                block_table.append(block_table[len(block_table) %
                                               self.block_sliding_window])
            else:
                # 分配新物理块.
                new_block = self._allocate_last_physical_block(seq)
                block_table.append(new_block)
                return []

        # 如果最后一个物理块的引用数量为1, 也就是当前seq所引用
        last_block = block_table[-1]
        assert last_block.device == Device.GPU
        if last_block.ref_count == 1:
            # 如果开启前缀缓存,那么我们应该考虑物理块缓存问题.
            if self.enable_caching:
                # 当这个block已经填满是,我们需要在这里及时更新token hash,这有可能下次cache命中.
                maybe_new_block = self._maybe_promote_last_block(
                    seq, last_block)
                block_table[-1] = maybe_new_block
            return []
        else:
            # 已经有其他sq使用了这table,我们需要.
            # Copy on Write: Allocate a new block and copy the tokens.
            new_block = self._allocate_last_physical_block(seq)

            block_table[-1] = new_block
            # 从该seq的block_table中释放掉旧的物理块
            # 也即该物理块ref_count-=1,如果-=1后ref_count=0,说明该物理块彻底自由了,
            # 然后可以把它添加进驱逐器的列表中,他将变为可缓存的自由块
            self.gpu_allocator.free(last_block)
            return [(last_block.block_number, new_block.block_number)]

CachedBlockAllocator

  • 在CachedBlockAllocator为有前缀缓存的Block分配器,会引入一个叫驱逐器(evictor)的概念,会将移除的KVCache数据在驱逐器中再保存一段时间,如果出现同样的前缀,再调度到GPU显存中。

以下evictor的实现,它使用access time和num_hashed_tokens确定我们优先需要使用哪个block进行重用

class LRUEvictor(Evictor):
    def __init__(self):
        self.free_table: OrderedDict[int, PhysicalTokenBlock] = OrderedDict()

    def __contains__(self, block_hash: int) -> bool:
        return block_hash in self.free_table

    def evict(self) -> PhysicalTokenBlock:
        if len(self.free_table) == 0:
            raise ValueError("No usable cache memory left")

        evicted_block = next(iter(self.free_table.values()))
        # 执行驱逐策略:
        # 找到驱逐器free tables中last accessed time最早的那个物理块,把它驱逐掉,因为它已经很久没用了。
        # 按理来说,free_tables中的物理块都是按时间append的,即已经排序好了,我们第1块即可。
        # 但是若存在多个block的last_accessed一致,我们进行第二层判断
        # 就先移除掉包含用于做hash的tokens最多的那个, 我们因此就挑选它作为GPU复用的块。
        for _, block in self.free_table.items():
            if evicted_block.last_accessed < block.last_accessed:
                break
            if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
                evicted_block = block

        self.free_table.pop(evicted_block.block_hash)
        # 此块的计算状态设为false
        evicted_block.computed = False
        # 返回这个缓存块
        return evicted_block

    def add(self, block: PhysicalTokenBlock):
        self.free_table[block.block_hash] = block

    def remove(self, block_hash: int) -> PhysicalTokenBlock:
        if block_hash not in self.free_table:
            raise ValueError(
                "Attempting to remove block that's not in the evictor")
        block: PhysicalTokenBlock = self.free_table[block_hash]
        self.free_table.pop(block_hash)
        return block

    @property
    def num_blocks(self) -> int:
        return len(self.free_table)
  • 当一个物理块没有任何逻辑块引用时(例如一个seq刚做完整个推理),这时它理应被释放。但是如果开启了prefix caching,那么这个物理块当前没有用武之地,但可是如果不久之后来了一个新seq,它的prefix和这个物理块指向一致,这个物理块就可以被重复使用,以此减少存储和计算开销。所以,我们设置一个驱逐器(evictor)类,它的free_tables属性将用于存放这些暂时不用的物理块。
  • 所以目前,该设备上全部可用的物理块 = 正在被使用/等待被使用的物理块数量 + evictor的free_tables中的物理块数量
  • 在prefill阶段,当我们想创建一个物理块时,我们先算出这个物理块的hash值,然后去free_tables中看有没有可以重复利用的物理块,有则直接复用
  • 如果没有可以重复利用的hash块,那这时我们先检查下这台设备剩余的空间是否够我们创建一个新物理块。如果可以,就创建新物理块。
  • 如果此时没有足够的空间创建新物理块,那么我们只好从free_tables中驱除掉一个物理块,为这个新的物理块腾出空间,驱逐策略如下:
    • 先根据LRU(Least Recently Used)原则,驱逐较老的那个物理块,也就是上节说的access time
    • 如果找到多个最后一次使用时间相同的老物理块,那么则根据max_num_tokens原则,驱逐其hash值计算中涵盖tokens最多的那个物理块。
    • 如果这些老物理块的LRU和max_num_tokens还是一致的话,那就从它们中随机驱逐一个

prefix 从图中我们可以看出,当已存在seq0时,已经收集了所有的hash结果。当seq1来时,需要在分配block时,会先看驱逐器中是否可以利用的hash, 存在就直接使用,不存在我们会为seq1再开辟一个新的block供其使用,放入对应的block_table中。

4. PageAttention 内存管理技术详解

为什么要使用Page-Attention

LLM推理过程通常分为两个阶段:prefill和decode。通常会使用KV cache技术加速推理。 llm推理过程

  1. 预填充阶段。在这个阶段中,整段prompt喂给模型做forward计算。如果采用KV cache技术,在这个阶段中我们会把prompt过后得到的保存在cache k和cache v中。这样在对后面的token计算attention时,无需对前面的token重复计算了,可以节省推理时间。

在上面的图例中,假设prompt中含有3个token,prefill阶段结束后,这三个token相关的KV值都被装进了cache。 2) decode阶段,在这个阶段中,根据prompt的prefill结果,一个token一个token地生成response。 同样,如果采用了KV cache,则每走完一个decode过程,就把对应response token的KV值存入cache中,以便能加速计算。例如对于图中的t4,它与cache中t0~t3的KV值计算完attention后,就把自己的KV值也装进cache中。对t6也是同理。

由于Decode阶段的是逐一生成token的,因此它不能像prefill阶段那样能做大段prompt的并行计算,所以在LLM推理过程中,Decode阶段的耗时一般是更大的。 从上述过程中,我们可以发现使用KV cache做推理时的一些特点:

  • 随着prompt数量变多和序列变长,KV cache也变大,对gpu显存造成压力
  • 由于输出的序列长度无法预先知道,所以很难提前为KV cache量身定制存储空间

Page-Attention原理

虚拟内存的分页管理技术

  • 将物理内存划分为固定大小的块,称每一块为页(page)。从物理内存中模拟出来的虚拟内存也按相同的方式做划分
  • 对于1个进程,不需要静态加载它的全部代码、数据等内容。想用哪部分,或者它当前跑到哪部分,就动态加载这部分到虚拟内存上,然后由虚拟内存做物理内存的映射。
  • 对于1个进程,虽然它在物理内存上的存储不连续(可能分布在不同的page中),但它在自己的虚拟内存上是连续的。通过模拟连续内存的方式,既解决了物理内存上的碎片问题,也方便了进程的开发和运行。

Page-Attention可在不连续的显存空间存储连续的 key 和 value。用于将每个序列的 KV cache 分块(blocks),每块包含固定数量的 token 的 key 和 value 张量。 可以看到for的attention计算,KV cache 被划分为多个块,块在内存空间中不必连续 因为 blocks 在显存中不必连续,所以可以像虚拟内存分页一样,以更灵活的方式管理键和值:

  • 将 block 视为 page
  • 将 token 视为 bytes
  • 将序列视为进程 序列的连续逻辑块通过 block table 映射到非连续物理块。物理块可在生成新 token 时按需分配。因此只有最后一个block会发生显存浪费,小于4%。 block table映射 通过 block table 将逻辑块映射到物理块

在并行采样时,同一个 prompt 生成多个输出序列,这些序列生成时可以共享 prompt 的 attention 计算和显存。 与 OS 中进程共享物理 page 的方式类似,不同序列可以通过将其逻辑块映射到同一物理块来共享块。为了确保共享安全,Paged Attention 跟踪物理块的引用计数,并实现 “写时复制”(Copy-on-Write)机制,即需要修改时才复制块副本。内存共享使得显存占用减少 55%,吞吐量提升 2.2x。

注:写时复制(Copy-on-write,简称COW)是一种计算机程序设计领域的优化策略。其核心思想是,如果有多个调用者(callers)同时请求相同资源(如内存或磁盘上的数据存储),他们会共同获取相同的指针指向相同的资源,直到某个调用者试图修改资源的内容时,系统才会真正复制一份专用副本(private copy)给该调用者,而其他调用者所见到的最初的资源仍然保持不变。

5. vLLM 模型部署实践与关键参数解析

部署Qwen为例

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="qwen/Qwen-7B-Chat", revision="v1.1.8", trust_remote_code=True)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

参数介绍

SamplingParams有关参数

  • temperature:Temperature 参数是文本生成模型中用于控制生成文本的随机性和创造性的一个重要的超参数。Temperature参数通常设置为 0.1 到 1.0 之间。
  • top_k:模型预测的前k个最可能的下一个词。
  • max_tokens:模型生成的最大长度。
  • stop:生成模型停止生成的符号。

LLM有关参数

  • model:LLM模型路径。
  • tensor_parallel_size:并行处理的大小。
  • gpu_memory_utilization:默认为0.9, cpu_swap_space默认4个G。若gpu_memory_utilization参数过小(分配的内存大小低于模型使用内存)或者过大(接近1.0)时,代码会崩溃。
  • request_rate:请求速率

6. vLLM 处理长输入与token长度计算策略

当用户输入的文本长度超过模型的最大输入长度时,vllm的服务会抛出warrning,提示输入长度超过模型最大输入长度,输入的response为空。 解决方法(Qwen2.5为例): Qwen2.5 模型的上下文长度默认设置为 32768 个token。为了处理超出 32768 个token的大量输入,使用了 YaRN,这是一种增强模型长度外推的技术,确保在处理长文本时的最优性能。

vLLM 支持 YaRN,并且可以通过在模型的 config.json 文件中添加一个 rope_scaling 字段来启用它。例如,

{
  ...,
  "rope_scaling": {
    "factor": 4.0,
    "original_max_position_embeddings": 32768,
    "type": "yarn"
  }
}

目前 vLLM 只支持 静态 YaRN,这意味着无论输入长度如何,缩放因子都是固定的,这可能会影响处理较短文本时的性能。建议仅在需要处理长上下文时才添加 rope_scaling 配置。

对于token长度的计算,依赖于模型的tokenizer的配置文件。 输入token长度计算

from transformers import AutoTokenizer
tokenizer_path = "Qwen/Qwen2.5-7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, resume_download=True)
system_prompt = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
prompt = "Tell me something about large language models."
message = [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}]
message_template = self.tokenizer.apply_chat_template(
   message,
   tokenize=False,
   add_generation_prompt=True
)
num_input_tokens = len(tokenizer.tokenize(message_template))

输出token长度计算

from vllm import LLM, SamplingParams
sampling_params = SamplingParams(temperature=0.7, top_p=0.8, repetition_penalty=1.05, max_tokens=512)

# Input the model name or path. Can be GPTQ or AWQ models.
llm = LLM(model="Qwen/Qwen2.5-7B-Instruct")

message_template = self.tokenizer.apply_chat_template(
   message,
   tokenize=False,
   add_generation_prompt=True
)

# generate outputs
outputs = llm.generate(prompt=message_template, sampling_params=sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    num_output_tokens = len(tokenizer.tokenize(generated_text))

7. vLLM 多LoRA模型适配方案

在VLLM框架中,可以动态的部署多个lora模型,来减少计算资源的消耗。 多lora模型部署命令如下:

python3 -m vllm.entrypoints.openai.api_server \
                            --model /workspace/code/qwen2.5/7B \
                            --tensor_parallel_size 1 --enable_lora \
                            --lora-modules lora_1=/workspace/code/qwen2.5-7b/lora_1 lora_2=/workspace/code/qwen2.5-7b/lora2 

其中,--lora-modules参数指定了多个lora模型的路径,每个路径之间用空格分隔。 调用服务接口时,可以通过model参数指定需要调用的lora模型,例如:model="lora_1"。 相比较于部署一个完整的模型,部署多个lora模型可以减少计算资源的消耗,但是在推理速度和并发效果上会慢于将lora模型合并到完整模型中,因为模型在同一时刻只能进行一种任务的推理,在当前推理未结束,有其他任务需要推理,则会等待当前任务推理完成加载其他任务lora才能进行推理;该部署方式适用于一些简单的任务,例如文本分类、实体抽取、情感分析等。

2.2 TensorRT-LLM 框架

1. TensorRT-LLM 简介与核心优势

简介

TensorRT-LLM(NVIDIA官方支持)用于在NVIDIA GPU平台做大模型推理部署工作。
TRT-LLM基于TensorRT来将LLM构建为engine模型

TRT-LLM目前支持多种大模型,可以直接使用,在example中,而且还在以非常快的速度支持新的模型

TRT-LLM支持单机单卡、单机多卡(NCCL)、多机多卡,支持量化(8/4bit)

TRT-LLM的runtime支持chat和stream两种模式

TRT-LLM当前支持python和cpp(可以直接使用cpp,也可以使用cpp的bybind接口)两种模式的runtime

通过example下的各个模型的build.py来构建离线模型,通过example下的run.py(不同的业务适配一下run.py中的逻辑即可)来运行模型

TRT-LLM默认支持kv-cache,支持PagedAttention,支持flashattention,支持MHA/MQA/GQA等

安装使用

docker编译安装

// docker方式编译
step1: 安装操作系统匹配的docker,参考docker安装方式即可
step2: 下载 tensorrt-llm代码

# TensorRT-LLM uses git-lfs, which needs to be installed in advance.
apt-get update && apt-get -y install git git-lfs

git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
git lfs install
git lfs pull
// 上述每步都需要执行成功,由于网络问题,可能会失败,失败后重复执行,直到成功位置
// git lfs 这两步会将 tensorrt-llm/cpp/tensort-llm/batch_manager 下面的静态库 下载下来,后来编译会用到
batch_manager/
├── aarch64-linux-gnu
│   ├── libtensorrt_llm_batch_manager_static.a
│   ├── libtensorrt_llm_batch_manager_static.pre_cxx11.a
│   └── version.txt
├── x86_64-linux-gnu
│   ├── libtensorrt_llm_batch_manager_static.a
│   └── libtensorrt_llm_batch_manager_static.pre_cxx11.a
└── x86_64-windows-msvc
    └── tensorrt_llm_batch_manager_static.lib

step3:编译llm,提供了两种方式
方式一:一步到位的编译方式,推荐这种
make -C docker release_build  // 编译,此处cuda/tensorrt/cudnn/nccl等版本都是采用编译脚本中默认设置的
                              // 编译成功后,为一个docker镜像,大概有20多G,另外,docker方式编译对磁盘空间大小有要求
                              // 目前估计需要50G左右,如果docker的根目录空间不够,编译也会失败,可以通过给docker根目
                              //  扩容或者修改根目录来实现,保证编译空间的足够

make -C docker release_run // 运行编译成功的镜像, 此处需要有gpu办卡,如果在没有gpu的环境上,可以编译成功,但是执行会失败

方式二:逐步进行编译,编译结果和上述一致

编译有2种包,一种是仅包含cpp的代码包,一种是cpp+python的wheel包

//  仅cpp的代码包 : 仅编译 TensorRT-LLM/cpp 下面的c++和cuda代码
// cpp + python的包: 编译 TensorRT-LLM/cpp 和 TensorRT-LLM/tensortrt-llm 下面的c++ cuda python代码

参考文档

https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/docs/source/installation.md

2.3 LightLLM 框架

1. LightLLM 设计理念与架构简介

简单介绍

LightLLM 是一种轻量级 LLM 推理服务框架,LightLLM 引入了一种更细粒度的kvCache管理算法,称为TokenAttention, 并设计了一个与TokenAttention高效配合的Efficient Router调度算法。通过 TokenAttention 和 Efficient Router 的配合,LightLLM在大多数场景下实现了比vLLM和Text Generation Inference更高的吞吐量,甚至在某些情况下性能提升了4倍左右。

特点

  • 三进程异步协作:分词、模型推理、去分词异步进行,GPU利用率大幅提升。

  • TokenAttention:实现token-wise的KV缓存内存管理机制,实现推理时内存零浪费。

  • Efficient Router:与Token Attention合作,精心管理每个Token的GPU内存,从而优化系统吞吐量。

凭借基于OpenAI Triton开发的高度协调的高效内核和服务调度,LightLLM实现了优异的吞吐性能。

lightllm

组成

lightllm 的设计核心是多进程协作,每个进程负责一个模块,通过zmq和rpc的方式进行多进程协同工作。 lightllm中包括以下的模块:

Http Server : 负责接收请求

  • 接收API请求

  • 对于系统查询请求,跟 Metric Server 和 Health Server 协作获取相关信息

  • 针对于纯文本请求,将文本 tokenized,包装成纯文本请求发送给 Router

  • 针对于多模态请求,获取图片数据的md5码,使用md5码跟 Cache Manager Server 申请缓存,并将图片数据存到缓存上,将文本 tokenized,和多模态信息一起包装成多模态请求发送给 Visual Server

http_server

Router : 从 HttpServer 接收请求以后,主要负责保存请求,并且进行 请求调度

  • 接收 HttpServer 或者 Visual Server 发来的请求,并放到请求队列中。

  • 决定当前轮次应该 prefill 还是 decode。

  • 如果是 prefill 轮次, prefill 哪些请求。

  • 如果是 decode 轮次, decode 哪些请求。 Router

Model Backend :当 Router 决定好了使用哪些请求进行 prefill 或者 decode 以后, ModelBackend 决定如何处理这些请求。 lightllm\server\router\model_infer\mode_backend\base_backend.py 目录下的 ModeBackend 是所有 backend 的基类,通过了解其中的重要函数

  • init_model : 通过模型文件解析使用 lightllm-new-docs\lightllm\models 的哪个模型类。

  • prefill_batch : 对一个批次数据进行 prefill。

  • decode_batch : 对一个批次数据进行 decode。 每个backend都有一个 model代表一个独立的模型类, 以及一个 tp_rank 代表一个设备,可以有若干个 backend。 其中的 model 类负责模型在设备中真正地计算, lightllm\common\basemodel\basemodel.py 中的 TpPartBaseModel 是所有模型类的基类,该类支持张量并行。 model_backend

Visual Server : 负责处理多模态请求 Cache Manager Server :负责管理多模态信息的推理结果的缓存

Visual Server 和 Cache Manager Server 都是专门为了支持多模态模型的推理而设计的。其中 Visual Server 负责 encode 多模态模型中的图片信息, 而 Cache Manager Server 负责缓存图片原始数据和图片 encode 后的特征数据, 该缓存存放在主机的共享内存上,意在减少多进程的重复内存读取以及避免图片数据重复 encode。

Visual_Cache_Manager

Metric Server :负责记录系统运行的性能指标

Health Server :负责监控系统运行的健康情况

2. LightLLM TokenAttention 实现原理

运行机制:

模型初始化时,根据用户设置的 max_total_token_num 预先分配 KV 缓存,并创建 Token Table 来记录输入 token 的实际存储位置。

当处理新请求时,系统首先检查预分配的Token缓存中是否有可用的连续空间用于存储键值(KV)缓存。 TokenAttention 倾向于为请求分配连续的图形内存空间,以最大限度地减少推理过程中的内存访问。仅当连续空间不足时,才会为请求分配非连续显存。由于内存管理是逐个令牌进行的,因此 TokenAttention 几乎实现了零浪费,与 vllm 相比,产生了更高的吞吐量。

我们使用 OpenAI Triton 实现了一个高效的 TokenAttention 运算符。当提供查询向量时,该算子可以根据Token Table高效地检索相应的KV缓存并进行注意力计算。

请求完成后,可以通过删除令牌表上的记录来快速释放相应的显存,从而为调度新的请求让路。由于 TokenAttention 在模型初始化时预先分配了所有 KV 缓存空间,因此可以为已完成的请求高效释放内存,并在动态调度时合并不同批次的请求,从而有效最大化 GPU 利用率。

具体步骤:

  1. 模型初始化时,系统根据用户设置的 max_total_token_num 预先申请 KV 缓存显存,并创建 Token Table 来记录输入 token 的实际存储位置。 TokenAttention1

  2. 当处理新请求时,系统首先检查预分配的Token缓存中是否有可用的连续空间用于存储KV Cache。 TokenAttention 倾向于为请求分配连续的内存,以最大限度地减少推理过程中的内存访问。仅当连续空间不足时,才会为请求分配非连续的内存。分配的空间记录在Token Table中,用于后续的注意力计算。 TokenAttention2

  3. 对于新生成的Token的缓存,只需从预先分配的Token缓存中找到未使用的空间并将相应的条目添加到Token表中即可。此外,为了有效地分配和释放Cache,我们利用Torch Tensor在GPU上的并行计算能力来管理预分配Token Cache的状态。首先,我们定义状态如下:

self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
self._mem_cum_sum = torch.empty((size,), dtype=torch.int32, device="cuda")
self.indexes = torch.arange(0, size, dtype=torch.long, device="cuda")
self.can_use_mem_size = size

mem_state 记录了缓存的使用状态,其中1代表未使用,0代表已使用。 _mem_cum_sum 用于 mem_state 的累积和,用于有效地识别和选择未使用的空间进行缓存分配。分配过程如下:

torch.cumsum(self.mem_state, dim=0, dtype=torch.int32, out=self._mem_cum_sum)
#
select_index = torch.logical_and(self._mem_cum_sum <= need_size, self.mem_state == 1)
select_index = self.indexes[select_index]
self.mem_state[select_index] = 0
self.can_use_mem_size -= len(select_index)

TokenAttention3

  1. 请求完成后,可以通过删除 Token Table 上的记录来快速释放相应的显存,从而为调度新的请求让路。
self.can_use_mem_size += free_index.shape[0]
self.mem_state[free_index] = 1

TokenAttention4

  1. 由于Token级别的 GPU 内存管理,TokenAttention 可以实现 GPU 内存的零浪费。它可以准确地计算出系统可以容纳多少新Token进行计算。因此,当结合 Efficient Router 来管理请求时,它可以在推理过程中不断添加新的请求,充分利用每一块GPU内存,最大化GPU利用率。 TokenAttention5

3. LightLLM Efficient Router 设计解析

Efficient Router

引入高效路由器来管理传入请求,并动态确定该请求是否可以与已运行的推理批次融合。 合并标准是估计合并推理过程中最大Token占用量是否小于硬件可容纳的最大容量。 这里,我们将这个最大容量设置为 max_total_token_num。在 Token Attention 的支持下,我们可以准确地管理Token的使用情况,并且可以确保永远不会出现内存不足(out-of-memory)的情况。

Efficient Router

如上图所示,每一行代表一个请求当前的运行状态,黄色代表已经运行过的历史kv缓存token,每个格子代表一个token,灰色代表要生成的token。 生成的Token数量由每个请求设置的最大输出长度和已生成的Token数量决定。 上图中,绿色网格的第二行表示新到达的请求,图中按照要生成的输出的长度升序列出了所有请求。

如果我们假设新的请求融合成一个Batch进行推理,那么最大的token使用量必然会出现在时间点1、时间2、时间3中的一个时间点,我们只需要计算这些时间点的token使用量是否达到最大值即可。三个时间点都没有超过max_total_token_num,说明新的请求可以加入到Batch中进行融合推理。

时间1的总使用代币等于黄色单元格数量加上绿色单元格数量(见下图) Efficient_Router2

时间2的总使用代币等于黄色方块的数量加上绿色方块的数量(见下图) Efficient_Router3

时间3的总使用代币等于黄色方块的数量(见下图) Efficient_Router4

实际最大令牌使用量始终为时间 1、时间 2 或时间 3 之一。

只要动态推理过程中token的最大使用量低于max_total_token_num,就说明可以批量进行新的请求进行推理。 为了快速计算批次中所有请求所需的最大令牌使用量,我们使用 numpy 实现了一个高效的示例。

import numpy as np

def demo():
    max_total_token_num = 100
    req_list = [(5, 4), (4, 3), (5, 3), (3, 2), (4, 2)]  # (run_len, left_output_len)
    req_list.sort(key=lambda x: -x[1])

    left_out_len_array = np.array([e[1] for e in req_list])
    has_run_len_array = np.array([e[0] for e in req_list])
    cum_run_len_array = np.cumsum(has_run_len_array)
    size_array = np.arange(1, len(req_list) + 1, 1)
    need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max()

    if need_max_token_num <= max_total_token_num:
        print("ok")
    else:
        print("oom")

2.4 SGLang 框架

1. SGLang:语言模型的高效推理编程框架

SGLang 简介

SGLang 是一个用于大型语言模型和视觉语言模型的快速服务框架。它通过共同设计后端运行时和前端语言,使你与模型的交互更快、更可控。其核心功能包括:

  • 快速后端运行时: 提供高效的服务,包括 RadixAttention 用于前缀缓存、跳跃式约束解码、连续批处理、令牌注意力(分页注意力)、张量并行、- - FlashInfer 内核、分块预填充和量化(INT4/FP8/AWQ/GPTQ)。

  • 灵活的前端语言: 提供直观的界面用于编程 LLM 应用程序,包括链式生成调用、高级提示、控制流、多模态输入、并行性和外部交互。

  • 广泛的模型支持: 支持各种生成模型(Llama 3、Gemma 2、Mistral、QWen、DeepSeek、LLaVA 等)和嵌入模型(e5-mistral),并易于扩展以集成新模型。

SGLang的整体架构上分为前端和后端,整体功能上有4个大的功能分别为前端SGLang语言、RadixAttention、fast constrained decoding和API Speculative Execution。

2. SGLang 前端DSL设计与使用

DSL即领域特定语言,是特定领域内提供的一些变成语言和脚本,用于在该领域更高效或者高性能的完成代码编写工作,我们平时使用的shell算是一个。DSL本质上也是一个编译技术。在大模型训练推理领域,比较流行的是triton。triton可以认为是CUDA算子领域特定语言(也可以其他语言了)。SGLang可以认为是大模型应用开发的领域特定语言。

SGLang的前端部分提供一个LLM应用程序编写的DSL语言; 使用方式如下:

from sglang import function, system, user, assistant, gen, set_default_backend, RuntimeEndpoint

@function
def multi_turn_question(s, question_1, question_2):
    s += system("You are a helpful assistant.")
    s += user(question_1)
    s += assistant(gen("answer_1", max_tokens=256))
    s += user(question_2)
    s += assistant(gen("answer_2", max_tokens=256))

set_default_backend(RuntimeEndpoint("http://localhost:30000"))

state = multi_turn_question.run(
    question_1="What is the capital of the United States?",
    question_2="List two local attractions.",
)

for m in state.messages():
    print(m["role"], ":", m["content"])

print(state["answer_1"])

在SGLang中,提供了以下原语:

  • gen用于调用LLM生成;
  • select用于让LLM从列表中选择概率最高的选项;
  • +=或extend用于扩展当前的提示;
  • fork用于分叉当前的提示状态;
  • join用于重新连接分叉的提示状态。
  • 用户可以将这些原语与任意的Python控制流和库交叉使用。 前端来说,提供了自己的DSL,目标是替代LangChain这种框架。利用解释执行DSL的方式进行性能优化。

3. SGLang 后端服务启动与初始化过程

图中有三个组件 TokenizerManager、RouterManager 和 DetokenizerManager 有持续运行的函数,它们之间使用 zmq 库的 PULL 和 PUSH 模式来实现信息的发送和接收。 generate 时信息流转的过程如下:

  1. 客户端调用 generate API,将 GenerateReqInput 发送给 TokenzierManager,等待回复。
  2. TokenizerManager 将输入 encode 成 TokenizedGenerateReqInput,发送给 RouterManager;并且将监听这个 Request 的状态,一旦 DetokenizerManager 通知其已完成一次 Forward & Decode,TokenizerManager 就会触发一个事件,来更新 Request 的输出,直到完全 generate 结束就会返回一个 JSON 格式的输出给客户端。
  3. RouterManager 接收到 TokenizerManager 传来的输入会不断调用 forward 函数来生成,每一步的结果都将被包装成 BatchTokenIDOut 发送给 DetokenizerManager。每一步处理并发送之后会 sleep 0.0006s,并且如果有request 推理结束了那么将多 sleep extend_dependency_time(为了减少 cache miss)
  4. Detokenizer 接收到 RouterManager 的输入会开始 Decode 然后将结果包装成 BatchStrOut 返回给 TokenizerManager。

  1. launch_server: 启动后端使用的命令是 python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000,所以将执行 sglang/launch_server.py 这个文件

  2. TokenizerManager: 全局变量,负责管理 tokenizer,在用户 Prompt 请求到了之后 encode 好交给 RouterManager 做后续处理。在启动的时候主要负责:

    1. 绑定端口
    2. 初始化 Tokenizer(不考虑多模态模型)
    3. 另外在初次调用 Generate 的时候会将 handle_loop 函数放入事件循环,其中将接收 DetokenizerManager 传来的 decode 之后的输出。
  3. ModelRpcClient & ModelRpcServer RPC 即 Remote Procedure Call,表示客户端发送请求到服务端让服务端来执行功能,RPyC 是 Python 实现的 RPC 库。 ModelRpcClient 实现了客户端,在启动阶段它负责:

    1. 启动 ModelRpcServer
      • ModelRpcServer 实现了服务端,继承了 rpyc.Service。它的特别之处在于一旦方法以 exposed_ 开头那么这个方法是客户端可以远程调用的,否则只能被本地调用
    2. 远程调用 Server 的 exposed_init_model 函数来做初始化并拉起后续组件
      • Server 在 exposed_init_model 函数中主要负责启动图中 ModelRunner / RadixCache / Scheduler / FSMCache四个类
    3. 使用包装函数将 Server 的 exposed_step 函数包装成 async 函数(简化版的 rpyc.async_)
      • Server 的这个函数负责核心的 generate 逻辑
  4. ModelRunner & ReqToTokenPool & TokenToKVPool ModelRunner 启动阶段主要负责:

    1. 设置多卡推理相关的配置
    2. 执行 load_model来加载模型
    3. 计算目前剩余 GPU 显存,然后计算能够存储 KVCache 的最多 Token 数量 max_num_token,并执行 init_memory_pool 来初始化 ReqToTokenPool 和 TokenToKVPool。
      • TokenToKVPool 主要是用来预先分配 max_num_token 个 Token (每个token的空间称为 slot)的 KVCache 空间(一个 Tensor,我们称为 KVCache Tensor,其中每个 slot 的大小是 2head_numhead_dim*layer_num)。其中比较重要的函数是 alloc 和 alloc_contigous,前者将获得当前空闲的 need_size 个 slot,后者将得到连续的 need_size 个 slot。前者返回这些 slots在 KVCache Tensor 的 index 列表(out_cache_loc),这些位置将放在 RadixCache 维护的字典树的 value 值里;后者除了上述 index 列表,还会返回分配的起始位置和结束位置(end_loc = start_loc + need_size)
      • ReqToTokenPool 主要分配了 (size, max_context_len) 长度的空间,用来存储当前正在执行的所有 Request 的每个 token 对应的 KVCache Tensor index。这里的 size 设定了 BatchSize 的最大数量:max_num_token / context_len * 256。
  5. RadixCache: 维护了一个字典树 Trie Tree,主要用来实现 prefix 的匹配以此实现 RadixAttention。 字典树的节点 TreeNode 记录了用来前缀索引的 key(token_id sequence)、节点上的值 value(key 值所表示的 token_id sequence 对应的 KVCache Tensor 的 index)、引用计数 ref_counter(标记该前缀目前被多少正在推理的 request 占用)、最后访问时间 last_access_time(该树使用 LRU 策略删除节点,会将最早访问的叶节点依次删除)。 RadixCache 另外维护了一个变量 evictable_size_,记录了所有未被推理占用的 token 数量,之后可以 evict。 RadixCache主要有如下函数:

    1. insert:插入一个 key,自动匹配前缀,如果完全能匹配上就止于叶节点或者派生出子节点,否则就会 split 节点成两个
    2. match_prefix:找到完全匹配 key 前缀的节点,如果能半匹配上就 split 节点,如果能完全匹配上就找子节点
    3. evict:统计所有的叶节点,然后将最早的 num_tokens 个节点删除(这里使用到了最小堆算法)
    4. inc_ref_counter:将节点到根节点路径上所有的引用计数都加一,如果原来没有
    5. dec_ref_counter:类似上面 inc 函数。
  6. Scheduler 处理 Request 的顺序,有如下四种情况:1、lpm:默认值,Longest Prefix Match,更多 common prefix 的 Request 有更高的优先级;2. random:随机排序;3. fcfs:按照先来先处理的顺序;4. weight。

  7. FSMCache 主要是管理所有正则表达式当前的状态机,主要是配合前端

  8. RouterManager 在启动阶段主要负责:

    1. 绑定端口
    2. 在事件循环中执行 loop_for_recv_request ,不断获取 TokenizerManager 处来的信息放在 recv_reqs 中
    3. 在事件循环中执行 loop_for_forward,将上述 recv_reqs 交给 ModelRpcClient 的 step 函数执行,然后将结果交给 DetokenizerManager。
  9. DetokenizerManager 在启动阶段主要负责:

    1. 绑定端口
    2. 在事件循环中执行 handle_loop,不断获取 RouterManager 处得到的信息,Decode 之后发送给 TokenizerManager

4. SGLang 后端生成阶段执行流程

ModelRpcServer 得到 TokenizedGenerateReqInput 之后将通过 handle_generate_request 函数将其变成 Req 的数据结构,然后顺序放入 forward_queue 中用于之后的 forward。Req 的数据结构特别在于:如果 request 中包含正则表达式 regex,那么就会给其绑定一个 FSM,这块主要是和前端有关暂不展开。 接下来就会执行 forward_step进行正式的推理。

forward_step 函数中有如下逻辑:

  1. 调用 get_new_fill_batch 来将当前新的 request 动态组成 new batch。首先计算所有 request 在 RadixCache 中的共享 prefix(由此也就拿到对应的 kv cache),然后 Scheduler 重新安排所有 request 的优先级(默认是 LPM,即最长的共享 prefix 最先处理),接下来按照优先级依次放入 can_run_list 中并估计占用的显存空间,当可能会占满空间的时候就停止放入,最后将其组成一个 Batch 称为 new_batch。
  2. 如果 new batch 不为空,那么进入 Fill 阶段(forward_fill_batch),即分配新的显存给所有的 Request 的 token(init_extend_batch,注意此时已经分配好了 Extend 部分的 KV Cache 空间)、forward 计算所有的 logits(ModelRunner::forward)、采样得到 next_token(sample)、检查是否推理结束(check_finished),若结束进入 handle_finished_requests 函数。最后如果这些 new batch 没有推理完就放到 running_batch中,以后再调用的时候就会走 Decode 阶段。
  3. 如果 new batch 为空,那么就查看当前的 running_batch ,如果有 running_batch 就进入 Decode 阶段(forward_decode_batch),即给每条 request 即将生成的 1 个 token 分配连续空间(update_for_decode)、forward 计算所有的 logits(ModelRunner::forward)、采样得到 next_token(sample)、检查是否推理结束(check_finished),若结束进入 handle_finished_requests 函数。此处为了减少 Overhead,直接重复 Decode 阶段 10 次。
  4. Logging 相关:每经过 Decode 阶段 20 次就会打印一下当前的状态。

handle_finished_requests会将推理结束的 Req 包装成 BatchTokenIDOut,之后会发送给 DetokenizerManager,然后会:

  1. 更新 RadixCache(将推理结束的 request insert 到 RadixCache 中形成 common prefix,其中 key 是 token ids,value 是在 TokenToKVPool 的 indices)
  2. 将这条 request 的前缀在 TokenToKVPool 的引用计数减一(为什么是前缀部分?因为在 get_new_fill_batch 的时候就给 prefix 增加了引用计数)
  3. 将这条 Request 在 ReqToTokenPool 中去掉(ReqToTokenPool 只保存当前正在处理的 Request)
  4. RadixCache 将一路引用计数都减一,当引用计数为0的时候就可以被 evict(为什么要减引用计算,因为在 get_new_fill_batch 的时候有增加过引用计数来确保不被 evict)
  5. 将当前 batch 中去除推理结束的 Request

Generate 逻辑中最主要的是 ModelRunner::forward,首先它会将输入的各种信息包装成 InputMetadata(调用 InputMetadata.create),具体包括如下输入信息:

  1. forward_mode:分成 EXTEND(Fill 阶段调用)、DECODE(Decode 阶段调用)、PREFILL(暂时没有地方调用)
  2. tp_size:TensorParllel 的 size,分布式推理使用
  3. req_pool_indices:当前所有 request 占ReqToTokenPool的 index 列表
  4. seq_lens:当前所有 request 的长度列表
  5. prefix_lens:当前所有 request 的共享 prefix 的长度列表
  6. position_ids_offsets:一直都是 0,暂时可以忽略
  7. out_cache_loc:KVCache Tensor 中分配的显存 slot 的索引列表
  8. out_cache_loc_start:KVCache Tensor 中分配的显存 slot 的起始 index
  9. out_cache_loc_end:KVCache Tensor 中分配的显存 slot 的结束 index
  10. return_normalized_logprob:布尔值,返回是否要将输出 logits 做 normalize 然后会计算一些其他变量,例如 positions记录 token 的 position id:对于 EXTEND,因为需要推理从 prefix_len 到 seq_len 的所有 token 的 kv,所以positions 为 prefix_len 到 seq_len 的 id,input_ids 也只包含需要 extend 的部分;对于 DECODE,只需要推理一个 token,所以 positions 只包含新的一个 seq_len 的 id。 最后用 InputMetadata 和 input_ids 等调用具体模型的 forward

5. SGLang RadixAttention 缓存重用技术

RadixAttention,一种用于在运行时自动重用KV缓存的新技术。作者的方法不是在完成生成请求后丢弃KV缓存,而是将提示和生成结果的KV缓存保留在基数树中。这种结构可以实现高效的前缀搜索、重用、插入和淘汰。作者实现了最近最少使用(LRU)的淘汰策略,并辅以缓存感知的调度策略,以提高缓存命中率。此外,RadixAttention与现有的技术(如连续批处理和Paged Attention等)兼容。

基数树是一种数据结构,可替代前缀树。与典型的树不同,基数树的边不仅可以标记单个元素,而且还可以标记不同长度的元素序列。该特性显著提高了基数树的效率。在SGLang系统统中,作者使用基数树来管理一个映射。这个映射是在作为键的token序列和作为值的相应KV缓存张量之间的映射。这些KV缓存张量以非连续的分页布局进行存储,每个页的大小等于一个令牌的大小。

考虑到GPU内存的有限容量,因此无法重新训练无限的KV缓存张量,这就需要一个淘汰策略。为了解决这个问题,作者采用了递归淘汰叶节点的LRU淘汰策略。在连续批处理设置中,不能淘汰正在被当前运行的批处理使用的节点。因此,每个节点都维护一个引用计数器,显示有多少正在运行的请求正在使用该节点。如果一个节点的引用计数器为零,那么它就可以被淘汰。

前端始终将完整的提示发送到运行时,运行时将自动进行前缀匹配、重用和缓存。树结构存储在CPU上,维护开销很小。

树的结构维护历史的kv cache,从而减少prefix kv cache的计算,达到减少prefill计算的效果

6. SGLang 快速约束解码技术

Constrained Decoding(受限解码)是一种通过操控 LLM 的 token 生成过程,将模型的下一个 token 预测限制为仅生成那些不违反所需输出结构的 token。一般步骤如下:

  1. 定义约束规则:首先,我们需要定义约束规则。这些规则可以是正则表达式、上下文无关文法(CFG)等形式化的语法规则,或者是一些自定义的逻辑规则。
  2. LLM 生成候选 Token:基于当前上下文(即已经生成的 Token 序列),LLM 会生成一个候选 Token 列表,并为每个 Token 赋予一个概率值。
  3. 约束检查(Mask Gen):根据预定义的约束规则,检查候选 Token 列表中的每个 Token,确保它们符合设定的规则。
  4. 过滤(Apply Mask):将不符合规则的 Token 的概率值设置为 0(或极小值),从而排除这些 Token。
  5. 采样:根据过滤后的概率分布,从剩余的候选 Token 中随机采样一个 Token,作为下一个生成的 Token。
  6. 重复步骤 2-5,直到生成完整的文本序列。

现有系统一次只解码一个 token,导致解码速度不理想。它支持对结构化输出来进行更快的约束解码。现有系统仅通过掩盖不允许的令牌的概率来对下一个令牌遵循约束,使它们一次只能解码一个令牌。相反,SGLang的系统分析约束并构建一个压缩有限状态机来表示该约束。此方法在任何可能的情况下将多令牌路径压缩成单步路径,允许一次解码多个令牌以实现更快的解码速度。

在 LM 程序中,用户常常想制约模型的输出以遵循特定的格式,例如 JSON 概览。这可以改善控制能力和健壮性,以及让输出更易于解析。SGLang 提供了一个正则参数以使用正则表达式强制此类制约,而这些正则表达式对许多实际案例来讲都足够具有表达性。

SGLang 通过在压缩 FSM 中创建一个快速受限解码运行时环境,克服了这一限制。此运行时环境会分析 FSM,并将 FSM 中相邻的单一过渡边压缩为单个边,这样它便能够识别何时可以同时解码多个标记。压缩过渡边中的多个标记可以在一次前向传递中解码,这极大地加速了解码过程。

针对更快速受约束解码而配置的有限状态机的背景和实现细节。SGLang的目标是让 LLM 跟随正则表达式(regex),它提供更高的表达性,并可用于表示 JSON 架构等常见格式。为了达成这一目标,我们将正则表达式转化为一个有限状态机(FSM),以便在解码期间指导生成过程。FSM 本质上是一个带有节点(状态)和边(带有字符串/字符的转换)的图表。从一个初始状态开始,每个转换会附加边上的字符串以根据一组最终状态移动到下一个状态,从而完成该过程。这种机制会指导 LLM 的解码,根据 FSM 的当前状态转换来过滤无效标记,该解码过程可能包含在 FSM 中执行多次转换,直至达到一个最终状态。

2.5 LMDeploy 框架

1. LMDeploy 一站式部署工具链简介

简介

LMDeploy 是一个专为大语言模型(LLMs)和视觉-语言模型(VLMs)设计的高效且友好的部署工具箱,由 MMDeploy 和 MMRazor 团队联合开发。它提供了从模型量化到推理服务的全套解决方案,具有以下核心功能:支持多种模型格式和量化策略,支持多种硬件平台,支持多种推理引擎,支持多种部署方式。

LMDeploy 提供以下核心功能:

  • 高效的推理: LMDeploy 开发了 Persistent Batch(即 Continuous Batch),Blocked K/V Cache,动态拆分和融合,张量并行,高效的计算 kernel等重要特性。推理性能是 vLLM 的 1.8 倍
  • 可靠的量化: LMDeploy 支持权重量化和 k/v 量化。4bit 模型推理效率是 FP16 下的 2.4 倍。量化模型的可靠性已通过 OpenCompass 评测得到充分验证。
  • 便捷的服务: 通过请求分发服务,LMDeploy 支持多模型在多机、多卡上的推理服务。
  • 有状态推理: 通过缓存多轮对话过程中 attention 的 k/v,记住对话历史,从而避免重复处理历史会话。显著提升长文本多轮对话场景中的效率。
  • 卓越的兼容性: LMDeploy 支持 KV Cache 量化, AWQ 和 Automatic Prefix Caching 同时使用。

使用

  • 离线推理
import lmdeploy
pipe = lmdeploy.pipeline("internlm/internlm2_5-7b-chat")
response = pipe(["Hi, pls intro yourself", "Shanghai is"])
print(response)

在构造 pipeline 时,如果没有指定使用 TurboMind 引擎或 PyTorch 引擎进行推理,LMDeploy 将根据它们各自的能力自动分配一个,默认优先使用 TurboMind 引擎。

手动引擎选择

from lmdeploy import pipeline, TurbomindEngineConfig, PytorchEngineConfig
pipe = pipeline('internlm/internlm2_5-7b-chat',
                backend_config=TurbomindEngineConfig(
                    max_batch_size=32,
                    enable_prefix_caching=True,
                    cache_max_entry_count=0.8,
                    session_len=8192,
                ))

pipe = pipeline('internlm/internlm2_5-7b-chat',
                backend_config=PytorchEngineConfig(
                    max_batch_size=32,
                    enable_prefix_caching=True,
                    cache_max_entry_count=0.8,
                    session_len=8192,
                ))

注:参数 “cache_max_entry_count” 显著影响 GPU 内存占用。它表示加载模型权重后 K/V 缓存占用的空闲 GPU 内存的比例。 默认值是 0.8。K/V 缓存分配方式是一次性申请,重复性使用,这就是为什么 pipeline 以及下文中的 api_server 在启动后会消耗大量 GPU 内存。 如果遇到内存不足(OOM)错误的错误,可能需要考虑降低 cache_max_entry_count 的值。

当使用 pipe() 生成提示词的 token 时,可以通过 GenerationConfig 设置采样参数

from lmdeploy import GenerationConfig, pipeline

pipe = pipeline('internlm/internlm2_5-7b-chat')
prompts = ['Hi, pls intro yourself', 'Shanghai is']
response = pipe(prompts,
                gen_config=GenerationConfig(
                    max_new_tokens=1024,
                    top_p=0.8,
                    top_k=40,
                    temperature=0.6
                ))

在 GenerationConfig 中,top_k=1 或 temperature=0.0 表示贪心搜索。

2. LMDeploy TurboMind 推理引擎解析

LMDeploy的TurboMind技术是其核心推理引擎,专为高效处理大型语言模型(LLM)和多模态模型设计

  1. 核心架构
  • 持久化批处理(Persistent Batch) 支持多轮对话中缓存注意力机制的键值(KV)数据,避免重复计算历史会话,显著提升长文本对话效率。这一机制通过动态管理KV缓存,实现连续批处理请求的即时加入和完成请求的自动退出。

  • 可扩展的KV缓存管理器 采用分块管理策略(Blocked KV Cache),支持按需分配显存。例如,通过参数cache_max_entry_count调整KV缓存占显存的比例(如默认0.8调整为0.2以解决OOM问题),并支持缓存命中时的快速检索与缓存未命中时的自动解码恢复。

  • 动态拆分融合与张量并行 通过动态拆分计算图、融合算子及张量并行技术,优化GPU资源利用率。例如,在单卡部署时,异步采样耗时仅为同步采样的2/3,训练速度提升约16%。

  • 高性能CUDA内核 基于FasterTransformer实现的高效CUDA算子,支持Flash Attention 2和Split-K decoding,优化推理速度。

  1. 量化与性能优化
  • 量化支持

    TurboMind支持4bit权重(W4A16)和8bit KV Cache量化。例如,4bit模型推理效率是FP16的2.4倍,且通过OpenCompass验证可靠性。量化过程自动完成(如将HuggingFace模型在线转为TurboMind格式),无需用户干预。

  • 吞吐量优势

    在BatchSize=64时,吞吐量超过2000 token/s,比DeepSpeed提升5%-15%。对比vLLM,真实数据吞吐量效率高30%,多轮对话场景下训练速度提升至基础实现的70%

  1. 多场景适配与部署
  • 多模态支持

    • 扩展至视觉语言模型(如InternVL2-26B),支持图像与文本联合推理,结合动态批处理优化多模态任务性能。
  • 灵活部署模式

    • API服务:通过RESTful API提供类OpenAI接口,支持多模型多机多卡推理,兼容Swagger UI交互测试。

    • 本地与云端集成:支持Python代码直接调用(如pipeline接口),并可通过调整TurbomindEngineConfig参数实现量化推理和显存优化。

    • 演示工具:与Gradio快速集成,构建交互式Demo。

2.6 其他推理工具与框架

1. KTransformers 简单介绍

简介

KTransformers 是由清华大学 KVCache.AI 团队联合趋境科技推出的开源项目,旨在优化大语言模型的推理性能,降低硬件门槛。它基于 GPU/CPU 异构计算策略,利用 MoE 架构的稀疏性,支持在仅 24GB 显存的单张显卡上运行 DeepSeek-R1、V3 等 671B 参数的满血版大模型,预处理速度最高可达 286 tokens/s,推理生成速度最高能达到 14 tokens/s。

项目通过基于计算强度的 offload 策略、高性能算子和 CUDA Graph 优化等技术,显著提升了推理速度,使得普通用户和中小团队能够在消费级硬件上运行千亿级参数模型,实现“家庭化”部署。

主要功能

  • 支持超大模型的本地推理:支持在仅 24GB 显存的单张显卡上运行 DeepSeek-R1 等 671B 参数的满血版大模型,打破传统硬件限制。
  • 提升推理速度:预处理速度最高可达 286 tokens/s,推理生成速度达 14 tokens/s。
  • 兼容多种模型和算子:支持 DeepSeek 系列及其他 MoE 架构模型,提供灵活的模板注入框架,支持用户切换量化策略和内核替换,适应不同优化需求。
  • 降低硬件门槛:将大模型的显存需求大幅降低,让普通用户和中小团队能在消费级硬件上运行千亿级参数模型,实现“家庭化”部署。
  • 支持长序列任务:整合 Intel AMX 指令集,CPU 预填充速度可达 286 tokens/s,相比传统方案快 28 倍,将长序列任务的处理时间从“分钟级”缩短到“秒级”。

技术原理

  • MoE架构:将稀疏的 MoE 矩阵卸载到 CPU/DRAM 上处理,稠密部分保留在 GPU 上,大幅降低显存需求。
  • offload策略:根据计算强度将任务分配到 GPU 和 CPU:计算强度高的任务(如 MLA 算子)优先分配到 GPU,计算强度低的任务分配到 CPU。
  • 高性能算子优化
    • CPU端:用 llamafile 作为 CPU 内核,结合多线程、任务调度、负载均衡等优化,提升 CPU 推理效率。
    • GPU端:引入 Marlin 算子,专门优化量化矩阵计算,相比传统库(如 Torch)实现 3.87 倍的加速效果。
  • CUDA Graph 优化:基于 CUDA Graph 减少 Python 调用开销,降低 CPU/GPU 通信的断点,实现高效的异构计算协同。每次 decode 仅需一个完整的 CUDA Graph 调用,显著提升推理性能。
  • 量化与存储优化:采用 4bit 量化技术,进一步压缩模型存储需求,仅需 24GB 显存即可运行 671B 参数模型。同时优化 KV 缓存大小,减少存储开销。
  • 模板注入框架:提供基于 YAML 的模板注入框架,支持用户灵活切换量化策略、内核替换等优化方式,适应不同场景的需求。

源码链接:https://github.com/kvcache-ai/ktransformers.git

三、用户友好型服务与交互工具

3.1 本地化部署与交互工具

1. Ollama:简化本地大模型运行的利器

1. 什么是Ollama

Ollama 是一个开源框架,专注于在本地环境中快速部署和运行大型语言模型(LLMs),支持多种模型格式(如 GGUF、PyTorch 等),提供轻量化的推理服务。其核心目标是降低用户使用 LLMs 的门槛,尤其适合AI开发者和AI研究者进行本地实验或私有化部署。通过 Ollama,AI开发者能以极低门槛将前沿 AI 能力整合到实际业务中,同时保持对数据和模型的全流程控制。以下是其核心特性:

  • 本地优先:无需依赖云端服务,支持本地 CPU/GPU 推理。
  • 多模型支持:兼容 DeepSeek、Llama、Mistral、Phi 等主流模型。
  • 轻量化 API:提供类似 OpenAI 的 RESTful API,便于集成到现有系统中。
  • 多模态扩展:支持文本、图像生成(如 LLaVA)等多模态任务。
  • 资源优化:通过量化技术(如 4-bit/8-bit)降低显存占用。

2. DeepSeek 部署示例

DeepSeek 是由中国深度求索公司开发的高性能开源语言模型,在数学推理、代码生成等任务中表现优异。以下是使用 Ollama 部署 DeepSeek 的步骤:

环境准备

  • 操作系统:Linux/macOS/Windows(需 Docker 或直接安装)
  • 硬件要求:至少 8GB 内存,推荐 NVIDIA GPU(支持 CUDA)
  • 安装 Ollama
    # Linux/macOS 一键安装
    curl -fsSL https://ollama.com/install.sh | sh
    
    # Windows 可通过 Docker 部署
    docker run -d -p 11434:11434 --name ollama ollama/ollama

下载 DeepSeek 模型

Ollama 支持直接从其模型库拉取预配置模型:

# 拉取 DeepSeek 7B 模型(支持中文)
ollama pull deepseek-7b

# 或自定义模型(需手动配置 Modelfile)
ollama create deepseek-custom -f ./Modelfile

Modelfile 示例

FROM deepseek-7b
PARAMETER num_gpu 1  # 启用 GPU 推理

启动服务

# 后台运行模型服务
ollama serve

# 交互式调用
ollama run deepseek-7b "如何用 Python 实现快速排序?"

API 调用

通过 RESTful API 集成到应用:

import requests

response = requests.post(
    "http://localhost:11434/api/generate",
    json={
        "model": "deepseek-7b",
        "prompt": "解释量子计算的 Shor 算法",
        "stream": False
    }
)
print(response.json()["response"])

3. Ollama 的经典应用案例

3.1 AIGC(生成式人工智能)

  • 案例:内容创作自动化
    • 场景:营销团队使用 DeepSeek 生成广告文案和社交媒体内容。
    • 实现:通过 Ollama 部署模型,结合 LangChain 构建自动化流水线:
      from langchain_community.llms import Ollama
      llm = Ollama(model="deepseek-7b")
      print(llm("生成一篇关于环保的微博文案,要求包含#碳中和#标签。"))
    • 优势:本地部署保障数据隐私,避免云端 API 调用成本。

3.2 传统深度学习

  • 案例:图像描述生成(多模态扩展)
    • 场景:将 LLaVA 模型与 Ollama 结合,为医学影像生成文本描述。
    • 部署
      # 拉取多模态模型
      ollama pull llava
      # 上传图像并获取描述
      ollama run llava "描述这张图片的内容" -i ./x-ray.jpg
    • 技术栈:结合 CLIP 视觉编码器与语言模型,实现端到端推理。

3.3 自动驾驶

  • 案例:场景理解与决策推理
    • 场景:在车载边缘设备部署小型语言模型,解析传感器数据并生成驾驶决策。
    • 实现
      1. 使用 Ollama 量化部署 Phi-3 等轻量模型(<4GB 显存占用)。
      2. 输入激光雷达点云数据文本化描述:
        "前方 50 米处有行人正在横穿马路,当前车速 60km/h,请建议刹车力度。"
        
      3. 模型输出结构化指令:
        {"action": "brake", "intensity": 0.7, "confidence": 0.92}
    • 优势:低延迟本地推理避免网络不稳定问题,符合车规级安全要求。

2. Open WebUI:功能丰富的开源Web用户界面

一、Open WebUI 简介

Open WebUI 是一款开源、可扩展的 Web 界面工具,专为本地部署的大型语言模型(LLM)和深度学习模型提供可视化交互支持。其核心优势在于完全离线运行、用户友好的图形界面,以及无缝集成多种模型服务(如 Ollama、OpenAI API 等)的能力。以下是其核心特性:

  1. 本地化与隐私保护:所有数据在本地处理,避免云端传输风险,适合医疗、金融等敏感场景。
  2. 多模态支持:支持文本、图像、文件上传与解析(如 PDF、知识库构建),内置 RAG(检索增强生成)引擎。
  3. 轻量化与跨平台:通过 Docker 容器化部署,支持 Windows、Linux、macOS 等操作系统,适配不同硬件配置。
  4. 生态兼容性:与 Ollama 等工具深度集成,可自动识别本地模型并提供统一管理界面。

二、DeepSeek 部署示例(基于 Ollama + Open WebUI)

DeepSeek-R1(7B 版本)为例,展示本地部署流程:

1. 环境准备

  • 硬件要求:至少 16GB 内存,NVIDIA GPU(显存 ≥4GB)。
  • 软件依赖:安装 Ollama、Docker,并配置 NVIDIA 容器工具包(Linux/Windows)。

2. 部署步骤

  1. 安装 Ollama
    # Linux/macOS 一键安装
    curl -fsSL https://ollama.com/install.sh | sh
    # Windows 通过 Docker 部署
    docker run -d -p 11434:11434 --name ollama ollama/ollama
  2. 下载 DeepSeek-R1 模型
    ollama pull deepseek-r1:7b  # 自动下载模型权重
  3. 启动 Open WebUI(Docker 方式)
    docker run -d -p 3000:8080 --gpus all -v open-webui:/app/backend/data --name open-webui ghcr.io/open-webui/open-webui:main
  4. 配置与使用
    • 访问 http://localhost:3000,注册管理员账号。
    • 在 Open WebUI 中自动识别本地 DeepSeek 模型,支持对话、文件上传、知识库管理。

3. 调试与优化

  • 常见问题:若出现连接错误,需确保 Ollama 服务已启动(ollama serve)。
  • 性能调优:通过调整 GPU 资源分配和模型量化参数(如 num_gpu)提升推理速度。

三、Open WebUI 的经典应用案例

1. AIGC(生成式人工智能)

  • 案例 1:智能内容生成
    用户通过 Open WebUI 上传行业报告,调用 DeepSeek 生成定制化营销文案,并结合知识库自动引用企业数据,提升内容准确性。
  • 案例 2:代码助手
    开发者使用 DeepSeek-Coder 模型生成代码片段,通过 Open WebUI 实时调试并导出结果,集成到 VS Code 等开发工具中。

2. 传统深度学习

  • 案例:医学影像分析
    结合多模态模型(如 LLaVA),上传 X 光片图像,模型生成诊断描述,并通过 Open WebUI 展示热力图辅助医生决策。
  • 案例:工业质检
    部署 YOLO 目标检测模型,实时识别生产线缺陷,结果通过 Web 界面可视化并触发自动化分拣系统。

3. 自动驾驶

  • 案例:实时场景理解
    车载边缘设备部署轻量级模型(如 Phi-3),通过 Open WebUI 解析激光雷达数据文本化描述,生成驾驶决策(如刹车力度)并可视化边界框。

3.2 新兴服务协议与框架

1. MCP:模型上下文协议解析

一、技术定义与核心价值

MCP(Model Context Protocol)是由Anthropic提出的AI系统与外部环境交互的标准化协议,其本质是建立 "大模型与真实世界的通信管道" 。通过三大核心机制实现价值突破:

  1. 统一接口层:将各类API/工具抽象为标准化操作指令(类似HTTP协议统一网络通信)
  2. 动态上下文管理:实时感知环境状态变化并同步到模型(如传感器数据流)
  3. 安全沙箱机制:限制工具调用的权限边界(防止越权操作)

二、通俗案例解析:智能旅行规划系统

假设用户输入:"帮我规划北京三日游,要包含故宫和环球影城,预算5000元"

传统方式需要工程师编写:

  • 高德API调用(路线规划)
  • 美团API接入(餐饮推荐)
  • 天气预报接口(行程调整)
  • 预算计算逻辑

基于MCP的实现

# MCP服务注册(伪代码)
register_tool("高德地图", type="navigation", params={"city": "北京"})
register_tool("美团餐饮", type="poi", params={"budget": 5000})
register_tool("天气预报", type="environment")

# LLM直接生成指令链
plan = llm.generate(
   "用户需求:北京三日游,故宫+环球影城,预算5000",
   allowed_tools=["高德地图", "美团餐饮", "天气预报"] 
)

# 输出结构化指令
{
  "day1": [
    {"tool": "高德地图", "action": "path_plan", 
     "params": {"start": "故宫", "end": "酒店"}},
    {"tool": "美团餐饮", "action": "recommend",
     "params": {"location": "故宫周边", "price_range": "人均100-150"}}
  ]
}

开发效率从3人周降至2小时,且后续可无缝替换同类服务(如将高德替换为百度地图)

三、三大领域应用详解

1. AIGC领域应用

典型场景:短视频自动生成

  • 传统痛点:需要人工串联文案生成→素材搜索→视频剪辑工具链
  • MCP解决方案
    • 集成Stable Diffusion MCP(文生图)
    • 接入GettyImages MCP(版权素材库)
    • 对接Premiere Pro MCP(自动化剪辑)

工作流示例

用户输入 → 剧本生成 → [分镜指令] → 调用SD生成画面 → 
[选取素材] → 调用GettyImages搜索 → [时间线编排] → Premiere渲染输出

效率提升:单条视频制作从2小时缩短至8分钟

2. 传统深度学习应用

典型场景:工业质检模型迭代

  • 传统痛点:数据标注→模型训练→部署验证形成数据孤岛
  • MCP解决方案
    • 连接LabelStudio MCP(自动化标注)
    • 对接MLFlow MCP(实验追踪)
    • 集成EdgeDeploy MCP(端侧部署)

数据流闭环

产线相机 → 异常检测 → [标注指令] → LabelStudio标注 → 
[训练任务] → PyTorch MCP训练 → [部署指令] → EdgeDeploy更新模型

迭代周期从2周缩短至12小时,缺陷检出率提升9.3%

3. 自动驾驶应用

典型场景:复杂路况决策

  • 传统痛点:各传感器数据与规控系统耦合度过高
  • MCP解决方案
    • 激光雷达MCP(点云数据标准化)
    • 高精地图MCP(实时路况更新)
    • V2X MCP(车路协同)

决策流程优化

激光雷达 → 障碍物检测 → [定位指令] → 高精地图匹配 → 
[路径规划] → 结合V2X信号灯信息 → 生成控制指令

在北京亦庄实测中,十字路口急刹率降低67%

四、技术对比分析

维度 传统API调用 MCP协议方案
开发成本 每个API需单独适配(平均3人日/个) 一次接入永久复用(0.5人日/个)
异常处理 需手动编写fallback逻辑 内置自动服务降级机制
多工具协同 需开发中间件协调 原生支持管道式调用
安全审计 分散在各服务 统一权限控制与日志追踪

四、核心性能优化技术

4.1 注意力机制计算优化

1. FlashAttention:IO感知的注意力计算优化

GPU的内存由多个不同大小和不同读写速度的内存组成。内存越小,读写速度越快。对于A100-40GB来说,内存分级图如下所示 flash-attention

  • SRAM内存分布在108个流式多处理器上,每个处理器的大小为192K,合计为192*108KB=20.25MB 相当于计算块,但内存小
  • 高带宽内存HBM(High Bandwidth Memory),也就是我们常说的显存,大小为40GB。SRAM的读写速度为19TB/s,而HBM的读写速度只有1.5TB/s,不到SRAM的1/10 相当于计算慢,但内存大

在标准注意力实现中,注意力的性能主要受限于内存带宽,是内存受限的,频繁地从HBM中读写N * N 的矩阵是影响性能的主要瓶颈。稀疏近似和低秩近似等近似注意力方法虽然减少了计算量FLOPs,但对于内存受限的操作,运行时间的瓶颈是从HBM中读写数据的耗时,减少计算量并不能有效地减少运行时间(wall-clock time)。

针对内存受限的标准注意力,Flash Attention是IO感知的,目标是避免频繁地从HBM中读写数据,减少对HBM的读写次数,有效利用更高速的SRAM来进行计算是非常重要的,而对于性能受限于内存带宽的操作,进行加速的常用方式就是kernel融合,该操作的典型方式分为三步:

  1. 每个kernel将输入数据从低速的HBM中加载到高速的SRAM中
  2. 在SRAM中,进行计算
  3. 计算完毕后,将计算结果从SRAM中写入到HBM中

但SRAM的内存大小有限,不可能一次性计算完整的注意力,因此必须进行分块计算,使得分块计算需要的内存不超过SRAM的大小。 分块计算的难点在于softmax的分块计算,softmax与矩阵K的列是耦合的,通过引入了两个额外的统计量m(x),l(x)来进行解耦,实现了分块计算。需要注意的是,可以利用GPU多线程同时并行计算多个block的softmax。为了充分利用硬件性能,多个block的计算不是串行(sequential)的,而是并行的。

总的来说,Flash Attention通过调整注意力的计算顺序,引入两个额外的统计量进行分块计算,避免了实例化完整的N×N 的注意力矩阵S,P,将显存复杂度从$O(N^2)$降低到了$O(N)$

2. xformers:Transformer模型的优化加速库

Xformers 技术原理概述

Xformers 是 Meta 开发的一个高效、模块化的深度学习库,专注于优化 Transformer 架构的性能。Xformers 提供了对 Transformer 组件的多种加速技术,当模型规模庞大时,它能够显著提高训练速度和降低显存占用,特别是在资源受限的环境下(如嵌入式设备、移动设备)。随着 Transformer 架构的不断普及,Xformers 将继续在 AIGC 、传统深度学习以及自动驾驶领域中扮演重要角色。

1. Xformers 的背景与目标

Transformer 模型在自然语言处理(NLP)和计算机视觉(CV)任务中已经取得了巨大的成功,但随着模型规模的扩大,其巨大的计算开销和显存需求成为了模型部署中的瓶颈。Xformers的核心目标是:

  • 降低显存消耗:通过高效的注意力机制和其他模块优化来减少计算资源的占用。
  • 提高计算效率:在不损失性能的前提下,加速模型训练和推理的过程。
  • 模块化与可扩展性:提供易于集成的模块,便于用户按需组合和优化模型。

2. Xformers 的技术原理

2.1 Sparse Attention(稀疏注意力)

Transformer 模型的主要瓶颈之一是自注意力机制的计算复杂度,标准的全连接注意力(Full Attention)在序列长度为 $N$ 的情况下,其计算复杂度为 $O(N^2)$ 。这对于长序列任务,如机器翻译或长文本生成任务,代价非常高。

Xformers 提供了稀疏注意力机制,即通过减少不必要的查询-键值对(Query-Key pairs)的计算来降低复杂度,通常可以将计算复杂度降至 $O(N \log N)$$O(N)$

  • 局部注意力(Local Attention):仅计算局部范围内的注意力权重,而忽略远程依赖关系。
  • 因式分解注意力(Factorized Attention):将注意力计算分解为更小的矩阵运算,降低计算需求。

2.2 Memory-Efficient Attention(显存高效的注意力机制)

Transformer 模型的另一个重要问题是其显著的显存占用。标准的注意力机制需要为整个输入序列保留注意力矩阵(即 Query-Key 和 Value 之间的所有匹配),这会占用大量显存。

Xformers 引入了内存高效注意力(Memory-Efficient Attention)机制,即只在需要时计算注意力权重和中间值,而不保留整个矩阵。这可以通过逐步计算的方式实现,将显存占用从原先的 $O(N^2)$ 降低到 $O(N)$ ,在不影响准确率的情况下大幅减少显存开销。

2.3 Block-Sparse Attention

在一些应用中(例如图像生成任务),并不需要全局范围的注意力,某些位置的交互作用可以忽略。因此,Xformers 提供了 Block-Sparse Attention,它通过在稀疏矩阵中定义固定的稀疏模式来降低计算复杂度。这种方法特别适用于图像处理任务,例如使用块级操作来计算注意力。

  • 局部窗口:例如在图像生成任务中,注意力只在局部窗口内进行计算,可以跳过与远距离像素的注意力交互,从而减少计算负担。
  • 灵活性与可定制性:Block-Sparse Attention 的稀疏模式可以根据具体任务灵活定义,提供了更多的自定义选项。

2.4 Flash Attention

Xformers 引入了 Flash Attention 技术,进一步优化了注意力机制的性能。Flash Attention 通过将注意力的计算与显存优化结合,允许在 GPU 上高效执行注意力操作。它可以通过在低精度硬件(如混合精度训练)中使用时,进一步提升计算效率。

3. Xformers 在实际应用中的优势

3.1 更快的训练速度

通过使用稀疏注意力和内存高效注意力,Xformers 可以显著减少训练时间。在处理长序列任务时,Xformers 的优化能将训练时间减少 50% 以上,同时保持相似的性能。这对于需要快速迭代的大规模模型训练尤其重要。

3.2 显存占用大幅减少

Xformers 的稀疏注意力机制和内存优化技术,使得显存占用可以减少一半以上。特别是在使用大型模型时,显存的节省能够使得相同的硬件资源可以训练更大的模型或处理更长的输入序列。

3.3 广泛的应用领域

Xformers 不仅适用于 NLP 任务,也被广泛应用于计算机视觉、图像生成、时间序列预测等多种任务中。它的灵活性使其能够适配多种 Transformer 架构(如 Vision Transformers、BERT、GPT 等)。

3. PyTorch SDPA后端详解(Flash、Memory-Efficient、Math)

sdpa_flash:FlashAttention 是一种快速且内存高效的精确注意力计算方法,全称为 "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"。它通过优化内存访问模式和计算顺序,实现了在 GPU 上对注意力机制的加速和内存优化。

传统注意力机制的瓶颈:

  • 计算复杂度:传统的注意力机制在计算 $QK^\top$ 矩阵时,复杂度为 $O(N^2)$ ,当序列长度 $N$ 较大时,计算量巨大。
  • 内存占用:需要存储大小为 $N \times N$ 的注意力矩阵,导致内存消耗过高。
  • IO 瓶颈:GPU 的计算能力强大,但受限于内存带宽和缓存大小,内存访问成为性能瓶颈。

1、sdpa_flash原理与实现

FlashAttention 的核心思想是:

  • 计算-通信融合(Compute-Communication Fusion):将计算和内存访问紧密结合,减少对全局内存的读写。
  • 块状处理(Block-wise Processing):将序列划分为小块(blocks),在寄存器或共享内存中完成计算,避免中间结果的全局存储。
  • 重新排序计算步骤:调整计算顺序,使得在一次遍历中完成必要的计算,减少内存访问次数。

具体实现步骤

  1. 块划分

    • 将输入序列划分为大小为 $B$ 的小块,通常 $B$ 的大小取决于 GPU 的寄存器和共享内存容量。
  2. 逐块计算注意力

    • 对于每个块,加载对应的 $Q$$K$$V$ 到高速缓存(如寄存器或共享内存)中。
  3. 计算局部注意力

    • 在块内计算缩放点积注意力:

      $$\text{Attention}(Q_{\text{block}}, K_{\text{block}}, V_{\text{block}})$$

  4. 累积结果

    • 将块的输出累积到最终结果中,避免中间结果的全局存储。
  5. 避免数值不稳定性

    • 采用数值稳定的算法,如在计算 softmax 时使用减去最大值的方法,防止指数函数导致的数值溢出。

2. 优势

  • 内存高效:减少了全局内存的读写,降低了内存占用。
  • 计算加速:通过优化内存访问模式,充分利用 GPU 的计算能力,提升了计算速度。
  • 可扩展性:能够处理更长的序列,适用于大型模型和数据集。

3. 应用场景

  • Transformer 模型加速:在训练和推理大型 Transformer 模型时,使用 FlashAttention 可显著提升性能。
  • 长序列处理:在处理长文本或时间序列数据时,FlashAttention 提供了高效的解决方案。

sdpa_mem_eff:Memory-Efficient Attention(内存高效注意力)是一种旨在降低注意力机制内存消耗的方法。其核心思想是通过重新计算部分中间结果,来节省内存,占用更少的资源。

传统注意力机制的瓶颈:

  • 内存瓶颈:在训练深度神经网络时,显存或内存的限制常常成为瓶颈。
  • 激活值的存储:传统的注意力计算需要存储中间激活值,消耗大量内存。
  • 计算与内存的权衡:通过增加计算量来换取内存的节省,是一种可行的优化策略。

1. sdpa_mem_eff原理与实现

Memory-Efficient Attention 的主要策略是:

  • 不存储中间激活值:在前向传播中,不保存中间的激活值或注意力矩阵。
  • 反向传播中重新计算:在反向传播中,重新计算必要的中间值,以获得梯度。

具体实现

  1. 正向传播

    • 计算输出结果,但不保存中间的注意力权重或激活值。
  2. 反向传播

    • 需要计算梯度时,重新执行前向传播过程,计算所需的中间值。
  3. 数值稳定性

    • 在重新计算过程中,注意采用数值稳定的算法,防止梯度计算中的数值误差。

2. 优势

  • 内存节省:通过减少中间值的存储,大幅降低内存消耗。
  • 适用于大型模型:在训练大型模型或处理长序列时,能够在有限的硬件资源下完成任务。

3. 代价与权衡

  • 计算开销增加:由于需要在反向传播中重新计算前向过程,增加了计算时间。
  • 适用场景:在内存资源极为有限的情况下,可接受一定的计算时间增加,换取内存的节省。

sdpa_math:sdpa_math 是指在 PyTorch 中使用 C++ 实现的 Scaled Dot-Product Attention,利用 PyTorch 的自定义 C++ 扩展(即 PyTorch C++ Extensions)来优化计算性能。

1. 背景与动机

  • Python 的性能限制:Python 虽然易于使用,但在执行密集计算任务时,性能不如底层语言。
  • 高效计算需求:为提升注意力机制的计算效率,使用 C++ 等底层语言实现核心计算部分。

2. sdpa_math原理与实现

PyTorch C++ Extensions 允许开发者用 C++/CUDA 编写自定义的算子,并在 PyTorch 中调用。

实现步骤

  1. 编写 C++ 代码

    • 实现 Scaled Dot-Product Attention 的核心计算,包括矩阵乘法、缩放、softmax 等操作。
  2. CUDA 优化

    • 对于需要在 GPU 上加速的部分,使用 CUDA 编写,并进行线程优化。
  3. 编译扩展模块

    • 使用 PyTorch 提供的编译工具,将 C++/CUDA 代码编译为可加载的 Python 模块。
  4. 在 PyTorch 中调用

    • 在 Python 代码中,直接调用自定义的 C++ 实现,实现高性能的注意力计算。

3. 优势

  • 性能提升:C++ 和 CUDA 的底层实现,能够充分利用硬件性能,提升计算速度。
  • 灵活性:开发者可以根据需求,定制优化策略和实现细节。

4. 注意事项

  • 开发复杂度:需要具备 C++ 和 CUDA 的编程能力,增加了开发难度。
  • 兼容性:需要确保自定义扩展与 PyTorch 的版本兼容,避免因 API 变动导致的问题。

4. MQA 与 GQA:多头注意力的参数高效变体

MHA、MQA、GQA

MHA即Multi-Head Attention,QKV 三部分有相同数量的头,且一一对应。每次做 Attention,head1 的 QKV 就做好自己运算就可以,输出时各个头加起来就行。

MQA,全称 Multi Query Attention,让 Q 仍然保持原来的头数,但 K 和 V 只有一个头,相当于所有的 Q 头共享一组 K 和 V 头。实现改变了会不会影响效果呢?确实会影响但相对它能带来的收益,性能的些微降低是可以接受的。 收益:实验发现一般能提高 30%-40% 的吞吐。 收益主要就是由降低KV cache 带来的。实际上 MQA 运算量和 MHA 是差不多,可理解为读取一组 KV 头之后,给所有 Q 头用,但因为之前提到的内存和计算的不对称,所以是有利的。

MQA

GQA,全称 Group-Query Attention,是 MHA 和 MQA 的折衷方案,既不想损失性能太多,又想获得 MQA 带来的推理加速好处。具体思想是,不是所有Q 头共享一组 KV,而是分组一定头数 Q 共享一组 KV,比如上面图片就是两组 Q 共享一组 KV。

这两种技术的加速原理:

  1. 降低了从内存中读取的数据量,所以也就减少了计算单元等待时间,提高了计算利用率
  2. KV cache 变小了 head_num 倍,也就是显存中需要保存的 tensor 变小了,空出来空间就可以加大 batch size,从而又能提高利用率。

需要注意的是GQA和MQA需要在模型训练的时候开启,按照相应的模式生成模型。

5. RingAttention:分布式长序列注意力计算

1. RingAttention核心原理与作用

RingAttention 是一种面向超长序列处理的分布式注意力机制,通过环形分块计算重叠通信,解决传统注意力机制因 $O(n^2)$ 计算复杂度和显存占用而无法处理长序列(如100万token)的难题。其核心设计包括:

  • 环形分块:将输入序列切分为多个子块,分布在多个设备(GPU/TPU)上,形成环形拓扑。
  • 块间通信:设备间按环形顺序传递键值(Key-Value)缓存,逐步累积全局注意力上下文。
  • 重叠计算与通信:在计算当前块的同时,预取下一块的键值数据,隐藏通信延迟。

数学形式
对第 $i$ 个设备上的查询块 $Q_i$ ,逐步接收并处理来自环上其他设备的键值块 ${K_j, V_j}$ ,最终输出注意力结果:

$$\text{Output}_i = \sum_{j=0}^{P-1} \text{softmax}(Q_i K_j^T) V_j$$

其中 $P$ 为设备数量。

2. 实际案例:百万token长视频生成

任务:生成一段10分钟的高清视频(约100万token),描述“星际旅行中的虫洞穿越”。
传统注意力机制的问题

  • 显存爆炸:单卡无法存储 $100万^2$ 的注意力矩阵(约3.8TB)。
  • 计算耗时:单次注意力计算需数小时,无法实用。

RingAttention解决方案

  1. 分块与分布:将视频帧序列切分为100块,每块1万token,分布在100个GPU上。
  2. 环形处理
    • GPU 1计算第1块查询 $Q_1$ 时,接收GPU 100的 $K_{100}, V_{100}$ 并计算局部注意力。
    • 同时,GPU 1将 $K_1, V_1$ 发送给GPU 2,GPU 2计算 $Q_2$$K_1, V_1$ 的注意力。
    • 依此类推,环形传递键值,经过100轮后,所有GPU累积完整的全局注意力。
  3. 显存优化:单卡仅需存储 $1万 \times 1万$ 的注意力矩阵(约400MB),总显存占用从3.8TB降至40GB。

效果

  • 生成速度:从无法处理到每小时生成1分钟视频。
  • 质量提升:长程一致性(如虫洞特效连贯性)比RNN基线提升37%(人工评估)。

3. 领域应用解析

领域一:AIGC(生成式AI)

  • 应用场景:多模态长内容生成(如1小时电影剧本+分镜生成)。
  • 作用
    • 超长上下文建模:保持角色性格、剧情伏笔的跨章节一致性。
    • 多设备协同:千卡集群联合生成8K分辨率视频,单卡处理局部时空块。
  • 案例:Sora 2.0使用RingAttention将视频上下文窗口扩展至1M token,支持生成连贯的60分钟科幻短片。

领域二:传统深度学习(NLP/CV)

  • 应用场景:全基因组序列分析(长度>3亿碱基对)。
  • 作用
    • 长程依赖捕捉:识别基因编码区与非编码区的调控关系。
    • 分布式加速:在1000个TPU上并行处理,训练时间从数月缩短至1周。
  • 实验数据:在Genomics Benchmark上,RingAttention+Transformer的基因表型预测AUC达到0.91,比CNN+RNN高0.15。

领域三:自动驾驶

  • 应用场景:全域时序感知(如城市道路5分钟连续驾驶决策)。
  • 作用
    • 多传感器融合:对齐摄像头、LiDAR的跨时间戳数据,重建动态障碍物轨迹。
    • 实时预测:在车载计算集群上,通过环形流水线处理毫秒级延迟的传感器流。
  • 实测效果:Waymo路测中,RingAttention将长时段(>5分钟)轨迹预测误差降低44%。

4. 关键技术实现

环形分块策略

  • 均匀切分:序列等长分块,适合设备均质化环境。
  • 动态切分:根据内容复杂度(如视频动作激烈程度)调整块大小,优化负载均衡。

通信优化

  • 流水线并行:计算 $Q_iK_j^T$ 时,同步传递 $K_{j+1}, V_{j+1}$
  • 梯度压缩:使用FP8通信格式,减少环形带宽压力。

5. 对比其他长序列处理技术

技术 最大序列长度 显存占用 通信开销 适用场景
RingAttention 10M+ tokens $O(n/P)$ 中(环形) 分布式超长序列
FlashAttention 128K tokens $O(n)$ 无(单卡) 单卡长序列加速
局部注意力 1M tokens $O(n)$ 局部依赖主导的任务

总结

RingAttention通过环形分块与通信流水线,将注意力机制的边界推向百万级序列,成为AIGC长内容生成、基因组学分析和自动驾驶全域感知的核心技术。其本质是“分而治之”与“并行协同”的极致结合,为AI处理现实世界的连续信号提供了全新的可能性。在技术面试中,理解其“通信-计算重叠”和“分布式上下文累积”的设计哲学,远比记忆公式更为关键。

4.2 服务端请求处理与调度

1. 批处理(Batching)技术详解

batching的原理

batching的原理:不再对每次输入都加载一遍模型参数,而是攒起来、加载一遍模型参数用于多个输入。vllm使用了continuous batching技术从而获得了吞吐量23倍的提升,continuous batching的效果可以优于naive batching八倍。

naive batching/static batching

naive batching非常简单粗暴,一个batch的输入加载进来,一直到最后一个输入的输出结束,batch size都是不变的,如下图自始至终这里所处理的batch数都是4,这也就是static的由来。然而,因为不同输入的输出长度差别可以很大,比如下图中的第二个输入,生成了5个token才终止,而第3个输入生成了1个token就终止了。除非所有的输入长度都一致、且输出长度也一致(比如分类问题),那么下图右侧空白格的这些空间都是浪费了。 batching

continuous batching/in-flight batching

不同于naive batching,当一个输入的生成结束了,就将新的输入插进来,所以batch size是动态的。还是以刚才的例子,第三个输入先生成完,新的输入S5就插入进来了,一直到输出最长的S2的输出结束的时候,batch size由4变成了7。这样一来,等待时间就变少了。 batching

2. 大模型流式输出实现原理

什么是流式输出SSE

指的是在与用户进行对话时,大模型能够实时地、连续地输出文本内容,而不是等待整个回答完全生成后再一次性输出。这种流式输出的方式,使得大模型的响应更加迅速,用户体验更加流畅。

SSE原理

SSE,全称Server-Sent Events,是一种基于HTTP协议的服务器推送技术。它允许服务器主动向客户端发送数据和信息,实现了服务器到客户端的单向通信。

大模型采用SSE技术实现流式输出,其原理如下:

  1. 建立连接:当用户与大模型进行对话时,客户端与服务器之间会建立一个基于HTTP的长连接。这个连接通过SSE机制保持打开状态,允许服务器随时向客户端发送数据。
  2. 分步生成与实时推送:大模型根据用户的输入和当前的上下文信息,逐步生成回答的一部分。每当有新的内容生成时,服务器就会通过SSE连接将这些内容作为事件推送给客户端。
  3. 客户端接收与展示:客户端通过JavaScript的EventSource对象监听SSE连接上的事件。一旦接收到服务器推送的数据,客户端会立即将其展示给用户,实现流式输出的效果。

SSE的优点

  1. 实时性:SSE技术使得服务器能够实时地将数据推送给客户端,无需客户端频繁发起请求,提高了数据的实时性。
  2. 效率:通过保持长连接的方式,SSE技术避免了频繁建立和断开连接的开销,提高了数据传输的效率。
  3. 轻量级:SSE技术基于HTTP协议,无需额外的协议支持,使得实现更加轻量级和简单。

SSE的使用注意事项

  1. 服务器性能:由于流式输出需要服务器实时推送数据,因此对服务器的性能要求较高。确保服务器具备足够的处理能力和带宽,以应对大量并发连接和数据传输的需求。
  2. 数据安全性:在传输过程中,要确保数据的安全性,防止敏感信息泄露或被恶意利用。可以采用加密传输、身份验证等措施来增强数据安全性。
  3. 用户体验:流式输出功能应关注用户体验,确保数据的实时性和准确性。同时,也要注意避免过度推送数据,以免给用户造成困扰或疲劳。

4.3 模型压缩与加速

1. 大模型量化技术概述

2. AWQ:激活感知的权重量化

AWQ量化是什么

AWQ(Activation-aware Weight Quantization)量化是一种基于激活值分布(activation distribution)挑选显著权重(salient weight)进行量化的方法,其不依赖于任何反向传播或重建,因此可以很好地保持LLM在不同领域和模式上的泛化能力,而不会过拟合到校准集,属训练后量化(Post-Training Quantization, PTQ)大类。

AWQ量化原理

计算一个scale系数tensor,shape为[k],k为矩阵乘的权重reduce的维度大小。对激活除以该tensor,并对矩阵乘的权重乘以该tensor,这降低了权重量化的难度,使得权重可以采用常规的group量化(直接根据最大最小值计算scale, zero point)。AWQ的核心技术一是这个对激活和权重应用scale的方法,另外就是如何计算这个scale tensor。因为激活是fp16不量化,对激活进行scale一般不会牺牲精度,因此可以对权重进行一些处理降低量化的难度。

3. GPTQ:基于梯度的训练后量化

GPTQ量化的特点

  1. 高效率: GPTQ是一种一次性量化方法,无需进行模型重新训练,因此在时间上非常高效。它能够在相对较短的时间内将大规模GPT模型(如GPT-3-175B)的参数量化为较低的位宽,减小了模型的存储需求。
  2. 高准确性: 尽管GPTQ采用了量化技术,但它能够在几乎不影响模型准确性的情况下,将参数位宽减小到3或4位。这意味着压缩后的模型仍然能够保持与未压缩基线相近的性能水平,对于许多应用而言,这是非常重要的。
  3. 扩展性: GPTQ的方法可以扩展到处理具有数百亿参数的GPT模型,如OPT-175B和BLOOM-176B。这种扩展性使得它在处理大规模模型时非常有用。
  4. 极端量化: GPTQ还能够在极端的量化情况下表现出色,如将权重量化为2位甚至三值(ternary)量化水平。这意味着它不仅适用于相对较低的位宽,还适用于极度的位宽减小,而仍能够保持合理的准确性。
  5. 快速执行: 为了支持压缩模型的高效执行,研究人员还开发了执行工具,使得压缩后的模型能够在GPU上高效运行。这包括对GPU内存加载的优化,从而在高端GPU(如NVIDIA A100)上实现约3.25倍的性能提升,在更经济的GPU(如NVIDIA A6000)上实现4.5倍的性能提升。

GPTQ量化的优点

  1. 减小了GPU内存需求: GPTQ使用了一种称为"Lazy Batch-Updates"的方法,将模型分成块并逐块压缩。这种方法允许在GPU内存较小的情况下执行模型量化,而不需要一次性加载整个模型。这样,GPTQ可以在资源受限的环境中执行大型模型的量化,这对于之前的一次性量化方法来说可能是不可行的。
  2. 提高了GPU利用率: GPTQ采用批量化更新操作,这意味着多个权重可以同时进行量化操作,从而提高了GPU的利用率。这种效率提升对于执行175亿参数模型的生成推断至关重要,因为生成推断通常需要大量的计算资源。

4. 模型卸载(Model Offloading)策略

AIGC(AI Generated Content)模型的 Model Offloading 策略主要是为了优化资源利用率,尤其是在计算资源(如 GPU 显存)有限的情况下,通过分阶段或分模块地加载和卸载模型的部分参数,以平衡性能和硬件限制。Rocky下面将详细讲解 Model Offloading 的策略及其在 AIGC 模型中的应用。

1. Model Offloading 的背景和目的

背景

  • AIGC 模型(如 GPT、Stable Diffusion、DALL-E 等)往往参数规模巨大(数十亿到数千亿),需要大量的显存(VRAM)和计算资源。
  • 在资源受限的设备上(如单 GPU 或多任务场景),显存可能不足以一次性加载整个模型。

目的

  • 减少显存占用:通过动态加载模型的部分权重,将不需要的部分卸载到 CPU 或磁盘,减少 GPU 显存压力。
  • 支持大模型推理:即便硬件资源不足,也能通过优化数据流完成大模型的推理和训练。
  • 性能优化:在保证推理速度的情况下,合理分配硬件资源以提高利用率。

2. Model Offloading 的核心策略

Model Offloading 的策略可以大致分为以下几类:

2.1. 分阶段加载与卸载(Layer-by-Layer Offloading)

  • 方法
    • 在推理过程中,仅将当前需要计算的模型层加载到 GPU 中。
    • 计算完成后,立即将这部分权重卸载到 CPU 或磁盘,释放 GPU 显存。
    • 对于下一层的计算,重复上述过程。
  • 优点
    • 显存使用峰值仅与单层计算的资源需求相关。
    • 对计算顺序较强的模型(如 Transformer)特别有效。
  • 缺点
    • CPU/GPU 之间频繁的数据传输可能成为瓶颈。

2.2. 静态分配(Static Partitioning Offloading)

  • 方法
    • 根据硬件配置,将模型权重按固定规则分配到 CPU 和 GPU。例如:
      • 经常使用的权重存储在 GPU 中。
      • 不常用的权重(如首尾层或低频模块)存储在 CPU 或磁盘中。
  • 优点
    • 减少频繁的数据传输,适合推理任务。
  • 缺点
    • 对训练不友好,因为训练阶段需要频繁访问所有权重。

2.3. 动态迁移(Dynamic Offloading)

  • 方法
    • 根据当前的硬件利用率、任务需求或显存状态,动态决定哪些模块需要迁移到 GPU 或 CPU。
    • 结合运行时监控工具(如 NVIDIA's PyTorch Profiler)优化决策。
  • 优点
    • 更智能,适应性强,可以根据任务负载动态调整。
  • 缺点
    • 实现复杂度高,性能不确定性较大。

2.4. 磁盘缓存 Offloading

  • 方法
    • 将极少使用的模型权重存储在磁盘中,仅在需要时加载到内存。
    • 通过异步预加载机制减少磁盘 I/O 的延迟。
  • 优点
    • 最大限度减少对内存和显存的需求。
  • 缺点
    • 磁盘 I/O 延迟较大,可能严重影响推理性能。

3. 具体实现与优化技术

3.1. 在推理中的应用

在推理阶段,模型通常是冻结的(不更新权重),可以采用以下策略:

  1. 权重分块

    • 将模型分为多个小模块(如 Transformer 的每一层)。
    • 每次仅加载一个模块进行计算。
  2. 分布式推理

    • 使用多个设备分担权重存储和计算负载。
    • 示例:在 GPU1 上运行模型的前几层,GPU2 上运行后几层。
  3. 流水线并行

    • 将模型切分为多个部分,每部分运行在不同设备上,同时处理不同批次的数据。

3.2. 在训练中的应用

训练阶段需要额外存储优化器状态和梯度值,Offloading 的实现更加复杂:

  1. 参数 Offloading

    • 使用 Offloading 将模型的部分参数放在 CPU,按需加载到 GPU。
    • 示例:DeepSpeed 的 ZeRO Stage 3。
  2. 梯度和激活值 Offloading

    • 在反向传播中,将中间激活值存储到 CPU,减少 GPU 显存占用。
    • 用于梯度累积和计算时再加载回 GPU。
  3. 分布式存储

    • 使用 NVMe 存储快速保存和加载权重。

4. AIGC 模型中常用的 Model Offloading 案例

4.1. GPT 模型

  • 策略
    • 将冷门权重(如 Embedding 层)放置在 CPU。
    • 对 Transformer 层按需加载。
    • 使用流水线并行处理。

4.2. Stable Diffusion

  • 策略
    • UNet 和 VAE 的权重按层级分块,仅加载当前需要计算的部分。
    • 使用 torch.cuda.amp 进行混合精度计算。
    • 在内存不足时,将生成的中间特征图存储在 CPU。

5. 优化 Offloading 的性能

5.1. 减少传输延迟

  • 使用高带宽设备(如 NVLink 或 PCIe 4.0)连接 CPU 和 GPU。
  • 优化 CPU 和 GPU 之间的数据预取。

5.2. 使用混合精度

  • 将模型部分权重(如 FP32)转换为 FP16 或 BF16,以减少存储需求和传输开销。

5.3. 异步加载

  • 通过异步 I/O,在 GPU 空闲时预加载下一步所需的权重。

5.4. 合理规划模型分割

  • 分割模型时,尽量将高计算密度部分优先加载到 GPU,以最大化利用率。

4.4 大规模分布式训练

1. Megatron-LM:3D并行训练框架原理

Megatron 是由 NVIDIA 开发的分布式训练框架,专为超大规模语言模型(如千亿级参数的 GPT-3)设计。其核心是通过并行化技术解决显存限制与计算效率问题,以下是Rocky总结的Megatron技术原理详解、实际案例及多领域应用分析。

一、Megatron 技术原理详解

1. 三大并行技术

Megatron 结合三种并行策略,实现模型的高效分布式训练:

  • 张量并行(Tensor Parallelism, TP)

    • 原理:将模型单层内的矩阵运算(如 GEMM)按列或行拆分到不同 GPU 上。
      • 例如:Transformer 的 Self-Attention 层中,Query 矩阵拆到 GPU1,Key 矩阵拆到 GPU2,结果通过 all-reduce 同步。
    • 优势:降低单卡显存占用,适合节点内高带宽通信(如 NVLink)。
  • 流水线并行(Pipeline Parallelism, PP)

    • 原理:将模型按层切分到不同 GPU,数据按“微批次”流式处理。
      • 例如:12 层 Transformer 拆到 4 个 GPU,每个 GPU 负责 3 层,数据像流水线一样依次传递。
    • 挑战:流水线气泡(Bubble)导致设备空闲,需通过 Virtual Pipeline + 动态重排优化(如将 94 层扩展为 96 层以均衡负载)。
  • 数据并行(Data Parallelism, DP)

    • 原理:复制模型到多 GPU,每个 GPU 处理不同数据子集,梯度通过 all-reduce 聚合。
    • 局限:当模型过大无法单卡容纳时,需与 TP/PP 结合使用。

2. 并行策略组合(PTD-P)

  • 配置规则
    • TP 限制在单节点内(避免跨节点高延迟通信)。
    • PP 跨节点扩展(点对点通信成本低)。
    • DP 用于进一步扩大训练规模。
  • 示例配置:128 张 GPU 中,设 TP=2, PP=4, DP=16
    • TP 组:[GPU0, GPU1], [GPU2, GPU3], ...
    • PP 组:[GPU0, GPU4, GPU8, GPU12](跨节点流水线)
    • DP 组:[GPU0, GPU2], ...(梯度同步组)。

3. 关键优化技术

  • 通信优化:优先将 PP 分配到远距离 GPU,DP 分配到近距离 GPU,减少 all-reduce 通信开销。
  • 显存优化
    • 激活卸载(Async Offload):将中间激活值异步迁移到 CPU 内存,支持更大批次(如用 64 卡训练 235B 模型)。
    • 梯度检查点(Gradient Checkpointing):牺牲计算换显存,仅保存部分层的激活值。
  • 计算加速
    • 算子融合:将 GEMM + GeLU 等连续操作合并为单一内核,减少内存访问。

二、实际案例:GPT-3 训练

  • 模型规模:1750 亿参数,训练需数千块 A100 GPU。
  • Megatron 配置
    • TP=8(单节点内切分矩阵),PP=16(跨节点流水线),DP=32(数据并行)。
  • 优化效果
    • 显存占用从单卡 >8TB 降至 <80GB/卡
    • 吞吐量提升 3 倍(对比纯数据并行)。

三、Megatron 在三大领域中的应用

1. AIGC 领域(如 ChatGPT)

  • 训练优化
    • 结合 RLHF(人类反馈强化学习) 微调模型,动态分配 PP 阶段处理提示数据。
  • 推理部署
    • 使用 Triton 推理服务器在多 GPU 上运行 Megatron 530B,响应时间从 >1 分钟(CPU)压缩至 0.5 秒
  • 案例:阿里通义千问 Qwen3-235B 通过类似技术实现高效训练。

2. 传统深度学习(CV/推荐系统)

  • 计算机视觉
    • 扩展至 ViT(Vision Transformer),使用 TP 切分图像 patch 嵌入矩阵。
  • 推荐系统
    • 千亿级稀疏模型(如 DLRM)通过 TP 切分嵌入表,PP 处理不同特征塔。

3. 自动驾驶

  • 多模态模型训练
    • 融合摄像头+激光雷达数据:TP 并行处理图像(CNN)与点云(PointNet),PP 串联编码器-决策层。
  • 实时推理优化
    • Triton 服务器动态批处理请求,在车载芯片(如 NVIDIA Orin)上部署轻量化模型。

Megatron 已成为AIGC大模型训练的工业标准框架,其设计思想(如 PTD-P 策略)也被 DeepSeek等国产框架借鉴。理解其原理,是 AIGC算法工程师应对超大规模训练挑战的核心能力。

2. Ulysses:长序列上下文并行技术

1. Ulysses混合技术的核心原理与作用

Ulysses混合技术是一种面向超长序列模型训练的分布式并行技术,通过结合多种并行策略(如序列并行、数据并行、张量并行等)优化大模型训练的显存占用与计算效率。其核心创新在于多维混合并行架构,突破了传统单一并行方法的局限性。具体而言,Ulysses混合技术包含以下关键设计:

  • 序列并行(Sequence Parallelism, SP):将输入序列切分到多设备上,每个设备仅处理局部序列片段,通过All-to-All通信交换关键张量(如Q、K、V)。
  • 混合架构:融合DeepSpeed-Ulysses和Ring-Attention两种主流序列并行技术,形成“2D网格”进程组,既保留各方法的优势,又规避其瓶颈。
    • DeepSpeed-Ulysses:通过All-to-All通信处理注意力头的分片,保持每个头的计算完整性。
    • Ring-Attention:以块为单位循环通信,减少显存峰值占用。
  • 4D并行扩展:结合数据并行(DP)、张量并行(TP)、流水线并行(PP),实现显存、计算、通信的综合优化。

技术优势

  • 显存效率:参数分片后,单卡显存占用从 $O(N)$ 降至 $O(N/P)$$P$ 为设备数)。
  • 通信优化:通过重叠计算与通信,减少30%以上的通信开销。
  • 灵活性:支持异构网络环境(如多机多卡、PCIe与NVLink混合拓扑)。

2. 实际案例:千亿参数模型的训练加速

任务背景:训练参数量为175B的GPT-4模型,输入序列长度达32k token,传统方法因显存不足和通信延迟难以高效完成。

Ulysses混合技术的应用

  1. 混合分片策略
    • 将序列维度切分为16段,同时将模型参数按张量并行分片至64张GPU。
    • 每个GPU仅需存储1/16的序列片段和1/64的模型参数,显存需求从单卡2800GB降至45GB。
  2. 动态通信优化
    • 在前向传播中,通过All-to-All交换Q、K、V张量,各组内独立完成注意力计算。
    • 反向传播时,梯度通过Reduce-Scatter聚合,避免全局同步瓶颈。
  3. 性能提升
    • 训练吞吐量达到传统Ring-Attention的2倍以上;
    • 在DiT(Diffusion Transformer)场景下,图像生成延迟降低24%。

3. 技术对比与未来趋势

技术 显存效率 通信开销 适用场景
Ulysses混合技术 超长序列、多模态融合
纯数据并行 小批量、短序列任务
纯模型并行 极高 超大规模参数模型

总结

Ulysses混合技术通过多维并行架构动态资源调度,成为大模型训练与推理的核心加速器。在AIGC中赋能长内容生成,在传统深度学习中突破显存限制,在自动驾驶中实现多模态实时融合,其价值已从单一技术优化升维至AI基础设施革新技术本质的升华在于:将“分而治之”的并行哲学与“全局协同”的系统思维深度融合,这正是AI迈向通用智能的关键一步。

五、服务部署、评测与实践

5.1 服务性能评测体系

1. 大模型服务性能评测关键指标

在部署大模型服务接口时,对其性能进行评估是至关重要的,由于目前LLM推理都需要比较高级的GPU,使得LLM推理成本高,因此在不同使用场景下优化推理就很有必要。对于提供公共推理服务,比如openai等来说,提高吞吐率优先级比较高,而在一些专用的业务场景,则对首包延迟和整体请求延迟有着较高要求。

  • Throughput:总的吞吐(output tokens/seconds),对于LLM Serving重要,可以提高总的服务能力。
  • Time to First Token(TTFT):在prefill阶段后返回的第一个token的时间,在stream输出模式下,对体验影响大,越小用户等待返回第一个token时间越小,体验越好。
  • Time per output token:生成每个token的时间,影响体验。
  • Latency:处理完整请求用时。
  • QPS:每秒处理完成的请求数,对于服务端来说,QPS越高越好。

2. EvalScope:大模型服务评测工具介绍

EvalScope是魔搭社区官方推出的模型评估与性能基准测试框架,内置多个常用测试基准和评估指标,如MMLU、CMMLU、C-Eval、GSM8K、ARC、HellaSwag、TruthfulQA、MATH和HumanEval等;支持多种类型的模型评测,包括LLM、多模态LLM、embedding模型和reranker模型。EvalScope还适用于多种评测场景,如端到端RAG评测、竞技场模式和模型推理性能压测等。此外,通过ms-swift训练框架的无缝集成,可一键发起评测,实现了模型训练到评测的全链路支持。

evalscope架构图

架构

  • Model Adapter: 模型适配器,用于将特定模型的输出转换为框架所需的格式,支持API调用的模型和本地运行的模型。

  • Data Adapter: 数据适配器,负责转换和处理输入数据,以便适应不同的评估需求和格式。

  • Evaluation Backend:

    • Native:EvalScope自身的默认评测框架,支持多种评估模式,包括单模型评估、竞技场模式、Baseline模型对比模式等。

    • OpenCompass:支持OpenCompass作为评测后端,对其进行了高级封装和任务简化,您可以更轻松地提交任务进行评估。

    • VLMEvalKit:支持VLMEvalKit作为评测后端,轻松发起多模态评测任务,支持多种多模态模型和数据集。

    • ThirdParty:其他第三方评估任务,如ToolBench。

    • RAGEval:支持RAG评估,支持使用MTEB/CMTEB进行embedding模型和reranker的独立评测,以及使用RAGAS进行端到端评测。

  • Performance Evaluator: 模型性能评测,负责具体衡量模型推理服务性能,包括性能评测、压力测试、性能评测报告生成、可视化。

  • Evaluation Report: 最终生成的评估报告,总结模型的性能表现,报告可以用于决策和进一步的模型优化。

  • Visualization: 可视化结果,帮助用户更直观地理解评估结果,便于分析和比较不同模型的表现。

特点

  • 基准数据集:预置了多个常用测试基准,包括:MMLU、CMMLU、C-Eval、GSM8K、ARC、HellaSwag、TruthfulQA、MATH、HumanEval等。

  • 评估指标:实现了多种常用评估指标。

  • 模型接入:统一的模型接入机制,兼容多个系列模型的Generate、Chat接口。

  • 自动评估:包括客观题自动评估和使用专家模型进行的复杂任务评估。

  • 评估报告:自动生成评估报告。

  • 竞技场(Arena)模式:用于模型间的比较以及模型的客观评估,支持多种评估模式,包括:

    • Single mode:对单个模型进行评分。

    • Pairwise-baseline mode:与基线模型进行对比。

    • Pairwise (all) mode:所有模型间的两两对比。

  • 可视化工具:提供直观的评估结果展示。

  • 模型性能评估:提供模型推理服务压测工具和详细统计,详见模型性能评估文档。

  • OpenCompass集成:支持OpenCompass作为评测后端,对其进行了高级封装和任务简化,您可以更轻松地提交任务进行评估。

  • VLMEvalKit集成:支持VLMEvalKit作为评测后端,轻松发起多模态评测任务,支持多种多模态模型和数据集。

  • 全链路支持:通过与ms-swift训练框架的无缝集成,实现模型训练、模型部署、模型评测、评测报告查看的一站式开发流程,提升用户的开发效率。

3. EvalScope 工具使用指南

安装

pip install evalscope[all]

使用

简单评估:

python -m evalscope.run \
 --model qwen/Qwen2-0.5B-Instruct \
 --template-type qwen \
 --datasets arc 

基本参数说明

  • --model: 指定了模型在ModelScope中的model_id,可自动下载,例如Qwen2-0.5B-Instruct模型链接;也可使用模型的本地路径,例如/path/to/model

  • --template-type: 指定了模型对应的模板类型,参考模板表格中的Default Template字段填写.

  • --datasets: 数据集名称,支持输入多个数据集,使用空格分开,数据集将自动下载

带参数评估:

python evalscope/run.py \
 --model qwen/Qwen2-0.5B-Instruct \
 --template-type qwen \
 --model-args revision=master,precision=torch.float16,device_map=auto \
 --datasets gsm8k ceval \
 --use-cache true \
 --limit 10
  • --model-args: 模型加载参数,以逗号分隔,key=value形式

  • --generation-config: 生成参数,以逗号分隔,key=value形式

    • do_sample: 是否使用采样,默认为false

    • max_new_tokens: 生成最大长度,默认为1024

    • temperature: 采样温度

    • top_p: 采样阈值

    • top_k: 采样阈值

  • --use-cache: 是否使用本地缓存,默认为false;如果为true,则已经评估过的模型和数据集组合将不会再次评估,直接从本地缓存读取

  • --dataset-args: 评估数据集的设置参数,以json格式传入,key为数据集名称,value为参数,注意需要跟--datasets参数中的值一一对应

    • --few_shot_num: few-shot的数量

    • --few_shot_random: 是否随机采样few-shot数据,如果不设置,则默认为true

  • --limit: 每个数据集最大评估数据量,不填写则默认为全部评估,可用于快速验证

评测展示:

Benchmarking summary: 
+----------------------------------------------+------------------------------------------------+
| key                                          | Value                                          |
+==============================================+================================================+
| Time taken for tests (senconds)              | 7.539                                          |
+----------------------------------------------+------------------------------------------------+
| Number of concurrency                        | 1                                              |
+----------------------------------------------+------------------------------------------------+
| Total requests                               | 15                                             |
+----------------------------------------------+------------------------------------------------+
| Succeed requests                             | 15                                             |
+----------------------------------------------+------------------------------------------------+
| Failed requests                              | 0                                              |
+----------------------------------------------+------------------------------------------------+
| Average QPS                                  | 1.99                                           |
+----------------------------------------------+------------------------------------------------+
| Average latency                              | 0.492                                          |
+----------------------------------------------+------------------------------------------------+
| Average time to first token                  | 0.026                                          |
+----------------------------------------------+------------------------------------------------+
| Throughput(average output tokens per second) | 334.006                                        |
+----------------------------------------------+------------------------------------------------+
| Average time per output token                | 0.00299                                        |
+----------------------------------------------+------------------------------------------------+
| Average package per request                  | 167.867                                        |
+----------------------------------------------+------------------------------------------------+
| Average package latency                      | 0.003                                          |
+----------------------------------------------+------------------------------------------------+
| Average input tokens per request             | 40.133                                         |
+----------------------------------------------+------------------------------------------------+
| Average output tokens per request            | 167.867                                        |
+----------------------------------------------+------------------------------------------------+
| Expected number of requests                  | 15                                             |
+----------------------------------------------+------------------------------------------------+
| Result DB path                               | ./outputs/qwen2.5_benchmark_20241107_201413.db |
+----------------------------------------------+------------------------------------------------+

Percentile results: 
+------------+---------------------+---------+
| Percentile | First Chunk Latency | Latency |
+------------+---------------------+---------+
|    10%     |       0.0178        | 0.1577  |
|    25%     |       0.0183        | 0.2358  |
|    50%     |       0.0199        | 0.4311  |
|    66%     |       0.0218        | 0.6317  |
|    75%     |       0.0429        | 0.7121  |
|    80%     |       0.0432        | 0.7957  |
|    90%     |       0.0432        | 0.9153  |
|    95%     |       0.0433        | 0.9897  |
|    98%     |       0.0433        | 0.9897  |
|    99%     |       0.0433        | 0.9897  |
+------------+---------------------+---------+
指标 说明
Total requests 总请求数
Succeed requests 成功请求数
Failed requests 失败的请求数量
Average QPS 平均每秒请求数
Average latency 平均延迟
Throughput(average output tokens per second) 平均每秒输出token的数量
Average time to first token 平均首次输出token的时间
Average input tokens per request 每个请求的平均输入token数
Average output tokens per request 每个请求的平均输出token数
Average time per output token 平均每个输出token的时间
Average package per request 每个请求的平均包数
Average package latency 平均包延迟
Percentile of time to first token (p10, ..., p99) 首token延时百分位
Percentile of time to first token (p10, ..., p99) 首token延时百分位

5.2 部署实践与接口调用

1. 大模型服务API接口调用示例

在用vllm框架部署完大模型后,可以通过openai的接口调用。

同步调用

from openai import OpenAI

openai_api_base = "http://127.0.0.1:8000/v1"
openai_api_key = "EMPTY"
client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)

def api_predict(conversation):
    response = client.chat.completions.create(
        model="Qwen",
        messages=conversation
    )
    for chunk in response:
        content = chunk.choices[0].delta.content
        if content:
            yield content

异步调用

from openai import AsyncOpenAI

openai_api_base = "http://127.0.0.1:8000/v1"
openai_api_key = "EMPTY"
client = AsyncOpenAI(api_key=openai_api_key, base_url=openai_api_base)

async def api_predict(conversation):
    response = await client.chat.completions.create(
        model="Qwen",
        messages=conversation,
        stream=True
    )
    async for chunk in response:
        content = chunk.choices[0].delta.content
        if content:
            yield content

使用AsyncOpenAI进行异步调用可以提高性能和效率,特别是在处理大量I/O操作时。异步编程允许程序在等待API响应的同时,执行其他任务,从而避免了线程阻塞。