Add challenge 114: Multi-Head Latent Attention Decode (Hard)

[claude]

Summary

Adds challenge 114: Multi-Head Latent Attention Decode (Hard) — one decode step of MLA, the attention variant used by DeepSeek-V2/V3.

Instead of caching per-head K and V, MLA caches a single low-rank latent vector c (rank 512) per position, shared by every head, plus a shared decoupled rotary key k_pe. The solver must use weight absorption — fold W_UK into the query and run the whole attention in latent space, applying W_UV once per head after the weighted sum — rather than reconstructing per-head keys and values, which would defeat the compressed cache.

q~_h  = q_nope_h @ W_UK[h]
s_h,t = (q~_h · c_t + q_pe_h · k_pe_t) / sqrt(head_dim + rope_dim)
a_h   = softmax(s_h)
o_h   = (sum_t a_h,t * c_t) @ W_UV[h]

Why this is interesting

• The shared cache means all heads stream the same kv_cache rows — reuse across heads is the whole performance story, unlike MHA/GQA where each head owns its keys.
• Three distinct work-distribution problems in one kernel: a per-head GEMV against W_UK, a softmax reduction over the cache, and a rank-512 accumulation followed by a second GEMV.
• Not an element-wise op, and no existing or pending challenge covers latent/compressed KV attention (checked against merged challenges and all open PRs; 114 is unclaimed).

Contents

• challenge.html — description, SVG of the shared latent cache layout, worked 2-head example, constraints
• challenge.py — reference impl (standard PyTorch ops only, CUDA/XLA-safe), 10 functional tests (single-position cache, zero query, zero cache, powers-of-2, non-powers-of-2, rank-512 realistic case), perf test at num_heads=128, seq_len=4,096, kv_lora_rank=512, head_dim=128, rope_dim=64 (~77 MB)
• Starters for all six frameworks

Validation

• The absorbed reference was cross-checked against a naive MLA formulation that explicitly reconstructs per-head K/V — max abs diff 3.6e-07.
• A CUDA solution was submitted via scripts/run_challenge.py on a T4: all tests passed (functional + performance). The solution is not committed.
• pre-commit run --all-files passes.

🤖 Generated with Claude Code