Module ml::attention_advanced

#include <std/ml/attention_advanced.h> source

Advanced attention: Multi-head Latent Attention (MLA/EG-MLA), Shifted-Window (Swin) attention, and Gated DeltaNet.

Built on std/ml/attention.h's mha_forward/attention_forward and std/ml/tensor.h's Tensor. MLA and EG-MLA record gradients and run on devices like attention.h's ops (their weights require grad: they train through trainer.h); the shifted-window attention and Gated DeltaNet are forward-only host loops. Scope notes on each section below.

Structs

Mla	Compresses K/V into a shared low-rank latent c_kv = X @ Wdkv before up-projecting back to full K/V width — the KV-cache-memory-saving idea MLA is built around. (Decoupled rotary keys, the other half of the real DeepSeek-V2 design, are out of scope — no positional encoding is baked into this library's attention at all, see attention.h.) Q = X @ Wq [seqLen, dModel] c_kv = X @ Wdkv [seqLen, dLatent] K = c_kv @ Wuk [seqLen, dModel] V = c_kv @ Wuv [seqLen, dModel] out = mha_forward(Q, K, V, numHeads) (Mla::forward).
EgMla	EG-MLA (Extended-Group MLA): shares one down-projected latent c_kv across 'numGroups' independent up-projection groups (grouped-query attention's KV-sharing idea layered onto MLA's low-rank compression). Group g owns WukGroups[g]/WuvGroups[g] ([dLatent, groupDim], groupDim = dModel/numGroups) and runs 'headsPerGroup' ordinary attention heads over its own up-projected K/V slice; results are concatenated back into one [seqLen, dModel] tensor. WukGroups/WuvGroups hold one tensor per group; their length is the group count.
GatedDeltaNet	Linear-recurrent attention replacement (Yang et al. 2024): a [keyDim, valueDim] associative-memory state S is updated one timestep at a time with a delta rule (write the residual between the new value and what the current state already predicts for this key, rather than just accumulating) plus a scalar gate controlling how much of the old state decays away each step:

Functions

windowed_attention_forward Swin's two core ideas — partition into fixed-size local windows and attend only within a window, then cyclically shift the partition so the next layer's windows straddle this layer's boundaries — carried over to this library's 1D [seqLen, dim] sequence shape rather than Swin's original 2D image-patch grid (this library has no 2D spatial tensor type; see cnn.h's FeatureMap for the closest thing, which this doesn't use). 'seqLen' must be divisible by 'windowSize'. 'shift' (0 for the unshifted variant, typically windowSize/2 for the shifted one) cyclically rotates rows before windowing and rotates the result back after, so windows at the boundary contain rows from both original neighbors on the next layer.