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.