deepseek-ai

/

DeepSeek-R1

like 7.79k

Follow

DeepSeek 29.7k

感谢 DeepSeek 的开源和贡献，DeepSeek 确实很棒，所以最近也好好拜读了下这边的工作，然后恰巧最近在做一些 kv cache 压缩的工作，所以重点看了下这部分，然后发现了一些问题，来请教一下。

问题：在 DeepSeek-R1 论文中说采用了类似 DeepSeek-V3 中的 MLA 架构，文中说在这个架构中使用了低秩的方式存储 Key 和 Value，但是我看了 huggingface 中的代码，发现和论文中说的不一致，这是为什么？

详细说明：

按 DeepSeek-V3 一文中对该架构的说明，假设输入的序列长度为 n，hidden_dim 维度为 d，即为一个 n×d 的 hidden_states，如果低秩的维度为 d_c，位置编码的维度为 d_rope， 那么按文中的意思最后存储的 kv 大小为 n×d_rope + n×d_c，但是在代码中，存储的并不是这么大，实际存储的要比这个大，实际存储的是 key 为 n×d，value 为 n×d_v，即存储的 kv 大小为 n×d + n×d_v（对应文中图的话应该是 compress_kv 映射后上面没有位置编码的 k 和 v 外加有位置编码的 k 会被存储，即 n×d_rope + n×d_nope + n×d_c，其中 d_v 为一个可配置参数）。即对应文中说的 compress_kv 是实际存的 kv cache，实际并没存，并且这个的作用实际只是作为一个中间计算的过渡。

另外，从 config 中查看相关参数得知，d=7168，heads = 128，d_v_head = 128，故 d_v = 16384，这比 d 还要大，且 d_rope_head = 64，d_nope_head = 128, 则 d_k = 24576，这样一来，存储的 kv 相比 kv 同维度大小（同为24576）要小，但也仅仅是压缩了约 16.67%。

但对于论文中说的方式，存储的 kv 相关参数为 d_c = 512，d_rope_head = 64，则如果以 key 和 value 都为 24576 的未压缩状态计算，完全的 kv cache 大小为 49152n，论文中低秩存储时，存储的大小对应为 8704n，则压缩了约 82.29%。

然而，代码中的实现方式为压缩了 16.67% 的这种方式，并非为压缩了 82.29% 这种实现方式，且这种中间采用低秩计算这种方式，实际也增加了计算开销，并且如果仅仅是为了降低 value 的维度，以这个注意力层的隐维度大小，不采用这种低秩过度，直接采用 4 个矩阵映射的方式不是也可以吗？我大概算了下，这两种实现方式的权重大小基本差不多，这种采用多步计算的方式权重倒是多一点，还是说，采用这种计算方式，是为了降低计算复杂度？

以上，就是我对这块代码种实现和文中描述不一致地方的疑问，感觉 kv cache 这块在在实现上还未完全优化，但也不确定是不是我理解的有问题，所以特来请教。