持久化状态机：INT4内存单元如何重构LLM注意力机制

注意力机制的内存瓶颈：为什么KV Cache成为LLM的核心挑战

大语言模型（LLM）的核心是注意力机制（Attention），它让模型能够在处理每个token时"回顾"上下文中的所有历史信息。Transformer中的注意力机制由Vaswani等人在2017年的论文《Attention Is All You Need》中提出，其核心公式为Attention(Q,K,V) = softmax(QK^T/√d_k)V。具体而言，多头注意力（Multi-Head Attention）将输入分别通过不同的线性投影映射为多组Q、K、V，每组称为一个"注意力头"，各头独立计算注意力后拼接并再次线性变换。这种设计让模型能同时关注不同子空间的信息模式——例如某些头可能学会追踪语法依赖，另一些头则捕捉语义相似性。在自回归推理过程中，模型每生成一个新token，都需要计算该token的query与所有历史token的key之间的相似度。

在自回归生成模式中，模型逐token地生成输出序列。每生成一个新token时，该token的embedding经过线性变换产生query向量，而所有已生成token（包括prompt中的token）的key和value向量需要参与注意力计算。如果不使用KV Cache，每生成一个新token都需要对所有历史token重新计算key和value，导致生成N个token的总计算量为O(N²)。KV Cache通过「空间换时间」的策略，将已计算的key和value存储下来，使每一步只需计算新token的key/value并追加到缓存中，将每步的增量计算降至O(N)。这种策略在短序列时非常高效，但当序列长度达到数万甚至数十万token时，缓存本身的内存开销成为了新的瓶颈。

然而，随着上下文长度的不断增长，传统注意力机制面临着严峻的内存与计算瓶颈——KV Cache（键值缓存）的体积会随序列长度线性增长，成为长上下文推理场景中的主要成本来源。为避免重复计算，工程上采用KV Cache策略：将每一层、每一步已经计算过的Key和Value张量缓存起来。对于一个拥有L层、H个注意力头、头维度为d的模型，处理长度为N的序列时，KV Cache的总内存占用为2×L×H×d×N×sizeof(dtype)。以Llama-2-70B为例，其80层、64头、128维的配置下，处理4096长度的序列，单个请求的KV Cache在FP16下就需要约10GB显存，这在高并发场景下会迅速耗尽GPU内存。

现代GPU的内存层次结构对KV Cache的管理有深远影响。以NVIDIA H100为例，其HBM3显存容量为80GB，带宽为3.35TB/s，而片上SRAM仅约50MB但带宽超过30TB/s。LLM推理的生成阶段（decode phase）是典型的内存带宽受限（memory-bound）任务：每生成一个token，需要从HBM中读取整个KV Cache进行注意力计算，计算访存比极低。判断一个操作是计算受限还是内存受限，关键指标是算术强度（Arithmetic Intensity），即每字节数据传输对应的浮点运算次数（FLOPs/Byte）。当算术强度低于硬件的"屋脊线"（roofline）拐点时，操作就是内存受限的。对于decode阶段的注意力计算，batch size为1时每个KV Cache元素仅参与一次乘加运算，算术强度约为1 FLOP/2 Bytes（FP16），远低于H100约200 FLOPs/Byte的拐点，因此几乎完全受制于内存带宽。这意味着KV Cache的大小不仅影响能服务多少并发请求（受总显存容量限制），还直接影响单次生成的延迟（受内存带宽限制）。vLLM提出的PagedAttention、以及FlashAttention对注意力计算的IO优化，都是在这一背景下诞生的工程解决方案。

近期出现的一个技术议题《Persistent State Machines: LLM Attention with INT4 In-Memory Cells》（持久化状态机：基于INT4内存单元的LLM注意力）提出了一个颇具启发性的思路：将注意力的历史状态视为一种可持久化的状态机，并通过INT4量化的内存单元来存储这些状态。其背后触及的技术方向值得深入探讨。

什么是"持久化状态机"式的注意力

从瞬时计算到持久状态

传统的Transformer注意力在每次推理时都需要将当前query与所有历史key进行点积运算，再对value加权求和。这本质上是一种"无状态"的重复计算模式——虽然KV Cache缓存了历史的键值，但整体架构仍将上下文视为一个需要每步重新扫描的序列。

"持久化状态机"的核心理念，是把注意力的历史信息压缩成一组随时间演化的持久状态（persistent state）。这与近年兴起的状态空间模型（State Space Models，如Mamba）、线性注意力（Linear Attention）以及RWKV等架构的思路一脉相承：用一个固定大小的隐状态来"总结"历史，而非保留完整的KV序列。

状态空间模型源于控制理论中的线性时不变系统，其连续形式为h'(t) = Ah(t) + Bx(t), y(t) = Ch(t) + Dx(t)，其中h是隐状态，x是输入，y是输出，A/B/C/D是系统参数矩阵。通过零阶保持（ZOH）等离散化方法，可将其转化为离散递推形式h_k = Āh_{k-1} + B̄x_k，这使得模型既能以递推方式高效推理（O(1)内存、O(N)时间），又能通过展开为卷积形式实现并行训练。Mamba在此基础上引入了选择性机制（Selective SSM），让A、B、C矩阵依赖于输入内容动态变化，从而突破了传统SSM在内容感知选择上的局限性。

传统SSM（如S4、H3）使用与输入无关的固定参数矩阵A、B、C，这使得模型无法根据当前输入内容选择性地记住或遗忘信息——这正是Transformer注意力机制的核心优势。Mamba的选择性机制（Selective SSM）通过让B、C矩阵和离散化步长Δ成为输入的函数（通过线性投影计算），实现了内容感知的状态更新。然而，输入依赖的参数破坏了线性时不变性，使得模型无法再通过FFT加速的卷积进行并行训练。Mamba通过硬件感知的并行扫描（parallel scan）算法解决了这一问题，在GPU上实现了高效的训练和推理。Mamba-2进一步将选择性SSM与结构化注意力矩阵（半可分矩阵）建立了理论联系，证明了SSM和注意力之间存在深层的数学等价性。

状态机建模的核心优势

将注意力建模为状态机的最大优势在于内存占用的恒定性。传统KV Cache随序列长度线性增长，而状态机的隐状态大小固定，这意味着理论上可以处理无限长的上下文而不增加内存成本。这对于需要长程记忆的应用（如长文档理解、多轮对话、Agent持续运行）挺重要的。

从线性注意力的视角来理解这种恒定性：标准softmax注意力的复杂度为O(N²d)，线性注意力的核心思想是用核函数φ替代softmax，即Attention(Q,K,V) ≈ φ(Q)(φ(K)^TV) / (φ(Q)φ(K)^T1)。通过改变计算顺序（先计算K^TV的外积和），可将复杂度降至O(Nd²)。更关键的是，φ(K)^TV可以写成递推形式S_t = S_{t-1} + φ(k_t)v_t^T，其中S_t是一个d×d的状态矩阵。这个状态矩阵就是「持久状态」的数学基础——它以固定大小d²存储了所有历史信息的压缩表示。RetNet、GLA（Gated Linear Attention）、Lightning Attention等工作都在这一框架下改进状态更新规则以提升表达能力。

INT4内存单元：极致量化如何压缩状态存储

4位整数存储带来的内存收益

该议题最引人注目的技术点在于使用INT4（4位整数）内存单元来存储状态。量化技术在LLM推理优化中已广泛应用——从FP16到INT8再到INT4，每一次精度降低都能带来数倍的内存节省和吞吐提升。

模型量化是将浮点数映射为低位整数的过程，核心公式为q = round(x/s) + z，其中s为缩放因子，z为零点偏移。主流量化方案包括：对称量化（零点为0）、非对称量化、分组量化（group-wise，每g个元素共享一组量化参数）。INT4量化意味着每个数值仅能取16个离散值（有符号为-8到7），动态范围极其有限。业界已有GPTQ、AWQ、QuIP#等方法针对静态权重实现了INT4量化且保持了较好的精度。其中，GPTQ基于最优脑量化（OBQ）框架，通过逐层量化并利用Hessian矩阵信息补偿量化误差；AWQ（Activation-aware Weight Quantization）则观察到权重中存在少量对激活值影响极大的"显著通道"，通过对这些通道施加缩放来保护关键信息；QuIP#则利用随机正交矩阵对权重进行预处理，使权重分布更加均匀从而更适合均匀量化。

将状态机的隐状态量化到INT4，意味着每个状态值仅占用4位存储空间。相比FP16的16位，这带来了4倍的内存压缩。对于需要长期驻留在内存中的持久状态而言，这种压缩尤为关键——它直接决定了单张GPU或边缘设备能够维持多少并发会话，或支持多长的有效上下文。

从硬件架构角度看，INT4内存单元的概念可能与存内计算（Processing-in-Memory, PIM）或近存计算（Near-Memory Computing）趋势相关。传统冯诺依曼架构中，数据需要从内存搬运到处理单元，这一过程的能耗往往远超计算本身（所谓的「内存墙」问题）。据估算，在先进工艺节点下，从片外DRAM读取一个64位数据的能耗约为一次64位浮点乘法运算的200倍以上。新兴的存内计算范式将计算逻辑嵌入到内存阵列中，利用模拟电路或数字逻辑直接在存储位置完成矩阵-向量乘法等运算。具体实现方式包括：基于SRAM的数字PIM（如Samsung的HBM-PIM、UPMEM的DRAM-PIM），利用存储单元阵列的位线电流累加实现乘累加操作；基于RRAM/ReRAM（阻变存储器）的模拟PIM，利用欧姆定律（V=IR）和基尔霍夫电流定律自然实现乘法和累加；以及基于闪存（Flash）的PIM方案。三星已在2021年推出了集成PIM功能的HBM2E产品（HBM-PIM），能在内存堆栈内部直接执行简单的向量运算。INT4格式与这类硬件的匹配度极高，因为低位宽运算降低了模拟计算的精度要求和数字计算的面积开销。如果持久状态以INT4格式常驻于PIM单元中，状态读取和更新可以在内存侧直接完成，极大地减少数据搬运开销，这可能是该议题标题中"In-Memory Cells"的深层含义所在。

精度与效率的权衡挑战

当然，INT4量化并非没有代价。极低精度可能引入量化误差，尤其是在状态需要长期累积、演化的场景下，误差可能随时间放大。与静态权重量化不同，状态量化面临独特挑战：状态值在每个时间步都在更新，量化-反量化操作形成的舍入误差会在递推过程中逐步累积，类似于数值积分中的截断误差传播问题。如何设计合适的量化方案（如分组量化、动态缩放因子）来在存储效率与状态保真度之间取得平衡，是这类方案能否实用化的关键技术难题。

状态量化的误差传播问题可以通过离散动力系统的稳定性分析来理解。考虑状态更新h_t = A·quantize(h_{t-1}) + B·x_t，每步量化引入的舍入误差ε_t的传播由矩阵A的谱性质决定。如果A的谱半径ρ(A) < 1，则历史误差会指数衰减，系统具有自稳定性；如果ρ(A) ≈ 1（如门控机制中常见），误差可能长期驻留甚至缓慢增长。实际工程中，可能的缓解策略包括：周期性状态重新量化校准（periodic recalibration）、引入随机舍入（stochastic rounding）使量化误差期望为零、使用混合精度方案对关键状态通道保持更高精度、以及设计量化感知的训练流程（quantization-aware training）让模型学会在低精度状态下维持信息保真度。

技术定位：与Mamba、RWKV等主流架构的关系

线性注意力与极致量化的结合

这一思路可以看作是线性注意力/SSM架构与极致量化技术的结合体。近年来，学术界和工业界一直在探索如何突破Transformer的二次复杂度瓶颈：

• Mamba/SSM：通过选择性状态空间机制实现线性复杂度的序列建模
• RWKV：将RNN的循环特性与Transformer的并行训练能力结合。RWKV（Receptance Weighted Key Value）由彭博创建，其核心是WKV（Weighted Key Value）注意力机制，通过时间衰减因子w实现指数加权的历史信息聚合。这个机制可以通过维护两个累积状态（分子和分母）实现O(1)的递推推理。RWKV已发展到第六代（Eagle/Finch），在多项基准测试中接近同规模Transformer的性能，同时推理内存仅为常数级。
• Linear Attention：将softmax注意力近似为可递归计算的线性形式

"持久化状态机+INT4"的组合，本质上是在这些架构基础上进一步压榨内存效率，特别适合部署在资源受限的环境中。

潜在应用场景分析

如果这一方向能够成熟落地，最直接的受益场景包括：

1. 边缘设备上的长上下文推理：手机、嵌入式设备等内存有限的平台。当前主流边缘AI芯片的内存配置十分有限：Apple A17 Pro的Neural Engine可用内存约为系统共享的数GB，高通骁龙8 Gen 3的Hexagon NPU可访问约12-16GB统一内存。即便是最轻量的7B参数模型，INT4权重本身就需要约3.5GB存储，留给KV Cache和运行时的空间非常有限。以Llama-3-8B处理32K上下文为例，FP16 KV Cache需要约16GB，即使INT8量化后仍需8GB，远超边缘设备的承受能力。这正是固定大小持久状态方案的核心价值所在——无论上下文多长，状态占用始终恒定，使得边缘设备上的长上下文推理成为可能。
2. 高并发推理服务：在同等显存下支持更多用户会话。以一个典型的云端推理服务为例，若使用80GB显存的A100 GPU，模型权重（如Llama-2-70B INT4量化后约35GB）占据约一半显存，剩余约45GB用于KV Cache。传统FP16 KV Cache下，每个4K上下文的会话占用约10GB，意味着仅能同时服务4-5个会话。若采用固定大小的INT4持久状态（假设状态大小等效于256 token的KV Cache），每个会话的状态占用可降至约160MB，理论上同一张卡可同时维持数百个活跃会话，吞吐量提升可达两个数量级。
3. 持续运行的AI Agent：需要维持长期状态记忆的自主智能体。当前的AI Agent框架（如AutoGPT、MetaGPT等）在处理长期任务时，通常依赖外部记忆系统（向量数据库检索、文本摘要压缩等）来规避上下文长度限制。这些方案引入了额外的检索延迟和信息损失。如果模型本身能通过持久化状态机维持一个不断更新的、紧凑的「工作记忆」，Agent就能在不依赖外部存储的情况下保持对长期任务上下文的连贯理解，实现更自然的持续推理。

理性看待：早期探索的价值与现实局限

需要客观指出的是，这一议题目前尚处于早期探索阶段，更像是一个技术设想而非经过大规模验证的成熟方案。

从工程角度看，将INT4量化应用于持续演化的状态（而非静态权重）本身就是一项挑战——权重量化是"一次性"的，而状态量化需要在每个时间步反复进行量化-反量化操作，这既可能引入额外计算开销，也可能导致误差累积。具体而言，假设每一步引入的量化误差为ε，经过T步递推后，累积误差在最坏情况下可达O(T·ε)（如果状态更新矩阵的谱半径接近1）或O(ε/(1-ρ))（如果谱半径ρ<1提供了自然的误差衰减）。这意味着系统设计时需要仔细控制状态更新的动态特性，确保误差不会无界增长。这些实际问题都需要严谨的实验来验证。

此外，从模型能力的角度看，固定大小状态相比完整KV Cache存在信息容量的理论上界。信息论的分析表明，一个d×d的INT4状态矩阵最多能存储4d²比特的信息，而完整的FP16 KV Cache存储N×d×16比特。当上下文信息量超过状态容量时，必然发生信息丢失。如何设计智能的信息压缩策略——让模型学会「记住什么、遗忘什么」——是这类架构需要解决的根本性问题。Mamba和RWKV的成功表明，经过合理设计的门控机制确实能让有限状态承载令人惊讶的信息量，但在需要精确回忆远距离细节（如「needle in a haystack」测试）的任务上，仍与完整注意力存在差距。近期的一些混合架构（如Jamba由AI21 Labs提出，交替使用Mamba层和Transformer注意力层）试图在两者之间取得平衡：用SSM层处理大部分序列建模任务以获得效率优势，同时保留少量注意力层用于需要精确远程检索的场景。这种混合策略可能是未来实用系统的务实选择。

总结：固定内存与极致压缩的设计哲学

"持久化状态机：INT4内存单元LLM注意力"这一议题，虽然尚处萌芽阶段，却精准地指向了当前LLM发展中的两大核心痛点：长上下文的内存成本与推理部署的资源约束。它将状态空间建模的思想与极致量化技术相结合，为突破Transformer的内存瓶颈提供了一个值得关注的技术方向。

无论这一具体方案最终能否落地，其背后所代表的"固定内存、持久状态、极致压缩"的设计哲学，都值得整个AI基础设施领域持续关注。在大模型规模与上下文长度不断膨胀的今天，如何用更少的内存做更多的事，将始终是一个核心命题。从更宏观的技术演进视角看，这一方向可能代表着从「存储所有历史」到「压缩性记忆」的范式转变——正如人类大脑并不逐帧记录所有感知输入，而是将经历压缩为结构化的记忆表征。未来的高效序列模型，很可能在完整注意力（精确但昂贵）与固定状态（高效但有损）之间找到最佳的混合架构，在不同层或不同注意力头上采用不同的记忆策略，实现精度与效率的帕累托最优。