The explanation for Multi-head Latent Attention https://www.stephendiehl.com/posts/post_transformers/#multi-head-latent-attention https://www.stephendiehl.com/p

The explanation for Multi-head Latent Attention https://www.stephendiehl.com/posts/post_transformers/#multi-head-latent-attention https://www.stephendiehl.com/posts/post_transformers/#multi-... does not match the definition in the DeepSeek-V2 paper https://arxiv.org/pdf/2405.04434#subsection.2.1 https://arxiv.org/pdf/2405.04434#subsection.2.1 MLA as developed by DeepSeek is a technique to reduce the memory footprint of the KV cache by storing only two vectors of size latent_dim and rope_dim per token and layer, instead of 2 * num_heads vectors of size head_dim. (DeepSeek-V3 has num_heads = 128 and head_dim = 128 vs latent_dim = 512 and rope_dim = 64, so a significant reduction https://arxiv.org/pdf/2412.19437#subsection.4.2 https://arxiv.org/pdf/2412.19437#subsection.4.2 ) What this article describes instead is some kind of two-step attention scheme I haven't seen before and that I think wouldn't work with causal masking (despite mask appearing in the example code) because either you allow an earlier token to attend to a latent that attended to a later token (creating backwards information flow) or the latents can only attend to a limited prefix of the sequence, after which they're frozen and useless. I wonder whether the author dreamed it up himself or whether someone else is actually using this somewhere.