P5412. 预算约束下基于全局排序的 KV 缓存稀疏保留

场次：2026年9月9日 · 来源：CodeFun2000 P5412

题目描述

在大型语言模型进行长文档生成、持续对话等实时推理时，每一层、每一个注意力头的键值缓存（KV 缓存）会随着序列变长而不断增大，显存占用因此快速上升。为在固定显存预算内控制缓存规模，需要对 KV 缓存做稀疏保留：所有层、所有头都只保留一部分 token 的 KV，不同层和不同头允许保留的 token 数量可以不同。

给定层数 $L$ 、每层头数 $H$ 、序列长度 $N$ 、显存预算 $M$ ，以及每个注意力头在全部 $N$ 个位置上的注意力分数。层索引 $l$ 、头索引 $h$ 和 token 下标 $t$ 均从 0 开始。另有一个当前正在生成的位置 $C$ 。请确定每一层、每一个头最终保留哪些 token 下标。

选择必须满足以下规则：

对任意层 $l$ 和任意头 $h$ ，位置 $C$ 的 token 都必须被保留。
在所有位置不是 $C$ 的候选中，按注意力分数从高到低进行全局排序，分数更高的优先保留。
当注意力分数相同时，依次按层索引 $l$ 、头索引 $h$ 、token 下标 $t$ 从小到大排序。
最终保留的 token 总数不能超过预算 $M$ 。

具体地，先从预算 $M$ 中扣除所有层所有头必须保留的 $L \times H$ 个位置 $C$ ，得到剩余预算 $M - L \times H$ 。再对所有非位置 $C$ 的 $(l,h,t)$ 按注意力分数降序、层索引升序、头索引升序、token 下标升序进行排序，取排序结果的前 $M - L \times H$ 项。若在并列分数处不能全部放入剩余预算，则按上述顺序截断。最后将选中的 token 下标与每层每头必留的位置 $C$ 合并，并按下标升序输出。

约束条件

层数 $L$ 和头数 $H$ 满足 $1 \le L,H \le 32$ 。
序列长度 $N$ 满足 $1 \le N \le 200$ 。
显存预算 $M$ 满足 $L \times H \le M \le L \times H \times N$ 。
当前生成位置 $C$ 满足 $0 \le C < N$ 。
所有注意力分数为浮点数，且满足 $0 \le \text{score} \le 1$ 。

输入格式

第一行包含四个整数 $L$ 、 $H$ 、 $N$ 、 $M$ ，相邻两个数之间用一个空格分隔。

接下来有 $L$ 行。第 $i$ 行包含 $H \times N$ 个浮点数，空格分隔，表示第 $i$ 层的全部注意力分数。该行按头索引从小到大的顺序拼接：前 $N$ 个数属于头 0，接下来的 $N$ 个数属于头 1，依此类推，直到头 $H-1$ 。

最后一行包含一个整数 $C$ ，表示当前正在生成的位置。

输出格式

输出 $L$ 行。第 $i$ 行对应第 $i$ 层，包含 $H$ 个字符串，相邻字符串之间用一个空格分隔。第 $j$ 个字符串表示第 $i$ 层第 $j$ 个头实际保留的 token 下标；下标按升序排列，并用逗号 , 连接。若该头只保留一个下标，则该字符串不含逗号。

示例

示例 1

输入：

2 3 8
1 0.9 0.5 0.8 0.2 0.3
7 0.6 0.4 0.0 0.1 0.2
2

输出：

1,2 0,2
0,1,2 2

说明：

每层每头必须保留位置 2，共保留 $L \times H = 4$ 个 token，剩余预算为 $M - L \times H = 8 - 4 = 4$ 。

在不是位置 2 的候选中，按分数从高到低取 4 个。分数最高的依次是：层 0 头 0 的 token 1，分数 0.9；层 0 头 1 的 token 0，分数 0.8；层 1 头 0 的 token 0，分数 0.7；层 1 头 0 的 token 1，分数 0.6。

将它们与必留位置 2 合并后，层 0 头 0 保留 [1,2]，层 0 头 1 保留 [0,2]，层 1 头 0 保留 [0,1,2]，层 1 头 1 只保留 [2]。按下标升序连接后得到输出。

示例 2

输入：

2 2 6
5 0.5 0.5 0.5
5 0.5 0.5 0.5
0

输出：

0,1 0,1
0 0

说明：

每层每头必须保留位置 0，共保留 $L \times H = 4$ 个 token，剩余预算为 $M - L \times H = 6 - 4 = 2$ 。

所有不是位置 0 的候选分数都相同，为 0.5，因此按层索引、头索引、token 下标升序排序。前 2 个候选是层 0 头 0 的 token 1 和层 0 头 1 的 token 1。

因此层 0 的两个头保留 [0,1]，层 1 的两个头只保留 [0]。按下标升序连接后得到输出。

示例 3

输入：

1 2 2 2
0.5 0.9 0.4 0.8
0

输出：

0 0

说明：

每层每头必须保留位置 0，共保留 $L \times H = 2$ 个 token。此时剩余预算为 $M - L \times H = 2 - 2 = 0$ 。

因为没有额外预算，所有不是位置 0 的候选都不会被保留。因此每个头只保留位置 0，输出为两个 0。

解题思路

第一步：理解问题本质

先把输入整理成一张表：层的编号是行，头的编号是列，每个格子里有 $N$ 个位置，位置 $t$ 的“价值”是它的注意力分数 $\text{score}[l][h][t]$ 。整张表一共有

T = L \times H \times N

个候选格子，而显存预算 $M$ 是按保留的 token 个数计的。

题目给的规则可以拆成两层：

硬约束：每个 $(l,h)$ 的当前位置 $C$ 必须保留。这一条与分数无关，先把 $L\times H$ 个名额花掉，剩下的 $\text{extra} = M - L \times H$ 才是可以自由分配的预算。
全局竞争：剩下的 extra 个名额不再预先分配给某个头，而是在所有 $(l,h,t)$ 、 $t \ne C$ 的候选里一起排队，分数高的先拿。所以某个头可能一个名额都拿不到——示例 1 中 $(l,h)=(1,1)$ 这个头就只保留了 $C$ 本身。

这个“全局共享预算”是本题的关键。如果误以为预算是“每层每头各分 $M/(LH)$ 个”，示例 1 的输出就会完全对不上。

排队规则是一个四级字典序：

(-\text{score},\ l,\ h,\ t)

按这个元组升序排，就是“分数从高到低；分数相同看层号小的；再看头号小的；最后看下标小的”。

第二步：暴力解法

最直接的想法是枚举“从 $T-LH$ 个候选里挑 extra 个”的所有组合，看哪种挑法分数之和最大：

from itertools import combinations

candidates = []  # (score, l, h, t)
for l in range(L):
    for h in range(H):
        for t in range(N):
            if t != C:
                candidates.append((score[l][h][t], l, h, t))

best = -1
for chosen in combinations(candidates, extra):
    total = sum(item[0] for item in chosen)
    if total > best:
        best = total

组合数是

\binom{T-LH}{\text{extra}}

本题 $T$ 最大 $32\times32\times200=204800$ ，extra 甚至可能接近 $2\times10^5$ ，这个数量级完全不可行。

不过暴力枚举给我们一个重要的观察：候选之间没有任何耦合。每个候选能不能保留，只取决于预算还剩几个名额，而不取决于已经选了谁。所以“每一步都拿当前分数最高的那个”一定可以得到最优解——这正是贪心成立的条件。

第三步：优化解法——按优先级维护一个小根堆

既然每次拿当前最优就行，可以用一个大小为 extra 的堆来维护“目前入选的集合”：

import heapq

# rank 越大表示优先级越高：分数大的优先，同分时层/头/下标小的优先
def rank(item):
    score, l, h, t = item
    return (score, -l, -h, -t)

heap = []  # 小根堆，堆顶是当前入选集合里优先级最低的那个
for item in candidates:
    if len(heap) < extra:
        heapq.heappush(heap, (rank(item), item))
    elif rank(item) > heap[0][0]:
        heapq.heapreplace(heap, (rank(item), item))

复杂度是 $O(T \log \text{extra})$ 。当 extra 很小时它比全排序快；但本题 extra 可能和 $T$ 同量级，堆的常数又不小，代码里还要构造两次键（一次比、一次存）。所以这里只用它说明思路，最终实现换成一次排序。

第四步：最优解法——一次排序 + 切片

把候选直接写成四元组：

(-float(score), layer, head, index)

Python 的元组比较是从左到右逐项比较，正好就是字典序。于是：

第一项取负，sort() 的升序就变成了“分数降序”；
后三项保持原值，升序正好对应“层号小、头号小、下标小优先”。

这一步为什么能直接用浮点数比较，而不需要写 abs(a-b) < 1e-9 这样的容差判断？因为题目说的“分数相同”指的是输入里写出来的数值相同，同样的十进制字面量解析成同一个浮点数，元组比较自然走到后三项；而排序内部只做 < 比较，不会因为浮点误差把两个真正不同的分数误判成相等。相反，如果自己加一个 EPS，反而会把 $0.5000001$ 和 $0.5$ 这种确实有先后的候选抹平。

排序之后：

取前 extra 个（extra = 0 时切片为空）；
按下标归到 kept[layer][head] 里；
每个头把必留的 $C$ 并进去，再按下标升序排序；
用 , 连接，H 个字符串用空格连接，共输出 $L$ 行。

注意第 3 步：入选顺序是分数顺序，而输出要求的是下标顺序，这两者不能混。一个头保留一个下标时，",".join 的结果自然没有逗号，不需要特判。

完整代码实现

"""
P5412. 预算约束下基于全局排序的 KV 缓存稀疏保留 - 多关键字排序 + 全局 Top-K

输入格式：
- 第 1 行：四个整数 L、H、N、M；
- 接下来 L 行：每行 H*N 个浮点数，按头编号拼接；
- 最后一行：一个整数 C。

输出格式：
- 共 L 行，第 i 行是 H 个用空格分隔的字符串，每个字符串是该头保留的
  token 下标，升序、用逗号连接。

时间复杂度：O(T log T)，T = L*H*N
空间复杂度：O(T)
"""

import sys


def solve():
    data = sys.stdin.buffer.read().split()
    if not data:
        return

    layers = int(data[0])
    heads = int(data[1])
    length = int(data[2])
    budget = int(data[3])

    # C 单独占最后一行，所以直接取最后一个 token；
    # 分数是一段连续的 L*H*N 个 token，不必关心它们怎么换行。
    current = int(data[-1])
    total = layers * heads * length
    scores = data[4:4 + total]

    # 必留项先占预算：每个 (层, 头) 都要保留 C，一共 L*H 个。
    extra = budget - layers * heads
    if extra < 0:
        extra = 0

    # 元组第一项取分数相反数 -> 升序排序等于分数降序；
    # 后三项是层、头、下标，天然升序，正好对应并列时的比较顺序。
    candidates = []
    pos = 0
    for layer in range(layers):
        for head in range(heads):
            for index in range(length):
                if index != current:
                    candidates.append((-float(scores[pos]), layer, head, index))
                pos += 1

    candidates.sort()

    # 归组：kept[layer][head] 收集该头被选中的非必留下标
    kept = [[[] for _ in range(heads)] for _ in range(layers)]
    for _, layer, head, index in candidates[:extra]:
        kept[layer][head].append(index)

    lines = []
    for layer in range(layers):
        cells = []
        for head in range(heads):
            # 必留的 C 并进来后统一升序：题目要求按下标输出，不是按分数输出
            bucket = kept[layer][head]
            bucket.append(current)
            bucket.sort()
            cells.append(",".join(str(x) for x in bucket))
        lines.append(" ".join(cells))

    sys.stdout.write("\n".join(lines) + "\n")


if __name__ == "__main__":
    solve()

NumPy 版本

以下版本使用 NumPy。只有在评测环境允许第三方库时才能提交；CodeFun2000 的普通 Python 环境不一定安装 NumPy。

思路与纯 Python 版完全一致，只是把「摊平 + 排序 + 归组」整段交给 NumPy：分数表直接读成 (L, H, N) 的三维数组，用布尔掩码一次摊平出 t ≠ C 的全部候选，np.lexsort 一次完成「分数降序 + 层升序 + 头升序 + 下标升序」四级排序，再用布尔索引把入选的 token 归回每个头。

"""
P5412. 预算约束下基于全局排序的 KV 缓存稀疏保留 - NumPy 版

把“分数表”直接读成一个 (L, H, N) 的三维数组，用 np.lexsort 一次完成
“分数降序 + 层升序 + 头升序 + 下标升序”的多关键字排序，
再用布尔索引把入选的 token 归到 (层 × 头) 的集合里。

时间复杂度：O(T log T)，T = L*H*N；空间复杂度：O(T)。
"""

import sys

import numpy as np


def solve():
    data = sys.stdin.buffer.read().split()
    if not data:
        return

    layers = int(data[0])
    heads = int(data[1])
    length = int(data[2])
    budget = int(data[3])

    # C 单独占最后一行；中间这段是 L*H*N 个分数，按“层 → 头 → 下标”的顺序排列
    current = int(data[-1])
    total = layers * heads * length
    scores = np.array(list(map(float, data[4:4 + total])), dtype=np.float64)
    scores = scores.reshape(layers, heads, length)

    # 必留项先占预算：每个 (层, 头) 都要保留 C，一共 L*H 个
    extra = max(0, budget - layers * heads)

    # 候选是 t != C 的全部 (层, 头, 下标)；用布尔掩码一次摊平
    layer_id, head_id, token_id = np.meshgrid(
        np.arange(layers), np.arange(heads), np.arange(length), indexing="ij"
    )
    keep = token_id != current
    cand_layer = layer_id[keep]
    cand_head = head_id[keep]
    cand_token = token_id[keep]
    cand_score = scores[keep]

    # np.lexsort 的关键字从“最次要”排到“最主要”：最后一列才是主键。
    # 主键取分数的相反数 => 升序排序等价于分数降序；后三列天然升序。
    order = np.lexsort((cand_token, cand_head, cand_layer, -cand_score))
    chosen = order[:extra]                    # extra = 0 时切片为空

    # kept 按 (层*H + 头, 下标) 存，最后一列（必留的 C）全部置真
    kept = np.zeros((layers * heads, length), dtype=bool)
    cell = cand_layer[chosen] * heads + cand_head[chosen]
    kept[cell, cand_token[chosen]] = True
    kept[:, current] = True

    # 输出要求按下标升序（不是按分数顺序），所以直接取每行的真值位置
    lines = []
    for layer in range(layers):
        cells = []
        for head in range(heads):
            tokens = np.flatnonzero(kept[layer * heads + head])
            cells.append(",".join(str(int(x)) for x in tokens))
        lines.append(" ".join(cells))

    sys.stdout.write("\n".join(lines) + "\n")


if __name__ == "__main__":
    solve()

示例推演

示例 1

$L=2$ 、 $H=2$ 、 $N=3$ 、 $M=8$ 、 $C=2$ 。必留项有 $L\times H=4$ 个，所以

\text{extra} = 8 - 4 = 4

把两行分数按头拆开（每行前 $N$ 个属于头 0，后 $N$ 个属于头 1）：

层	头	$t=0$	$t=1$	$t=2$
0	0	0.1	0.9	0.5（必留）
0	1	0.8	0.2	0.3（必留）
1	0	0.7	0.6	0.4（必留）
1	1	0.0	0.1	0.2（必留）

除 $t=2$ 以外的 8 个候选按 $(-\text{score}, l, h, t)$ 升序排列：

名次	分数	层	头	下标
1	0.9	0	0	1
2	0.8	0	1	0
3	0.7	1	0	0
4	0.6	1	0	1
5	0.2	0	1	1
6	0.1	0	0	0
7	0.1	1	1	1
8	0.0	1	1	0

第 6、7 名都是 $0.1$ ，层号 $0<1$ ，所以 $(0,0,0)$ 排在 $(1,1,1)$ 前面。取前 4 名，再各自并上必留的 $2$ ：

层 0 头 0： $\{1\}\cup\{2\}=\{1,2\}$ → 1,2
层 0 头 1： $\{0\}\cup\{2\}=\{0,2\}$ → 0,2
层 1 头 0： $\{0,1\}\cup\{2\}=\{0,1,2\}$ → 0,1,2
层 1 头 1： $\varnothing\cup\{2\}=\{2\}$ → 2

输出：

1,2 0,2
0,1,2 2

示例 2

$L=2$ 、 $H=2$ 、 $N=2$ 、 $M=6$ 、 $C=0$ ，于是 $\text{extra}=6-4=2$ 。

所有候选分数都是 $0.5$ ，排序完全由 $(l,h,t)$ 决定。 $t\ne 0$ 的候选按顺序是：

(0,0,1),\ (0,1,1),\ (1,0,1),\ (1,1,1)

前 2 个是 $(0,0,1)$ 与 $(0,1,1)$ ，都落在层 0：

层 0 头 0： $\{1\}\cup\{0\}=\{0,1\}$ → 0,1
层 0 头 1： $\{1\}\cup\{0\}=\{0,1\}$ → 0,1
层 1 头 0、层 1 头 1：都只有必留的 $0$ → 0

输出：

0,1 0,1
0 0

示例 3

$L=1$ 、 $H=2$ 、 $N=2$ 、 $M=2$ 、 $C=0$ ，于是 $\text{extra}=2-2=0$ 。

额外名额为 0，候选一个都不入选，两个头都只剩必留的位置 $0$ ：

0 0

这也说明：额外名额为 0 时输出不是空字符串，每个头仍然要输出它的必留下标。

复杂度分析

解法	时间复杂度	空间复杂度	说明
枚举组合	$O\!\left(\binom{T-LH}{\text{extra}}\cdot \text{extra}\right)$	$O(\text{extra})$	组合数爆炸，完全不可行
逐轮线性扫描	$O(\text{extra}\cdot T)$	$O(T)$	`extra` 与 $T$ 同为 $10^5$ 量级时约 $4\times10^{10}$
小根堆部分选择	$O(T\log \text{extra})$	$O(\text{extra})$	`extra` 很小时更优，但常数与可读性较差
一次全排序（本题）	$O(T\log T)$	$O(T)$	$T=LHN\le 204800$ ，实测不足 0.2 秒

其中 $T=L\times H\times N$ 。 $T$ 最大为 $32\times32\times200=204800$ ，排序完全在可接受范围内。

易错点总结

1. 把预算按层按头均分

预算 $M$ 是全局的。如果按 $M/(LH)$ 给每个头均分，示例 1 里每个头只有 2 个名额，扣掉必留的 $C$ 就只剩 1 个额外名额：层 1 的两个头会分别输出 0,2 和 1,2，而正确答案是 0,1,2 和 2。

2. 忘记先扣掉必留的 $C$

候选集合里必须排除 $t=C$ ，同时预算要减去 $L\times H$ 。只做其中一件，要么重复计数，要么名额多算/少算。

3. 排序键的方向写反

分数要降序（键里取负），层、头、下标要升序。写成 (score, l, h, t) 会把分数最低的排到最前面；把后三项也取负，又会变成“层号大的优先”。

4. 输出顺序误用分数顺序

入选是按分数排的，输出却要求每个头按下标升序。必须归组后重新 sort()，不能直接按入选顺序打印。

5. `extra = 0` 时输出空字符串

即使一个额外名额都没有，每个头也保留着 $C$ ，输出仍然要有 $L$ 行、每行 $H$ 个字符串。

6. 输入布局读错

第 $i$ 层的 $H\times N$ 个数是按头拼接的一整行，不是 $N$ 行。按“每行一个 token”去解析会整体错位。

7. 用整数截断分数来比较

把分数乘 100 取整再比较，会把 $0.901$ 和 $0.899$ 判成并列，破坏“分数高的优先”。直接比较解析出来的浮点数即可。

8. 单个下标时多打了逗号

",".join(["2"]) 的结果就是 "2"，不需要额外处理；但若自己拼字符串（如 f"{x},"），就会多出一个逗号。

扩展思考

每个头有独立配额：如果题目改成“第 $h$ 个头最多保留 $q_h$ 个”，全局竞争就退化成每头各自取前 $q_h$ 名， $L\times H$ 次小规模部分选择即可。
只取前 $k$ 名且 $k$ 很小：用堆或快速选择（quickselect）可以做到 $O(T)$ ，但需要自己处理四级比较，代码复杂度会上升。
分数会变化：如果生成过程中分数在更新，需要反复取 Top-K，可以用堆增量维护，每次只调整发生变化的那些候选。
预算单位不是“个数”：如果不同层的 KV 缓存大小不同（例如分组查询注意力），问题就变成按“权重”而不是“个数”分配预算，贪心仍然可用，但要按“分数/权重比”排序，和分数背包的做法一致。
答案为什么唯一：四级键把 $(l,h,t)$ 全部包含进去了，任意两个候选的键都不会相同，所以排序结果唯一，不需要依赖排序算法的稳定性。

P5412. 预算约束下基于全局排序的 KV 缓存稀疏保留

题目描述

输入格式

输出格式

示例

示例 1

示例 2

示例 3

解题思路

第一步：理解问题本质

第二步：暴力解法

第三步：优化解法——按优先级维护一个小根堆

第四步：最优解法——一次排序 + 切片

完整代码实现

NumPy 版本

示例推演

示例 1

示例 2

示例 3

复杂度分析

易错点总结

1. 把预算按层按头均分

2. 忘记先扣掉必留的 $C$

3. 排序键的方向写反

4. 输出顺序误用分数顺序

5. `extra = 0` 时输出空字符串

6. 输入布局读错

7. 用整数截断分数来比较

8. 单个下标时多打了逗号

扩展思考

相关题目

自测

题目描述​

输入格式​

输出格式​

示例​

示例 1​

示例 2​

示例 3​

解题思路​

第一步：理解问题本质​

第二步：暴力解法​

第三步：优化解法——按优先级维护一个小根堆​

第四步：最优解法——一次排序 + 切片​

完整代码实现​

NumPy 版本​

示例推演​

示例 1​

示例 2​

示例 3​

复杂度分析​

易错点总结​

1. 把预算按层按头均分​

2. 忘记先扣掉必留的 CCC​

3. 排序键的方向写反​

4. 输出顺序误用分数顺序​

5. extra = 0 时输出空字符串​

6. 输入布局读错​

7. 用整数截断分数来比较​

8. 单个下标时多打了逗号​

扩展思考​

相关题目​

自测​

题目描述

输入格式

输出格式

示例

示例 1

示例 2

示例 3

解题思路

第一步：理解问题本质

第二步：暴力解法

第三步：优化解法——按优先级维护一个小根堆

第四步：最优解法——一次排序 + 切片

完整代码实现

NumPy 版本

示例推演

示例 1

示例 2

示例 3

复杂度分析

易错点总结

1. 把预算按层按头均分

2. 忘记先扣掉必留的 $C$

3. 排序键的方向写反

4. 输出顺序误用分数顺序

5. `extra = 0` 时输出空字符串

6. 输入布局读错

7. 用整数截断分数来比较

8. 单个下标时多打了逗号

扩展思考

相关题目

自测