#CCSP2026C2. 大模型 KVCache 缓存系统 - 加强版

    ID: 735 Type: Default 3000ms 512MiB Tried: 8 Accepted: 1 Difficulty: 10 Uploaded By: Tags>CCSP贪心字符串Trie树树结构树链剖分数据结构单调栈颜色段均摊/区间推平并查集

大模型 KVCache 缓存系统 - 加强版

根据不同做法,我们设置了加强版,修改了 $n$ 和 $\sum m_i$ 的上限。

时间限制: 3.0 秒

空间限制: 512 MB

问题描述

如今主流的大语言模型均采用 Decoder-only 的 Transformer 架构。该架构的主要部件包括自注意力(Self Attention)和前馈神经网络(Feed Forward Network, FFN)。

大模型推理系统通常部署在 GPU 上,需要服务多个请求。为了方便,本题不考虑多并发场景,即所有请求存在确定的先后顺序,完成上一个请求的服务后才能继续服务下一个请求。一个请求经过 chat template 和 tokenizer 的处理后会形成一个 token 序列。对于该请求的 token 序列,系统的一次推理服务分为 Prefill 和 Decode 两个阶段。在 Prefill 阶段,模型会为该请求的每个 token 生成 Attention 所需的 Key/Value 张量,供后续 Decode 阶段使用,这个 Key/Value 张量称为 KVCache。每个请求的每个 token 都有与之对应的 KVCache,且在同一模型下大小固定。请注意,KVCache 与 token 前缀相关。例如,对于请求 ACAC 和请求 BCBC,两个请求中 token CC 的 KVCache 并不相同。

在 Prefill 阶段,每次计算 KVCache 的成本较高。但是,如果一个请求与某个历史请求拥有相同的 token 前缀,那么该请求就可以复用这些与历史请求前缀命中下的 KVCache,然后从未命中的第一个 token 开始,继续生成所有的 KVCache。

由此,可以利用内存相比于显存更便宜、容量更大等特点,设计一个 KVCache 缓存系统。该缓存系统分为显存和内存两个部分。具体地,本题将每个非空 token 前缀的最后一个 token 对应的 KVCache 视为一个缓存项,占用一个容量单位。该缓存系统中的内存缓存池容量是有限的,在任何时刻其存储的 KVCache 缓存项都不能超过一个上限 CC。在第一个请求到达前,内存缓存池为空。每个请求到达时按以下顺序进行操作:

  1. 对该请求进行 tokenize,得到 token 序列。
  2. 根据 token 序列,进行前缀匹配:在 KVCache 的内存缓存池中找到尽可能多的前缀 KVCache,得到该请求前缀命中的 token 数,然后把这些 KVCache 传输到显存中。请注意,这里仅匹配完全命中的前缀,首个缺失位置之后的缓存项不计入命中。
  3. 在 GPU 中执行 Prefill 操作,计算出所有未前缀命中的 token 对应的 KVCache。Prefill 结束后,显存中有该请求所有 token 对应的 KVCache。
  4. 考虑将该请求的 KVCache 写回内存的缓存池中。注意,一个已经被淘汰或在产生时未被接纳的历史 KVCache,之后不能凭空重新进入该缓存池;只有当后续请求的 Prefill 再次经过该前缀、重新得到这段 KVCache 时,才能再次接纳它。新生成的 KVCache 可以丢弃,不强制写入内存缓存池。简单地说,请求处理结束后,可以从“处理请求前仍在内存缓存池中的 KVCache”与“本次请求整条路径上的 KVCache”的并集中,任意选择不超过 CC 个保留在缓存池中,其余 KVCache 均淘汰。

为了方便,本题不考虑在显存中进行跨请求缓存。即每个请求结束后,显存中的 KVCache 都将被删除。

给定 nn 个先后到达的请求 $\mathrm{req}_1,\mathrm{req}_2,\cdots,\mathrm{req}_n$,以及 KVCache 缓存池的大小 CC,对于任意一个缓存池的淘汰策略,我们都能得到第 ii 个请求前缀命中的 token 数,记为 hih_i。那么一个 KVCache 缓存系统淘汰策略的 Cache 命中率即为:

$$\frac{\sum\limits_{i=1}^{n}h_i}{\sum\limits_{i=1}^{n}|\mathrm{req}_i|}$$

其中 ∣reqi∣|\mathrm{req}_i| 表示第 ii 个请求的 token 数。

在实际的 KVCache 缓存系统中,因为无法得知后续请求的状态,KVCache 的保存机制一般是 LRU(Least Recently Used,最近最少使用)或近似 LRU。然而,如果能得到所有请求及其到达顺序,那么就可以利用完整的未来信息决定接纳和淘汰策略。算出最优策略下的 KVCache 命中率对分析 KVCache 淘汰策略有着非常重要的意义。现在你需要完成这个任务。

输入格式

从标准输入读入数据。

第一行包含两个整数 n,Cn,C,分别表示请求数量,以及 KVCache 缓存池容量。

接下来 nn 行,第 ii 行由若干整数组成。前三个整数为 pi,ki,mip_i,k_i,m_i,接下来是 mim_i 个整数 ai,1,ai,2,⋯ ,ai,mia_{i,1},a_{i,2},\cdots,a_{i,m_i},表示第 ii 个请求的 token 序列由请求 pip_i 的前 kik_i 个 token 依次加上 ai,1,ai,2,⋯ ,ai,mia_{i,1},a_{i,2},\cdots,a_{i,m_i} 这 mim_i 个 token 形成。保证 0≤pi<i0\le p_i<i,0≤ki≤∣reqpi∣0\le k_i\le|\mathrm{req}_{p_i}|,0≤mi0\le m_i,ki+mi≥1k_i+m_i\ge1。特别地,请求 00 视作空序列,且 ∣req0∣=0|\mathrm{req}_0|=0。当 pi=0p_i=0 时,第 ii 个请求的 token 序列就是 ai,1,ai,2,⋯ ,ai,mia_{i,1},a_{i,2},\cdots,a_{i,m_i}。

输出格式

输出到标准输出。

输出共一行,包含两个非负整数 a,ba,b(需要保证 a≥0,b>0a\ge0,b>0),使用一个空格隔开。a/ba/b 表示最优情况下的 KVCache 命中率。请注意,输出需要化简,即保证 gcd⁡(a,b)=1\gcd(a,b)=1。

3 1
0 0 2 1 2
1 1 1 3
1 2 0
1 3

样例 1 解释

三个请求依次为:

[1, 2]
[1, 3]
[1, 2]

处理请求 1 后保留 KV(1)\mathrm{KV}(1),请求 2 和请求 3 均可命中一个 token,因此总命中数为 22。请求 1 不可能命中任何 token;由于缓存容量为 11,请求 2 和请求 3 分别至多命中一个 token,所以不存在更优策略。

4 1
0 0 1 10
0 0 1 20
1 1 0
2 1 0
1 4

样例 2 解释

四个请求依次为:

[10]
[20]
[10]
[20]

一种最优策略是在请求 1 后保留 KV(10)\mathrm{KV}(10),并 bypass 请求 2 产生的 KV(20)\mathrm{KV}(20),从而在请求 3 命中一次。由于 KV(20)\mathrm{KV}(20) 没有被保留,它不能在请求 4 前凭空恢复。反过来选择保留 KV(20)\mathrm{KV}(20) 也只能命中一次。

子任务

本题无数据梯度,需要通过全部数据获得所有分数。

对于所有数据,$1\le n\le {\color{red}5\times 10^5},\ 0\le C\le10^9,\ \sum m_i\le {\color{red}2\times 10^6},\ 1\le a_{i,j}\le 10^6$。

保证 $0\le p_i<i,\ 0\le k_i\le|\mathrm{req}_{p_i}|,\ 0\le m_i,\ k_i+m_i\ge1$。

提示

本题输入量较大,请采用效率较高的读入方式。