Skip to main content

recognize_attention

Function recognize_attention 

Source
pub fn recognize_attention(stmts: &[Stmt<'_>]) -> Option<AttentionSpec>
Expand description

Recognize the single-query GQA attention block. The desugar is a Splice that: Lets the output Seq (first New), zero-fills it, then a head-loop (While ≤ H whose body Lets a scores New(Seq)) computes scores = Q·K / √HD, a softmax (Pow(E, …)), and the weighted sum of V, SetIndexing the output. It does NOT mask with the -1e30 sentinel (that’s moe). Dims come from the loop bounds (H = head-loop bound, L = score-loop bound, HD = score-inner bound) and the kvh divisor (group = H/KV); the three externally-indexed operands, in program order, are Q, K, V. The masked up to P form sets pos.