vyre_libs/nn/attention/
gated_delta.rs

1//! Recurrent gated delta-rule attention with explicit matrix state.
2
3use thiserror::Error;
4use vyre_foundation::ir::{DataType, Expr, Node, Program, UnOp};
5
6use super::gated_delta_layout::{self, GatedDeltaSpec};
7use crate::region::wrap_anonymous;
8
9const OP_ID: &str = "vyre-libs::nn::recurrent_gated_delta";
10
11/// Invalid recurrent gated delta construction.
12#[derive(Debug, Clone, PartialEq, Eq, Error)]
13pub enum RecurrentGatedDeltaError {
14    /// A required dimension is zero.
15    #[error(
16        "recurrent gated delta requires nonzero batch, sequence, head, key, and value dimensions"
17    )]
18    EmptyShape,
19    /// Value heads cannot evenly repeat key heads.
20    #[error("recurrent gated delta value_heads={value_heads} must be divisible by key_heads={key_heads}")]
21    InvalidHeadGrouping {
22        /// Key/query head count before repetition.
23        key_heads: u32,
24        /// Value/state head count.
25        value_heads: u32,
26    },
27    /// A flattened tensor size exceeds u32 indexing.
28    #[error("recurrent gated delta tensor element count overflows u32; split the tensor")]
29    ElementCountOverflow,
30    /// Source dtype lacks the required conversion contract.
31    #[error("recurrent gated delta supports F16, BF16, or F32 activations; got {dtype:?}")]
32    UnsupportedDtype {
33        /// Rejected activation dtype.
34        dtype: DataType,
35    },
36}
37
38/// Build a token-recurrent gated delta update.
39#[allow(clippy::too_many_arguments)]
40pub fn recurrent_gated_delta(
41    query: &str,
42    key: &str,
43    value: &str,
44    decay_log: &str,
45    beta_logits: &str,
46    state_input: &str,
47    output: &str,
48    state_output: &str,
49    batch: u32,
50    sequence: u32,
51    key_heads: u32,
52    value_heads: u32,
53    key_dim: u32,
54    value_dim: u32,
55    eps: f32,
56    dtype: DataType,
57) -> Result<Program, RecurrentGatedDeltaError> {
58    recurrent_gated_delta_impl(&GatedDeltaSpec {
59        query,
60        key,
61        value,
62        decay_log,
63        beta_logits,
64        state_input,
65        output,
66        state_output,
67        batch,
68        sequence,
69        key_heads,
70        value_heads,
71        key_dim,
72        value_dim,
73        eps,
74        dtype,
75    })
76}
77
78/// Build a fixed-size-64 chunk schedule for gated delta prefill.
79///
80/// The schedule retains the exact recurrent dependency inside each causal
81/// lower-triangular tile. Its final tile is padded structurally and guarded,
82/// so padding cannot read inputs, modify state, or appear in the output.
83#[allow(clippy::too_many_arguments)]
84pub fn chunked_gated_delta(
85    query: &str,
86    key: &str,
87    value: &str,
88    decay_log: &str,
89    beta_logits: &str,
90    state_input: &str,
91    output: &str,
92    state_output: &str,
93    batch: u32,
94    sequence: u32,
95    key_heads: u32,
96    value_heads: u32,
97    key_dim: u32,
98    value_dim: u32,
99    eps: f32,
100    dtype: DataType,
101) -> Result<Program, RecurrentGatedDeltaError> {
102    super::gated_delta_chunked::chunked_gated_delta_impl(&GatedDeltaSpec {
103        query,
104        key,
105        value,
106        decay_log,
107        beta_logits,
108        state_input,
109        output,
110        state_output,
111        batch,
112        sequence,
113        key_heads,
114        value_heads,
115        key_dim,
116        value_dim,
117        eps,
118        dtype,
119    })
120}
121
122/// Token-recurrent gated delta arithmetic.
123///
124/// Q and K are L2-normalized in F32. `decay_log` is exponentiated and
125/// `beta_logits` is passed through sigmoid. Matrix state remains F32; activation
126/// output converts once to `dtype`. `state_input` is preserved and
127/// `state_output` receives the continued generation.
128fn recurrent_gated_delta_impl(
129    spec: &GatedDeltaSpec<'_>,
130) -> Result<Program, RecurrentGatedDeltaError> {
131    let counts = spec.counts()?;
132    let GatedDeltaSpec {
133        query,
134        key,
135        value,
136        decay_log,
137        beta_logits,
138        state_input,
139        output,
140        state_output,
141        sequence,
142        key_heads,
143        value_heads,
144        key_dim,
145        value_dim,
146        eps,
147        ref dtype,
148        ..
149    } = *spec;
150
151    let qk_index =
152        |dim: Expr| gated_delta_layout::qk_index(sequence, key_heads, key_dim, Expr::var("token"), dim);
153    let state_index = |key_index: Expr, value_index: Expr| {
154        gated_delta_layout::state_index(key_dim, value_dim, key_index, value_index)
155    };
156    let value_index = |dim: Expr| {
157        gated_delta_layout::value_index(sequence, value_heads, value_dim, Expr::var("token"), dim)
158    };
159    let scalar_index = gated_delta_layout::scalar_index(sequence, value_heads, Expr::var("token"));
160    let output_index = value_index(Expr::var("value_index"));
161
162    let init_state = Node::loop_for(
163        "key_index",
164        Expr::u32(0),
165        Expr::u32(key_dim),
166        vec![Node::loop_for(
167            "value_index",
168            Expr::u32(0),
169            Expr::u32(value_dim),
170            vec![Node::Store {
171                buffer: state_output.into(),
172                index: state_index(Expr::var("key_index"), Expr::var("value_index")),
173                value: Expr::load(
174                    state_input,
175                    state_index(Expr::var("key_index"), Expr::var("value_index")),
176                ),
177            }],
178        )],
179    );
180    let norm_sums = vec![
181        Node::let_bind("query_sum", Expr::f32(0.0)),
182        Node::let_bind("key_sum", Expr::f32(0.0)),
183        Node::loop_for(
184            "key_index",
185            Expr::u32(0),
186            Expr::u32(key_dim),
187            vec![
188                Node::let_bind(
189                    "query_component",
190                    Expr::cast(
191                        DataType::F32,
192                        Expr::load(query, qk_index(Expr::var("key_index"))),
193                    ),
194                ),
195                Node::let_bind(
196                    "key_component",
197                    Expr::cast(
198                        DataType::F32,
199                        Expr::load(key, qk_index(Expr::var("key_index"))),
200                    ),
201                ),
202                Node::assign(
203                    "query_sum",
204                    Expr::add(
205                        Expr::var("query_sum"),
206                        Expr::mul(Expr::var("query_component"), Expr::var("query_component")),
207                    ),
208                ),
209                Node::assign(
210                    "key_sum",
211                    Expr::add(
212                        Expr::var("key_sum"),
213                        Expr::mul(Expr::var("key_component"), Expr::var("key_component")),
214                    ),
215                ),
216            ],
217        ),
218        Node::let_bind(
219            "query_scale",
220            Expr::mul(
221                Expr::UnOp {
222                    op: UnOp::InverseSqrt,
223                    operand: Box::new(Expr::add(Expr::var("query_sum"), Expr::f32(eps))),
224                },
225                Expr::UnOp {
226                    op: UnOp::InverseSqrt,
227                    operand: Box::new(Expr::f32(key_dim as f32)),
228                },
229            ),
230        ),
231        Node::let_bind(
232            "key_scale",
233            Expr::UnOp {
234                op: UnOp::InverseSqrt,
235                operand: Box::new(Expr::add(Expr::var("key_sum"), Expr::f32(eps))),
236            },
237        ),
238    ];
239    let decay_state = Node::loop_for(
240        "key_index",
241        Expr::u32(0),
242        Expr::u32(key_dim),
243        vec![Node::loop_for(
244            "value_index",
245            Expr::u32(0),
246            Expr::u32(value_dim),
247            vec![Node::Store {
248                buffer: state_output.into(),
249                index: state_index(Expr::var("key_index"), Expr::var("value_index")),
250                value: Expr::mul(
251                    Expr::load(
252                        state_output,
253                        state_index(Expr::var("key_index"), Expr::var("value_index")),
254                    ),
255                    Expr::var("decay"),
256                ),
257            }],
258        )],
259    );
260    let value_update = Node::loop_for(
261        "value_index",
262        Expr::u32(0),
263        Expr::u32(value_dim),
264        vec![
265            Node::let_bind("memory", Expr::f32(0.0)),
266            Node::loop_for(
267                "key_index",
268                Expr::u32(0),
269                Expr::u32(key_dim),
270                vec![Node::assign(
271                    "memory",
272                    Expr::add(
273                        Expr::var("memory"),
274                        Expr::mul(
275                            Expr::load(
276                                state_output,
277                                state_index(Expr::var("key_index"), Expr::var("value_index")),
278                            ),
279                            Expr::mul(
280                                Expr::cast(
281                                    DataType::F32,
282                                    Expr::load(key, qk_index(Expr::var("key_index"))),
283                                ),
284                                Expr::var("key_scale"),
285                            ),
286                        ),
287                    ),
288                )],
289            ),
290            Node::let_bind(
291                "delta",
292                Expr::mul(
293                    Expr::sub(
294                        Expr::cast(
295                            DataType::F32,
296                            Expr::load(value, value_index(Expr::var("value_index"))),
297                        ),
298                        Expr::var("memory"),
299                    ),
300                    Expr::var("beta"),
301                ),
302            ),
303            Node::loop_for(
304                "key_index",
305                Expr::u32(0),
306                Expr::u32(key_dim),
307                vec![Node::Store {
308                    buffer: state_output.into(),
309                    index: state_index(Expr::var("key_index"), Expr::var("value_index")),
310                    value: Expr::add(
311                        Expr::load(
312                            state_output,
313                            state_index(Expr::var("key_index"), Expr::var("value_index")),
314                        ),
315                        Expr::mul(
316                            Expr::mul(
317                                Expr::cast(
318                                    DataType::F32,
319                                    Expr::load(key, qk_index(Expr::var("key_index"))),
320                                ),
321                                Expr::var("key_scale"),
322                            ),
323                            Expr::var("delta"),
324                        ),
325                    ),
326                }],
327            ),
328            Node::let_bind("attention_output", Expr::f32(0.0)),
329            Node::loop_for(
330                "key_index",
331                Expr::u32(0),
332                Expr::u32(key_dim),
333                vec![Node::assign(
334                    "attention_output",
335                    Expr::add(
336                        Expr::var("attention_output"),
337                        Expr::mul(
338                            Expr::load(
339                                state_output,
340                                state_index(Expr::var("key_index"), Expr::var("value_index")),
341                            ),
342                            Expr::mul(
343                                Expr::cast(
344                                    DataType::F32,
345                                    Expr::load(query, qk_index(Expr::var("key_index"))),
346                                ),
347                                Expr::var("query_scale"),
348                            ),
349                        ),
350                    ),
351                )],
352            ),
353            Node::Store {
354                buffer: output.into(),
355                index: output_index,
356                value: Expr::cast(dtype.clone(), Expr::var("attention_output")),
357            },
358        ],
359    );
360    let mut token_body = norm_sums;
361    token_body.extend([
362        Node::let_bind(
363            "decay",
364            Expr::UnOp {
365                op: UnOp::Exp,
366                operand: Box::new(Expr::cast(
367                    DataType::F32,
368                    Expr::load(decay_log, scalar_index.clone()),
369                )),
370            },
371        ),
372        Node::let_bind(
373            "beta_logit",
374            Expr::cast(DataType::F32, Expr::load(beta_logits, scalar_index)),
375        ),
376        Node::let_bind(
377            "beta",
378            Expr::div(
379                Expr::f32(1.0),
380                Expr::add(
381                    Expr::f32(1.0),
382                    Expr::UnOp {
383                        op: UnOp::Exp,
384                        operand: Box::new(Expr::UnOp {
385                            op: UnOp::Negate,
386                            operand: Box::new(Expr::var("beta_logit")),
387                        }),
388                    },
389                ),
390            ),
391        ),
392        decay_state,
393        value_update,
394    ]);
395    let token_schedule = Node::loop_for("token", Expr::u32(0), Expr::u32(sequence), token_body);
396    let body = vec![
397        Node::let_bind("head_index", Expr::InvocationId { axis: 0 }),
398        Node::if_then(
399            Expr::lt(Expr::var("head_index"), Expr::u32(counts.head)),
400            vec![
401                Node::let_bind(
402                    "batch_index",
403                    Expr::div(Expr::var("head_index"), Expr::u32(value_heads)),
404                ),
405                Node::let_bind(
406                    "value_head",
407                    Expr::sub(
408                        Expr::var("head_index"),
409                        Expr::mul(Expr::var("batch_index"), Expr::u32(value_heads)),
410                    ),
411                ),
412                Node::let_bind(
413                    "key_head",
414                    Expr::div(Expr::var("value_head"), Expr::u32(counts.group)),
415                ),
416                init_state,
417                token_schedule,
418            ],
419        ),
420    ];
421
422    Ok(Program::wrapped(
423        gated_delta_layout::gated_delta_buffers(spec, &counts),
424        [64, 1, 1],
425        vec![wrap_anonymous(OP_ID, body)],
426    ))
427}
vyre_libs/nn/attention/gated_delta.rs

vyre_libs/nn/attention/
gated_delta.rs
