题目内容

在大语言模型（$LLM$）的自回归推理中，$KV\ Cache$ 用于缓存历史 $token$ 的 $Key$ 和 $Value$ 张量，其显存占用随序列长度线性增长。为支持长上下文，工业界采用动态稀疏策略：设定容量上限 $K$，仅保留注意力分数最高的 $K$ 个 $token$，超出时淘汰分数最低者。

任务描述

实现一个稀疏 $KV\ Cache$ 管理器，支持：

1.$ADD$：插入新 $token$ 的 $KV$ 对，附带注意力分数

2. 自动剪枝：当缓存大小超过$K$ 时，立即淘汰注意力分数最低的 $token$（分数相同则淘汰位置编号较小的）
4. $QUERY$：按位置编号升序返回当前缓存中的所有 $KV$

输入描述

第一行：$K$（整数，$1 \leq K \leq 100000$），最大缓存容量

$N$（整数，$1 \leq N \leq 100000$），操作数量

接下来$N$ 行，每行为以下两种格式之一：

1. ADD pos score [key_line] [value_line]

   • pos: 整数，token 位置（唯一，严格递增，$0 \leq pos \leq 10^9$）
   • score: 浮点数，注意力分数（$0.0 \leq score \leq 1000.0$）
   • key_line: 4 个浮点数，Key 向量
   • value_line: 4 个浮点数，Value 向量

2. QUERY

   • 查询当前缓存状态

输出描述

• $ADD$ 操作：

• 若插入后触发剪枝（缓存大小由 $K$ 变为$K+1$ 后淘汰一个），输出一行：PRUNED pos（pos 为被淘汰的位置编号）

  若未触发剪枝，无输出

• $QUERY$ 操作：

  第一行输出当前缓存中的 $token$数量 $M$

  接下来 $M$组，每组$3$ 行： pos score

  k1 k2 k3 k4

  v1 v2 v3 v4

  按 pos 严格升序排列

样例1

输入

3 7
ADD 0 5.0
1.0 2.0 3.0 4.0
0.1 0.2 0.3 0.4
ADD 1 2.0
2.0 2.0 2.0 2.0
0.5 0.5 0.5 0.5
ADD 2 8.0
3.0 3.0 3.0 3.0
0.9 0.9 0.9 0.9
QUERY
ADD 3 6.0
4.0 4.0 4.0 4.0
0.0 0.0 0.0 0.0
QUERY
ADD 4 9.0
5.0 5.0 5.0 5.0
1.0 1.0 1.0 1.0

输出

3
0 5.0
1.0 2.0 3.0 4.0
0.1 0.2 0.3 0.4
1 2.0
2.0 2.0 2.0 2.0
0.5 0.5 0.5 0.5
2 8.0
3.0 3.0 3.0 3.0
0.9 0.9 0.9 0.9
PRUNED 1
3
0 5.0
1.0 2.0 3.0 4.0
0.1 0.2 0.3 0.4
2 8.0
3.0 3.0 3.0 3.0
0.9 0.9 0.9 0.9
3 6.0
4.0 4.0 4.0 4.0
0.0 0.0 0.0 0.0
PRUNED 0

说明

• position 唯一且递增（按添加顺序）
• score 为浮点数，范围 [0, 100]
• 当需要剪枝时，淘汰 score 最低的 token（分数相同则淘汰 position 较小的）
• K/V 维度固定为 4（k_dim = v_dim = 4）