题目内容

在大型语言模型进行长文档生成、持续对话等实时推理时，每一层、每一个注意力头的键值缓存（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 2 3 8
0.1 0.9 0.5 0.8 0.2 0.3
0.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 2 6
0.5 0.5 0.5 0.5
0.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。