博客

  • LLM 前沿技术跟踪:PagedAttention 升级版 vAttention

    近年来,大型语言模型(LLM)在各个领域都取得了显著进展,但其高昂的计算资源消耗和内存占用成为了制约其应用的关键瓶颈。PagedAttention 作为一种动态内存分配机制,在一定程度上解决了 LLM 的内存问题,但也存在一些缺陷,例如内存不连续导致的效率低下和开发难度增加。

    本文将介绍一种全新的内存管理技术——vAttention,它可以看作是 PagedAttention 的升级版本,有效地解决了这些问题,并显著提升了 LLM 的性能。

    PagedAttention 的局限性

    PagedAttention 通过动态分配内存来应对 LLM 自回归生成时内存需求不确定的问题。然而,这种方法存在以下几个不足:

    • 内存不连续: PagedAttention 会导致内存分配不连续,需要在框架中引入额外的内存管理代码,增加了开发难度。
    • 注意力计算效率降低: kv cache 在注意力计算中无法连续读取,需要手动修改注意力内核代码进行分段读取,降低了计算效率。
    • 性能劣于 FlashAttention: 在小数据量时,PagedAttention 的速度明显慢于 FlashAttention,可能是由于 cache 命中率高,连续内存读取效率更高导致的。

    vAttention:突破传统内存管理的限制

    vAttention 巧妙地利用了虚拟内存机制,将预分配虚拟内存和实际分配物理内存分离,并拓展了分配小内存的 PyTorch 算子,将内存管理从注意力内核中剥离出来。

    vAttention 的核心思想:

    1. 预留虚拟内存: 创建足够长的虚拟内存,但并不分配物理内存,利用系统本身的虚拟内存管理机制。
    2. 按需物理内存分配: 每次优先分配一个物理内存页面,仅当请求已使用完其先前分配的所有物理内存页面时,再分配新的物理内存,并映射到虚拟内存中。
    3. 利用低级 CUDA 支持: vAttention 利用 CUDA 低级 API,可以对虚拟和物理内存进行细粒度控制,实现更灵活的内存管理。
    4. 拓展 PyTorch 接口: vAttention 扩展了 PyTorch 缓存分配器,允许应用程序为张量保留虚拟内存缓冲区,而无需提前提交物理内存。
    5. 请求级 KV-cache 索引: vAttention 使用唯一的整数标识符来定位请求的子张量,并进行相应的内存分配和释放。

    vAttention 的优势:

    • 连续内存: vAttention 保证了 KV 缓存存储在连续的虚拟内存中,无需提前提交物理内存,避免了内存碎片化问题。
    • 性能提升: vAttention 能够无缝接入 FlashAttention 等高性能注意力计算内核,并通过优化内存管理机制,显著提升了 LLM 的性能。

    性能评估

    实验结果表明,vAttention 在吞吐量和推理时间方面都取得了显著提升,性能比 vLLM 高出 1.97 倍。

    总结

    vAttention 作为 PagedAttention 的升级版本,有效地解决了传统内存管理机制的缺陷,为 LLM 的应用提供了更强大的支持。它不仅提高了 LLM 的性能,也简化了开发流程,为 LLM 的进一步发展提供了新的方向。

    注: 本文内容主要参考了 akaihaoshuai 的知乎文章,并加入了个人理解和分析。

  • 一篇文章看懂Redission原理

    Redission 是一个基于 Redis 的 Java 客户端,它提供了一系列的分布式数据结构和服务,方便开发者在分布式环境下进行数据操作和通信。本文将深入探讨 Redission 的原理,并以可重入锁、锁重试和 WatchDog 机制、MutiLock 原理为例进行详细讲解。

    ☃️可重入锁原理

    Redission 的可重入锁利用 Redis 的 Hash 结构实现,它使用一个大 Key 来表示锁是否存在,并使用多个小 Key 来记录当前持有锁的线程信息。

    加锁逻辑:

    1. 判断锁是否存在: 如果锁不存在,则使用 redis.call('hset', KEYS[1], ARGV[2], 1) 命令将锁信息写入 Redis 的 Hash 结构中,并设置过期时间。
    2. 判断锁是否被当前线程持有: 如果锁存在,则使用 redis.call('hexists', KEYS[1], ARGV[2]) 命令判断当前线程是否持有该锁。如果是,则使用 redis.call('hincrby', KEYS[1], ARGV[2], 1) 命令将锁的 value 值加 1,表示该线程再次获得了锁。
    3. 设置过期时间: 使用 redis.call('pexpire', KEYS[1], ARGV[1]) 命令为锁设置过期时间。

    释放锁逻辑:

    释放锁时,使用 redis.call('hincrby', KEYS[1], ARGV[2], -1) 命令将锁的 value 值减 1。当 value 值减至 0 时,表示该线程不再持有锁,锁被释放。

    可重入机制:

    Redission 的可重入锁通过记录每个线程持有的锁次数来实现可重入机制。当一个线程第一次获得锁时,锁的 value 值为 1。如果该线程再次尝试获得锁,则 value 值会加 1,表示该线程再次获得了锁。只有当 value 值减至 0 时,该线程才真正释放锁。

    ☃️锁重试和WatchDog机制

    Redission 的锁重试机制是指当线程尝试获得锁失败时,会不断重试直到获得锁。WatchDog 机制则是为了防止锁在持有线程意外宕机时无法释放,而引入的一种自动续约机制。

    锁重试:

    Redission 的锁重试机制通过 while(true) 循环实现,每次循环都会尝试获得锁。如果获得锁成功,则退出循环;否则,会根据 waitTimeleaseTime 参数来控制重试频率和重试时间。

    WatchDog 机制:

    WatchDog 机制通过一个定时任务来实现,该定时任务会定期检查锁的剩余时间,并在剩余时间不足时进行续约。WatchDog 机制的核心代码如下:

    ttlRemainingFuture.onComplete((ttlRemaining, e) -> {
        if (e != null) {
            return;
        }
    
        // lock acquired
        if (ttlRemaining == null) {
            scheduleExpirationRenewal(threadId);
        }
    });

    这段代码会在锁获得成功后,启动一个定时任务,该定时任务会根据 internalLockLeaseTime 参数来设置续约时间。当定时任务触发时,会调用 renewExpirationAsync 方法来进行锁续约。

    ☃️MutiLock原理

    为了提高 Redis 的可用性,我们通常会搭建集群或者主从模式。在主从模式下,如果主机在将锁信息同步到从机之前宕机,则新的主机会丢失锁信息,导致锁失效。

    为了解决这个问题,Redission 提出了 MutiLock 锁,它将锁信息写入到所有 Redis 节点中,只有所有节点都写入成功,才算加锁成功。

    MutiLock 加锁原理:

    1. 将多个锁添加到一个集合中: Redission 会将需要加锁的所有锁添加到一个集合中。
    2. 循环尝试获取锁: Redission 会使用 while 循环,不断尝试获取集合中的所有锁。
    3. 设置总加锁时间: Redission 会设置一个总加锁时间,该时间等于需要加锁的个数乘以 1500 毫秒。
    4. 判断加锁是否成功: 如果在总加锁时间内,所有锁都获取成功,则加锁成功;否则,会再次进行重试。

    MutiLock 的优势:

    • 提高锁的可靠性: MutiLock 锁将锁信息写入所有 Redis 节点,即使某个节点宕机,也不会导致锁失效。
    • 提高锁的可用性: MutiLock 锁可以提高锁的可用性,即使某个节点宕机,其他节点仍然可以正常提供服务。

    参考文献

人生梦想 - 关注前沿的计算机技术 acejoy.com 🐾 步子哥の博客 🐾 背多分论坛 🐾 借一步网
Page Stats: PV: 1 | UV: 1
Last updated: 2025-07-07 04:25:04
沪ICP备2024052574号-1