Problems
Loading...
Step 2 of 11
Gated MLA
Hard

Gated Multi-head Latent Attention compresses each hidden token once and reconstructs keys and values when attention runs. Let hidden states XX have shape (B,S,D)(B,S,D), let HH be the number of heads, let Dh=D/HD_h=D/H, and let LL be the latent width. The projections produce

Q=XWqT∈RB×S×D.Q=XW_q^{\mathsf T}\in\mathbb{R}^{B\times S\times D}. C=XWcT∈RB×S×L.C=XW_c^{\mathsf T}\in\mathbb{R}^{B\times S\times L}. K=CWkT,V=CWvT,K=CW_k^{\mathsf T},\qquad V=CW_v^{\mathsf T},

where both KK and VV have shape (B,S,D)(B,S,D). Split QQ, KK, and VV into HH heads of width DhD_h. For each head, compute

A=softmax⁡(QKTDh+M),A=\operatorname{softmax}\left(\frac{QK^{\mathsf T}}{\sqrt{D_h}}+M\right),

where MM masks later sequence positions when causal mode is enabled and is zero otherwise. These layers use no positional encoding. Concatenate the head contexts into O~∈RB×S×D\widetilde{O}\in\mathbb{R}^{B\times S\times D}, then apply the channel gate and output projection:

Y=[sigmoid⁡(XWgT)⊙O~]WoT.Y=\left[\operatorname{sigmoid}(XW_g^{\mathsf T})\odot\widetilde{O}\right]W_o^{\mathsf T}.

Return a dictionary with output and latent_cache. The output has shape (B,S,D)(B,S,D), and the cache has shape (B,S,L)(B,S,L); both preserve the input floating-point dtype and device.

Loading visualization...

Examples

Input: hidden_states = [[[2]]], query_projection = [[1]], latent_down_projection = [[1]], key_up_projection = [[1]], value_up_projection = [[1]], output_gate_projection = [[0]], output_projection = [[1]], num_heads = 1, causal = true

Output: {"output": tensor([[[1.0]]]), "latent_cache": tensor([[[2.0]]])}

Explanation: The token is cached as the scalar latent value 2. Single-token attention reads that value, and the output gate scales it by 0.5.

Input: hidden states shape (1, 2, 2), latent width = 1, number of heads = 1, causal = True

Output: {"output": tensor of shape (1, 2, 2), "latent_cache": tensor of shape (1, 2, 1)}

Input: hidden states shape (1, 3, 2), latent width = 1, number of heads = 2, causal = False

Output: {"output": tensor of shape (1, 3, 2), "latent_cache": tensor of shape (1, 3, 1)}

Hint 1

Reshape projected tensors to batch, heads, sequence, head width before computing attention.

Hint 2

A causal mask excludes positions above the main diagonal before softmax.

Hint 3

Merge the head outputs back to model width before applying the channel gate and output projection.

Requirements

  • Return a dictionary with exactly output and latent_cache.
  • The output has the same shape, dtype, and device as the hidden states.
  • The cache shape is (batch, sequence, latent width).
  • Do not mutate any input tensor.

Constraints

  • Model width is divisible by the positive head count.
  • Projection matrices have compatible input and output widths.
  • The causal flag is a Boolean value.
  • All tensor inputs share a floating-point dtype and device.
Prev2 / 11Next