题目内容

实现 MLA（Multi-head Latent Attention）的前向传播。MLA 是 DeepSeek-V2 提出的注意力机制，核心是把 Key 和 Value 压缩到低维潜变量，从而大幅减少推理时的 KV 缓存。

本题同时包含解耦 RoPE：内容部分（NoPE）与位置部分分开。为简化，旋转之后的位置向量直接给出，不需要自行实现 RoPE。

给定输入 $X$（$N \times d$）。按头计算，每头内容维度为 $d_h$。可见计算步骤为：

1. 低秩压缩（KV 共享潜变量，位置部分单独给出）：
   • $C_{kv} = X W_{DKV}$，$C_q = X W_{DQ}$。
   • 内容 Query / Key / Value 由潜变量上投影：$Q_c = C_q W_{UQ}$，$K_c = C_{kv} W_{UK}$，$V = C_{kv} W_{UV}$。
2. 位置部分 $Q_r$、$K_r$ 由输入中的对应投影得到（已经是旋转后的向量）。
3. 每个头把内容分数和位置分数相加，再做因果掩码与 softmax：

$j \le i$ 的位置保留，$j > i$（未来 token）置为 $-\infty$。然后 $\mathrm{weights} = \mathrm{softmax}(\mathrm{scores})$，$\mathrm{head} = \mathrm{weights}\, V$。

4. 各头拼接后经输出投影 $W_o$ 得到最终输出。

输入描述

第一行输入七个整数 $N, d, n_h, n_q, d_c, d_h, d_{rope}$（$1 \le N \le 12$，$1 \le d \le 8$，$1 \le n_h, n_q, d_c, d_h, d_{rope} \le 4$），依次表示序列长度、模型维度、头数、Query 潜变量维度、KV 潜变量维度、每头内容维度、位置维度。

随后按行输入浮点数矩阵，元素空格分隔，顺序为：

• $X$：$N \times d$
• $W_{DQ}$：$d \times n_q$
• $W_{UQ}$：$n_q \times (n_h \cdot d_h)$
• $W_{DKV}$：$d \times d_c$
• $W_{UK}$：$d_c \times (n_h \cdot d_h)$
• $W_{UV}$：$d_c \times (n_h \cdot d_h)$
• $W_{QR}$：$d \times (n_h \cdot d_{rope})$，用于得到已经旋转后的 $Q_r = X W_{QR}$
• $W_{KR}$：$d \times d_{rope}$，用于得到已经旋转后的 $K_r = X W_{KR}$
• $W_o$：$(n_h \cdot d_h) \times d$

位置向量不再单独输入，按上述两个投影得到。

输出描述

输出 $N$ 行，每行 $d$ 个浮点数，表示 MLA 前向结果，保留 $4$ 位小数。

样例 $1$

输入

2 2 1 1 1 1 1
1.0000 0.0000
0.0000 1.0000
1.0000
0.0000
1.0000
0.0000
1.0000
1.0000
1.0000
1.0000
0.0000
0.0000
1.0000
1.0000 0.0000

输出

0.0000 0.0000
0.5000 0.0000

说明

单头。$C_{kv} = X W_{DKV}$ 得到低维潜变量，内容分数与位置分数相加后再 softmax。因果掩码把上三角置为 $-\infty$，因此第二个 token 只能看到自己。

样例 $2$

输入

2 2 1 1 1 1 1
1.6243 -0.6118
-0.5282 -1.0730
0.8654
-2.3015
1.7448
-0.7612
0.3190
-0.2494
1.4621
-2.0601
-0.3224
-0.3841
1.1338
-1.0999 -0.1724

输出

2.3022 0.3609
1.5942 0.2499

说明

维度与样例 $1$ 相同，权重改为非平凡数值。$d_c = 1$，远小于展开后的 $n_h \cdot d_h$ 在更大配置下的规模；KV 仍先压到潜变量再上投影。