Skip to main content

recognize_causal_attention

Function recognize_causal_attention 

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

Recognize Let C be causal attention of Q over K and V with H heads KV groups from P — a block of T query rows: the Let from be P and its block check from ≥ 1 and from + (T − 1) ≤ L at the top, the output’s zero fill, then a row loop (to T) whose body binds the row’s length from + row − 1 and runs the single-query head loop over it.