下图标注了 MLA 的一种计算方式，橙色虚线部分可以被包裹到重计算中，入口为 q1，kv1 和 k_rope1（不保存 kv1 和 k_rope1 而是保存 kv1_and_k_rope 也可以，它们的大小一样，没区别）。flash_attn 比较特殊，除了保存输入之外还会保存输出。

如果要做 TP，需要保持 down_proj 为完整矩阵，在 up_proj 矩阵做列切分，在 o_proj 做行切分。如果要做 CP，需要在 attention 的位置插入 CP 策略。

MLA 推荐阅读： https://www.cnblogs.com/rossiXYZ/p/18827618