TinyServe: Query-Aware Cache Selection for Efficient LLM Serving

发表时间: 2025-10 · arXiv:2509.12211 (MM 2025)

原文: https://arxiv.org/abs/2509.12211

Dong Liu (Yale University)
Yanxuan Yu (Columbia University)

速读

一句话结论
TinyServe 是一个专为小型大语言模型设计的轻量级推理服务框架，它通过引入基于查询感知的键值缓存动态选择机制，在几乎不损失精度的前提下，实现了最高 3.4 倍的解码提速和 2 倍以上的显存节省。

要解决什么问题
大语言模型在自回归解码阶段面临严重的显存带宽和延迟瓶颈。在生成每一个新 token 时，模型都需要读取历史上下文中所有的键值向量（KV cache）来计算注意力分数。随着上下文长度的增加（如达到 16K 或 32K），这种每次解码都要全量扫描高带宽显存（HBM）的操作会导致极大的内存搬运开销。现有的静态剪枝或统一缓存淘汰策略（如只保留最近的 token 或固定丢弃某些 token）存在一个致命盲区：不同的查询向量（Query）对历史 token 的依赖是动态变化的。一个在大多数解码步中看似无关紧要的 token，可能会在遇到特定查询时瞬间变得至关重要。如果采用一刀切的缓存清理，会导致关键信息丢失、模型精度下降；如果全量保留，则显存读取效率极低。此外，在 7B 以上规模的模型上直接验证和迭代底层系统设计的成本过高，使得系统研究人员难以低成本地观测和优化注意力机制的底层行为。

怎么做的
TinyServe 的核心思路是查询感知的页面级稀疏注意力，即根据当前生成的查询向量，动态且精准地只从显存中加载最相关的 KV 数据块。为了绕开全量扫描 HBM 的带宽卡点，作者将 KV cache 划分为固定大小的页面，并为每个页面提取极小的数据摘要（元数据）存放在高速缓存（SRAM 或 L2）中，通过在高速缓存中进行低成本的粗筛，决定去 HBM 中拉取哪些真正的页面。具体而言，该机制由三个关键设计构成：第一，页面级边界框元数据。将历史键向量按大小为 $S$ 划分为 $P$ 个页面 $\mathcal{K}_j$。每个页面只维护一个通道级别的最大值和最小值作为元数据 $\phi(\mathcal{K}_j) = (m_j, M_j)$。第二，基于查询的动态相关性打分。在解码步 $t$，利用当前的查询向量 $q_t$ 与各个页面的元数据计算相关性得分。该得分本质上是评估该查询与该页面内所有键向量点积的理论上限：

$$\begin{aligned} r(q_t, \phi(\mathcal{K}_j)) = \sum_{i=1}^d \begin{cases} q_{t,i} \cdot M_{j,i}, & \text{if } q_{t,i} \ge 0 \\ q_{t,i} \cdot m_{j,i}, & \text{if } q_{t,i} < 0 \end{cases} \end{aligned}$$


系统根据该得分选出排名前 $K$ 的页面集合 $S_t$。第三，融合的 CUDA 算子。作者开发了一个单趟执行的底层算子，将元数据打分、Top-K 页面筛选、稀疏显存加载以及最终的掩码注意力计算全部融合在一起：

$$\mathrm{SparseAttn}(q_t) = \sum_{j \in S_t} \sum_{k_i \in \mathcal{K}_j} \mathrm{softmax}(q_t^\top k_i) \cdot v_i$$
这种设计使得单步解码的延迟从全量加载的开销，转变为轻量级元数据扫描开销加上部分页面的加载与计算开销，从而在物理层面大幅降低了 HBM 的带宽压力。

效果如何
实验在 8 张 A100 80GB GPU 上进行，评测了从 125M 到 1.3B 参数规模的多个模型（如 TinyLLaMA-125M、GPT2-345M、LLaMA-1.3B），上下文长度最高达 32K。对比基线分为两类：一类是工业级推理系统，包括 vLLM（代表 PagedAttention 路线）、TGI 和 TensorRT-LLM；另一类是前沿的缓存压缩研究方法，包括 StreamingLLM（代表滑动窗口路线）、SnapKV 和 PyramidKV（代表静态与聚类剪枝路线），以及全量缓存。作者将 TinyServe 的核心机制直接集成到了 vLLM 中进行对比。在 LongBench 长文本任务和 2048 个 token 的严格缓存预算下，TinyServe 在保持与全量缓存几乎一致精度的同时，平均解码速度提升了 2.1 到 3.4 倍。在模拟 1024 个并发请求的生产级多用户负载测试中（使用 GPT2-345M），TinyServe 的吞吐量达到了 28.6 请求/秒，显著高于原生 vLLM 的 18.4 请求/秒，且 P50 延迟从 45.2 毫秒降低至 32.1 毫秒。该方法的代价与局限在于页面大小的设定是一个硬性权衡：页面设置过大（如 64）会降低元数据扫描的开销，但会导致边界框估算变得粗糙，拉取大量无效 token 从而导致模型困惑度上升；页面设置过小（如 4）虽然精度极高，但元数据数量激增会拖慢扫描速度，实验表明页面大小为 16 是最佳平衡点。此外，边界框估算天然带有近似误差，该误差的上限受限于页面内键向量的方差大小。

主要贡献

大型语言模型（LLMs）在对话、检索、摘要和代码生成等现代AI应用中发挥着核心作用。然而，在长上下文或高吞吐量条件下，由于自回归解码期间键值（KV）缓存访问带来的高内存和延迟开销，高效地部署LLMs仍然具有极大的挑战性。尽管近期的系统（如vLLM、TGI和FasterTransformer）引入了分页注意力、投机解码和缓存重排序等复杂策略，但理解LLM训练和推理的内部动态仍然很困难。系统研究人员通常只能将模型视为黑盒，且在没有大型GPU集群的情况下难以验证假设或进行设计迭代。

为了解决这一问题，本文提出了TinyServe，这是一个轻量级且可扩展的服务框架，专为部署微型LLMs（例如TinyLLaMA，GPT2-345M）而设计。本文的核心创新点包括：
1. TinyServe框架：提出了一个能够使用微型且高效的架构进行快速、可解释的训练和推理服务框架。与以往的模拟框架不同，它支持结构化的KV稀疏性、基于插件的Token选择机制以及硬件高效的注意力内核，能够在完全可控的环境中执行实时解码。
2. 查询感知（Query-Aware）的KV选择机制：引入了一种利用边界框元数据来估计查询与KV缓存块之间注意力相关性的机制。该机制能够在不修改模型的情况下，以极低的开销动态选择最相关的KV部分，从而减少内存移动并保持准确性。
3. 高保真度的系统级重现：通过在标准数据集和诊断数据集上进行广泛的实验，证明了TinyServe能够忠实地复制在大型模型中观察到的关键延迟-准确性权衡，为资源受限硬件上的LLM服务行为研究提供了可访问、可重复的基础。

背景知识与设计原则

微型模型在LLM服务研究中的作用：虽然大多数LLM研究集中在数十亿参数的大规模模型上，但近期研究（如TinyStories和TinyLLaMA）表明，125M至350M参数的微型语言模型在经过适当训练后，可以捕获大型模型中的许多语言特性。本文延续了这一思路，将微型LLM重新用于训练和推理级别的分析，表明即使在小规模下，也能忠实地重现Token级别的延迟、缓存重用行为和准确率下降等现象。

LLM推理分析与加速现状：现有的系统级框架（如vLLM的PagedAttention、FasterTransformer的自定义CUDA内核）在优化大模型推理效率方面做出了贡献，但对这些系统进行性能分析计算成本高昂，且部署的复杂性往往掩盖了细粒度的系统洞察。TinyServe通过使用小型LLM在真实服务场景下重现流式注意力、动态批处理和量化解码等关键服务栈组件，能够以极低的计算成本快速测试架构变更的假设。

服务导向的基准测试：受算法推理和可解释性基准测试的启发，本文采用了针对服务领域的特定任务（如复制、计数、罕见Token召回）来对系统进行压力测试。这些基准测试能够精确评估系统干预（如剪枝或近似）的效果，并帮助隔离服务效率低下的根本原因。

推理时间由解码阶段主导：LLM推理包含预填充（prefill）和解码（decode）两个阶段。在预填充阶段，所有提示Token被转换为键（$K$）、查询（$Q$）和值（$V$）向量并存入KV缓存。在解码阶段，每生成一个新Token，都会产生一个新的$Q$并与所有存储的$K$向量进行比较以计算注意力分数。由于解码是逐Token进行的，且每次都需要读取整个KV缓存，因此它占据了绝大部分的延迟——特别是当序列长度达到16K或32K时。

查询感知稀疏性的优化动机：尽管先前的研究表明只有一小部分KV Token对准确预测至关重要【索引1，LoRA: Low-Rank Adaptation of Large Language Models+2022+ICLR+URL】【索引2，H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models+2023+URL】，但本文观察到，这一关键Token集合在不同的查询之间存在显著差异。如图1所示，某些Token可能在大多数解码步骤中影响极小，但在与特定查询对齐时会瞬间变得至关重要。为了在计算和内存预算有限的微型LLM中高效支持推理，本文通过查询感知稀疏性来优化自注意力机制：即以当前查询向量为条件，动态仅选择最相关的KV Token。这种机制消除了存储和关注无关Token的开销，同时通过保留与当前解码步骤相关的上下文来维持准确性。在TinyServe中，查询感知路由是在页面（page）粒度上实现的，通过维护轻量级的元数据（存储的Key向量的通道级最大最小值），实现了在极小内存移动下高效选择Top-$K$页面。

图1：查询感知Token选择的动机。每个查询向量关注KV页面的不同子集。统一的缓存保留会导致不必要的内存读取，而查询感知路由通过聚焦高相关性区域来实现动态稀疏性。
图1：查询感知Token选择的动机。每个查询向量关注KV页面的不同子集。统一的缓存保留会导致不必要的内存读取，而查询感知路由通过聚焦高相关性区域来实现动态稀疏性。

方法细节

系统架构概览：TinyServe是一个专为在严格的内存和延迟限制下服务微型语言模型而设计的轻量级框架。它不仅仅是一个基准测试工具，而是一个支持稀疏感知注意力、模块化Token选择和高效KV缓存重用的实时服务环境。

三大核心组件设计：系统围绕三个核心组件组织：（1）查询感知KV检索器（Query-Aware KV Retriever）：在解码时根据当前查询向量和页面级元数据动态选择相关的键值块，减少不必要的内存访问。（2）模块化调度管道（Modular Scheduling Pipeline）：调度循环处理传入的查询，并将它们路由通过可配置的插件（例如，基于熵的提前退出、Token级剪枝、近似注意力）。这种模块化设计允许在不修改核心模型的情况下尝试不同的稀疏策略。（3）稀疏注意力执行器（Sparse Attention Executor）：使用融合的CUDA内核对选定的KV页面高效计算注意力，支持FP16/INT8 KV格式和多GPU调度。

TinyServe的执行流程：在TinyServe中，每个解码步骤都会激活管道：查询向量用于对KV页面进行评分，获取排名靠前的页面，执行稀疏注意力，并且插件模块可能会触发剪枝或提前停止。这种设计平衡了灵活性和效率，并支持微型LLM中稀疏策略的静态部署和研究原型设计。

训练加速支持：除了推理优化，TinyServe还为训练加速和细粒度分析提供了专门的支持。对于训练场景，TinyServe实现了梯度感知的内存管理，在反向传播期间根据梯度大小选择性地保留KV缓存条目。这种方法在微调期间将内存占用减少了高达$40\%$，同时保持了训练的稳定性。

细粒度性能分析系统：分析系统集成了微秒级精度的层级性能监控，跨不同模型层跟踪注意力模式、内存访问模式和计算瓶颈。这使得能够对训练动态进行详细分析，并有助于识别前向和后向传递中的优化机会。分析数据通过轻量级的插桩钩子收集，为训练过程增加的开销极小（$<2\%$）。

分布式训练支持：对于分布式训练场景，TinyServe支持具有可配置通信模式的异步梯度同步，允许研究人员在不修改核心训练循环的情况下尝试不同的并行化策略。这对于探索资源受限硬件上的高效训练策略特别有价值。

标准Transformer解码层的注意力计算：在标准Transformer解码器层中，解码步骤$t$的注意力计算涉及一个新的查询$q_t \in \mathbb{R}^d$关注所有过去的键$K_{<t} = \{k_1, k_2, \dots, k_{t-1}\}$：<br />

$$ \mathrm{Attn}(q_t, K, V) = \sum_{i=1}^{t-1} \mathrm{softmax}(q_t^\top k_i) \cdot v_i $$

推理延迟瓶颈：这一过程在推理期间是延迟关键的，主要因为两个瓶颈：一是内存移动，需要从高带宽内存（HBM）加载所有的$k_i, v_i$；二是无结构访问，注意力计算需要全键扫描，没有缓存预取模式。

结构化内存布局：为了解决这个问题，TinyServe通过将Token分组为固定大小的页面引入了结构化的内存布局。令$K = \cup_{j=1}^P \mathcal{K}_j$被划分为$P = \lceil t / S \rceil$个大小为$S$的页面。每个页面$\mathcal{K}_j$存储一个小的元数据摘要$\phi(\mathcal{K}_j)$，以实现相关性估计。

问题公式化定义：本文定义了一个相关性函数$r : \mathbb{R}^d \times \mathbb{R}^{2d} \to \mathbb{R}$，使得：

$$ r(q_t, \phi(\mathcal{K}_j)) \approx \max_{k \in \mathcal{K}_j} q_t^\top k $$


然后，选择一个页面索引子集$S_t \subseteq \{1, \dots, P\}$，使得：

$$ S_t = \mathrm{TopK}_j ~ r(q_t, \phi(\mathcal{K}_j)) \quad \mathrm{with} \quad |S_t| = K $$
注意力计算随后仅在选定页面的并集上进行：
$$ \mathrm{SparseAttn}(q_t) = \sum_{j \in S_t} \sum_{k_i \in \mathcal{K}_j} \mathrm{softmax}(q_t^\top k_i) \cdot v_i $$

相关性函数实例化：本文将$r$实例化为一个定向边界框估计器，它使用每个维度的边界：

$$\begin{aligned} \begin{array}{rlr} \displaystyle \phi(\mathcal{K}_j) = (m_j, M_j) \in \mathbb{R}^{2d}, & & \\ \displaystyle r(q_t, \phi(\mathcal{K}_j)) = \sum_{i=1}^d \{ q_{t,i} \cdot M_{j,i}, \mathrm{~if~} q_{t,i} \geq 0 & & \\ r(q_t, \phi(\mathcal{K}_j)) = \sum_{i=1}^d \{ q_{t,i} \cdot m_{j,i}, \mathrm{~if~} q_{t,i} < 0 \end{array} \end{aligned}$$

硬件执行模型假设：假设每个页面$\mathcal{K}_j$驻留在HBM中，并作如下假设：从HBM获取页面的成本为$\tau_{\mathrm{hb}} \cdot S$个周期；驻留在缓存中的元数据$\phi(\mathcal{K}_j)$存储在SRAM或L2中，成本$\tau_{\mathrm{meta}}$可忽略不计；页面选择成本为$O(P \cdot d)$，但可以融合到GPU上的单个内核中。

有效延迟成本计算：假设选择了$K$个页面，有效延迟成本变为：

$$ \mathrm{Latency}_t = \underbrace{\tau_{\mathrm{meta}} \cdot P}_{\mathrm{lightweight scan}} + \underbrace{\tau_{\mathrm{hb}} \cdot K \cdot S}_{\mathrm{KV load}} + \tau_{\mathrm{attn}}(K \cdot S) $$

结构感知设计的优势：这种结构感知的设计确保了：查询依赖的缓存激活；内存感知调度（例如，将热页面保留在共享内存中）；以及降低了HBM带宽压力。

系统实现意义：TinyServe在不需要重新训练架构的情况下实现了动态的查询感知稀疏性。模块化实现直接集成到TinyServe的内核循环中，并允许硬件敏感的调度：例如，将热页面保留在共享内存中，或限制$K$以匹配张量核心的粒度。TinyServe的内核设计详见如下算法：

# Algorithm 1: Fused Query-Aware Sparse Attention Kernel
# Require: Query vector q_t in R^d, Page metadata {\phi_j = (m_j, M_j)}_{j=1}^P, KV-cache {k_i, v_i}_{i=1}^L
# Ensure: Output vector o_t in R^d

# Step 1: Relevance scoring over page metadata (in L2/shared)
for j in range(1, P + 1): # in parallel
    s_j = 0
    for i in range(1, d + 1):
        q_i = q_t[i]
        s_j += q_i * (M_{j,i} if q_i >= 0 else m_{j,i})

# Step 2: Top-K page selection (shared heap or radix select)
S_t = TopK(s_1, ..., s_P)

# Step 3: Sparse KV gather (HBM access)
K_selected, V_selected = [], []
for j in S_t: # in parallel
    # Fetch page K_j = {k_{j,1}, ..., k_{j,S}} from HBM
    # Append keys to K_selected, values to V_selected
    pass

# Step 4: Attention computation over selected KV pairs
for i in range(1, len(K_selected) + 1):
    a_i = q_t.T @ k_i
\alpha = softmax(a)
o_t = sum(\alpha_i * v_i for i in range(len(K_selected)))
return o_t

内存效率量化模型构建：为了量化查询感知稀疏性下的内存访问节省，本文构建了一个概率成本模型，该模型考虑了（1）元数据开销，（2）选定的KV Token，以及（3）跨步重用。该分析为通过本文方法可实现的性能提升提供了理论边界。

变量定义：设$L$为总缓存长度（Token数）；$S$为页面大小（每页Token数）；$K$为选定的页面数；$M$为每个Token的内存（字节）；$\rho$为相邻解码步骤之间选定页面的重用概率。

单步解码内存移动公式：每个解码步骤的内存移动为：

$$ \mathrm{Load} = 2M \cdot \left( \frac{L}{S} + \rho \cdot K \cdot S \right) $$

变量解释：其中$\frac{L}{S}$个页面存储$\min/\max$元数据（两个长度为$d$的向量），$\rho$考虑了摊销重用——即每步只有$\rho K$个页面是新加载的。

与全缓存注意力的归一化对比：为了与全缓存注意力进行比较，本文进行归一化：

$$ \mathrm{Memory ~ Fraction} = \frac{1}{S} + \rho \cdot \frac{K \cdot S}{L} $$

理论最优边界：对于最优页面大小$S^* = \sqrt{L/K}$，内存比例变为：

$$ \mathrm{Memory ~ Fraction}^* = \frac{2\sqrt{K/L}}{S^*} + \rho \cdot \frac{K \cdot S^*}{L} = 2\sqrt{\frac{K}{L}} \cdot \rho $$

理论下界与实际收益：这提供了内存移动的理论下界。对于典型值（$K = 0.3P$, $L = 32K$, $S = 16$），与全缓存注意力相比，实现了约$8\times$的内存移动减少。

查询感知准确率分析：边界框估计器在相关性评分中引入了近似误差$\epsilon$：

$$ \epsilon = \max_{k \in \mathcal{K}_j} q_t^\top k - r(q_t, \phi(\mathcal{K}_j)) $$

误差上界推导：在对查询和键分布的合理假设下，可以界定此误差：

$$ \mathbb{E}[\epsilon] \le \frac{d \cdot \sigma^2}{S} \cdot \sqrt{\log(S)} $$


其中$\sigma^2$是页面内键向量的方差。这个边界确保了本文的近似方法在实现显著内存节省的同时保持高准确率。

实验环境

  • 数据集名称、规模及用途：

    • 语言建模任务：PG19、WikiText-103、C4-News，用于基础语言能力测试。
    • 长上下文任务：LongBench（包含NarrativeQA, Qasper, GovReport, TriviaQA, HotpotQA），Passkey retrieval，用于评估长上下文理解和检索能力。
    • 推理任务：MMLU、LAMBADA、多轮QA、代码生成。
    • 服务工作负载：模拟生产环境中的多用户场景，包含512-2048个并发请求。

  • 模型架构关键参数：评估涵盖了多个不同参数规模的模型，包括TinyLLaMA-125M、GPT2-345M、OPT-350M、GPT2-774M和LLaMA-1.3B。对于大于2048 Token的序列，使用1024 Token重叠的滑动窗口处理。

  • 硬件配置：8张NVIDIA A100 80GB GPUs。
  • 软件配置：CUDA 11.8, cuDNN 8.7, PyTorch 2.0.1, 操作系统为Ubuntu 20.04 LTS。基线系统包括生产系统（vLLM、TGI、TensorRT-LLM）和研究方法（StreamingLLM窗口大小2048、SnapKV聚类大小64、PyramidKV top-k=512），以及剪枝基线（FullCache、SoftPrune阈值0.1、EntropyStop阈值0.5）。vLLM集成了本文的查询感知页面选择机制进行对比测试。

实验结果

  • 多用户工作负载下的服务栈评估：
    • 实验内容：比较修改后的vLLM（TinyServe）与原始vLLM基线。模拟了512-2048个并发请求，请求到达服从泊松分布（平均间隔50ms）。
    • 实验结果：表3显示，TinyServe在P50延迟（32.1ms）、P99延迟（89.3ms）、吞吐量（28.6 req/s）和GPU利用率（91.2%）上均优于vLLM、TGI和TensorRT-LLM。图2和图3展示了TinyServe在会话管理中实现了更高的缓存重用率，同时最小化了迁移成本。
    • 分析结论：TinyServe在真实的多用户并发服务场景中展现出优越的会话管理能力和整体性能。
图2：TinyServe会话管理能力的可视化演示。左图：3D会话流动态显示了带有颜色编码的时间进展和会话节点的螺旋状会话演变。右图：8个关键指标（缓存命中、会话重用、迁移、负载均衡、响应时间、吞吐量、内存效率、可扩展性）的性能比较雷达图，展示了TinyServe相较于基线系统的全面优势。
图2：TinyServe会话管理能力的可视化演示。左图：3D会话流动态显示了带有颜色编码的时间进展和会话节点的螺旋状会话演变。右图：8个关键指标（缓存命中、会话重用、迁移、负载均衡、响应时间、吞吐量、内存效率、可扩展性）的性能比较雷达图，展示了TinyServe相较于基线系统的全面优势。
图3：会话管理性能，显示了不同会话大小下的缓存重用率和迁移开销。TinyServe实现了更高的缓存重用率，同时最小化了迁移成本。
图3：会话管理性能，显示了不同会话大小下的缓存重用率和迁移开销。TinyServe实现了更高的缓存重用率，同时最小化了迁移成本。
  • 全面对比评估：
    • 实验内容：在固定的2048 Token预算下，评估所有LongBench任务的准确率、延迟、吞吐量和KV缓存命中率。
    • 实验结果：表1和图4（雷达图）显示，TinyServe在各个模型规模（125M到1.3B）上，在保持或微幅提升准确率（如GPT2-345M上LongBench准确率提升1.1%）的同时，显著降低了延迟，并保持了最高的KV命中率（如TinyLLaMA上达到96.2%）。
    • 分析结论：由于其查询感知选择机制，TinyServe在延迟和准确性之间展现出卓越的权衡。
图4：TinyLLaMA（左）和GPT2-345M（右）的准确率、延迟、吞吐量和KV命中率雷达图。误差条显示了95%的置信区间。所有指标均是越高越好。
图4：TinyLLaMA（左）和GPT2-345M（右）的准确率、延迟、吞吐量和KV命中率雷达图。误差条显示了95%的置信区间。所有指标均是越高越好。
  • 跨模型加速分析：
    • 实验内容：在增加的上下文长度（高达32k Token）下评估端到端解码延迟。
    • 实验结果：图5显示，与FullCache基线相比，TinyServe平均实现了$2.1\times - 3.4\times$的加速，显著优于基于剪枝的基线方法。
    • 分析结论：TinyServe在长上下文解码中具有显著的加速优势。
图5：在32k提示长度和2048 Token预算下，不同基线之间的相对解码延迟加速（↓）。误差条代表5次运行的标准差。
图5：在32k提示长度和2048 Token预算下，不同基线之间的相对解码延迟加速（↓）。误差条代表5次运行的标准差。

  • 任务级全面评估：

    • 实验内容：使用GPT2-345M在2048预算下测试LongBench的所有五个具体任务。
    • 实验结果：表4表明，TinyServe在NarrativeQA、Qasper、TriviaQA、HotpotQA和GovReport上均取得了最高的加速比（2.03x - 2.31x），且准确率几乎没有损失。
    • 分析结论：TinyServe在各种具体长上下文任务中表现出一致的高效性。

  • KV缓存效率与访问分解：

    • 实验内容：可视化随时间变化的KV缓存利用率并分析内存访问模式。
    • 实验结果：图6显示TinyServe维持了更高的命中率和更少的Token驱逐。图7显示FullCache频繁达到HBM带宽限制，StreamingLLM仍有突发加载，而TinyServe由于页面级选择，内存访问模式更平滑且显著低于HBM阈值。
    • 分析结论：TinyServe保留了高相关性Token，避免了缓存刷新，实现了更高的有效重用率和更低的内存带宽压力。
图6：解码时间内的KV重用（上下文长度=32k，解码长度=2k）。TinyServe保持了更高的命中率和更少的Token驱逐。
图6：解码时间内的KV重用（上下文长度=32k，解码长度=2k）。TinyServe保持了更高的命中率和更少的Token驱逐。
图7：不同缓存策略下解码步骤的KV缓存访问带宽。由于不剪枝的完全KV重用，FullCache表现出持续的高内存带宽使用，经常达到HBM限制。StreamingLLM通过丢弃早期Token比FullCache有所改进，但仍表现出突发的内存加载。得益于查询感知的页面级KV选择，TinyServe显示出更平滑且显著更低的访问模式，保持远低于HBM阈值。垂直虚线标记了Token重用或解码阶段的关键过渡。
图7：不同缓存策略下解码步骤的KV缓存访问带宽。由于不剪枝的完全KV重用，FullCache表现出持续的高内存带宽使用，经常达到HBM限制。StreamingLLM通过丢弃早期Token比FullCache有所改进，但仍表现出突发的内存加载。得益于查询感知的页面级KV选择，TinyServe显示出更平滑且显著更低的访问模式，保持远低于HBM阈值。垂直虚线标记了Token重用或解码阶段的关键过渡。

  • 服务合成诊断：

    • 实验内容：通过三种合成任务测试系统行为：重复任务（测试注意力重用效率）、罕见Token召回（测试缓存配置对低频词的影响）、注意力混叠（测试重叠上下文中的位置混淆）。
    • 实验结果：表5显示，TinyServe在重复任务（97.5%）、罕见Token（93.8%）和混叠（88.2%）上的准确率非常接近FullCache，显著优于StreamingLLM和SoftPrune。
    • 分析结论：TinyServe在各种极端服务场景下表现出行为的一致性和鲁棒性。

  • 消融实验（插件与组件）：

    • 实验内容：表6评估了禁用单个插件（查询路由器、页面管理器、缓存融合、多GPU）的影响。表2评估了查询感知、页面级、边界框和融合内核组合的影响。
    • 实验结果：移除任何组件都会导致延迟增加或准确率/命中率下降。全配置的TinyServe实现了最佳的延迟（9.3ms）和命中率（91.7%）。
    • 分析结论：TinyServe的各个系统组件共同作用，缺一不可。

  • 页面大小与多GPU扩展：

    • 实验内容：表7测试了页面大小（4到64）对延迟和准确率的影响。表8评估了从1到8个A100 GPU的吞吐量扩展。
    • 实验结果：页面大小为16时提供了最佳的延迟（9.3ms）和准确率（命中率91.7%）权衡。多GPU扩展在8张卡上达到了7.68x的加速（96.0%效率）。
    • 分析结论：16是默认的最佳页面大小；系统具有近乎线性的扩展能力，验证了内核融合和GPU间缓存重用的有效性。

补充细节

超参数搜索与配置：系统进行了广泛的超参数搜索以获得最佳配置：页面大小（Page Size）测试了[4, 8, 16, 32, 64]，基于延迟-准确率权衡选择了16；选择比例（Selection Ratio）评估了[0.1, 0.2, 0.3, 0.5]，选择了0.3以获得最佳性能；批处理超时（Batch Timeout）测试了[10ms, 25ms, 50ms, 100ms]，为服务场景选择了50ms。

代码可用性：所有实验均固定随机种子（seed=42）以确保可重复性。完整的实现、数据集和评估脚本将公开，包括带有模块化插件系统的TinyServe框架、所有基线实现（vLLM, TGI, TensorRT-LLM适配器）、评估脚本以及预处理的数据集和模型检查点。

结论

本文介绍了TinyServe，一个轻量级且可扩展的服务系统，用于微型语言模型的高效推理和训练加速。TinyServe通过模块化支持Token选择、缓存稀疏性、融合注意力内核和训练优化，解决了LLM服务中的系统级瓶颈（如KV缓存饱和和解码时延迟）。系统的核心是一种查询感知的页面选择机制，它使用边界框元数据近似注意力相关性，以极小的开销实现选择性的KV访问。该方法在不影响准确率的情况下实现了显著的延迟和内存减少。通过其内核级优化、多GPU扩展和即插即用架构，TinyServe能够在资源受限的硬件上实现快速、可重复的实验。它为LLM服务研究提供了一个实用的基础，既支持微型模型的实时部署，也支持在没有全规模模型成本的情况下对服务机制进行有原则的评估。

参考文献引用汇总（方法细节部分）

  • 在说明部分KV Token对准确预测至关重要时，引用了【索引5，LoRA: Low-Rank Adaptation of Large Language Models+2022+ICLR+URL: https://openreview.net/forum?id=nZeVKeeFYf9】和【索引35 ，H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models+2023+URL: https://arxiv.org/abs/2306.14048】。原文描述为 ：“While prior work has demonstrated that only a fraction of KV tokens are critical for accurate predictions [5, 35], we observe that the set of critical tokens varies significantly across queries.”（尽管先前的研究表明只有一小部分KV Token对准确预测至关重要[5, 35]，但我们观察到，这一关键Token集合在不同的查询之间存在显著差异。）