ferrotensor_backend/engine/
attention.rs

1//! Rotary embeddings and fused (flash) attention.
2
3use super::{dummy_storage, groups64, uniform, ComputeEngine};
4use crate::{UnifiedTensor, WgpuContext};
5use wgpu::Buffer;
6
7impl ComputeEngine {
8    /// Rotate-half RoPE (theta 10000) applied in place.
9    ///
10    /// `tensor` is `[batch, seq, ...]`, `[seq, hidden]` or a single token
11    /// `[hidden]`; every head vector of token `s` gets position `pos + s`.
12    pub fn rope(&self, ctx: &WgpuContext, tensor: &UnifiedTensor, head_dim: usize, pos: usize) {
13        self.rope_theta(ctx, tensor, head_dim, pos, 10000.0);
14    }
15
16    /// [`rope`](Self::rope) with base frequency `theta` (the `rope_theta` of
17    /// a Hugging Face config, e.g. 100000 for SmolLM2).
18    pub fn rope_theta(
19        &self,
20        ctx: &WgpuContext,
21        tensor: &UnifiedTensor,
22        head_dim: usize,
23        pos: usize,
24        theta: f32,
25    ) {
26        let (seq_len, tokens) = match tensor.shape.len() {
27            0 | 1 => (1, 1),
28            2 => (tensor.shape[0], tensor.shape[0]),
29            _ => (tensor.shape[1], tensor.shape[0] * tensor.shape[1]),
30        };
31        let total = tensor.shape.iter().product::<usize>();
32        let heads = total / (tokens * head_dim).max(1);
33        let params = uniform(
34            ctx,
35            bytemuck::cast_slice(&[
36                head_dim as u32,
37                seq_len as u32,
38                pos as u32,
39                heads as u32,
40                theta.to_bits(),
41                0,
42                0,
43                0,
44            ]),
45        );
46        let pairs = tensor.shape.iter().product::<usize>() / 2;
47        self.run(
48            ctx,
49            &self.rope,
50            &[(0, tensor.buffer()), (1, &params)],
51            groups64(pairs),
52        );
53    }
54
55    /// Single-head attention: `q`, `k`, `v`, `out` are `[seq_len, head_dim]`.
56    ///
57    /// Runs the v2 kernel on `[1, seq_len, 1, head_dim]` views of the inputs.
58    pub fn flash_attention(
59        &self,
60        ctx: &WgpuContext,
61        q: &UnifiedTensor,
62        k: &UnifiedTensor,
63        v: &UnifiedTensor,
64        out: &UnifiedTensor,
65        causal: bool,
66    ) {
67        let [q4, k4, v4, o4] = [q, k, v, out].map(single_head_view);
68        self.flash_attention_v2(ctx, &q4, &k4, &v4, &o4, causal, 0.0, 0, None);
69    }
70
71    /// Fused attention forward pass with grouped-query attention, causal
72    /// masking, an optional additive mask and attention dropout.
73    ///
74    /// * `q`, `out` - `[batch, seq_q, num_heads, head_dim]`
75    /// * `k`, `v` - `[batch, seq_k, num_kv_heads, head_dim]`; `num_heads`
76    ///   must be a multiple of `num_kv_heads`
77    /// * `causal` - key `j` is visible to query `i` only when `j <= i`
78    /// * `dropout_prob`, `seed` - attention dropout; pass `0.0` to disable
79    /// * `mask` - additive mask `[batch, heads, seq_q, seq_k]`, where size-1
80    ///   dimensions broadcast
81    ///
82    /// `head_dim` must be at most 128. See
83    /// [`flash_attention_v2_with`](Self::flash_attention_v2_with) for
84    /// KV-cache offsets and sliding windows.
85    #[allow(clippy::too_many_arguments)]
86    pub fn flash_attention_v2(
87        &self,
88        ctx: &WgpuContext,
89        q: &UnifiedTensor,
90        k: &UnifiedTensor,
91        v: &UnifiedTensor,
92        out: &UnifiedTensor,
93        causal: bool,
94        dropout_prob: f32,
95        seed: u64,
96        mask: Option<&UnifiedTensor>,
97    ) {
98        let opts = AttentionOptions {
99            causal,
100            dropout_prob,
101            seed,
102            ..AttentionOptions::default()
103        };
104        self.flash_attention_v2_with(ctx, q, k, v, out, &opts, mask);
105    }
106
107    /// [`flash_attention_v2`](Self::flash_attention_v2) with every option.
108    #[allow(clippy::too_many_arguments)]
109    pub fn flash_attention_v2_with(
110        &self,
111        ctx: &WgpuContext,
112        q: &UnifiedTensor,
113        k: &UnifiedTensor,
114        v: &UnifiedTensor,
115        out: &UnifiedTensor,
116        opts: &AttentionOptions,
117        mask: Option<&UnifiedTensor>,
118    ) {
119        let dims = AttentionDims::new(q, k);
120        let params = dims.params(ctx, opts, mask);
121        let dummy;
122        let mask_buf = match mask {
123            Some(m) => m.buffer(),
124            None => {
125                dummy = dummy_storage(ctx);
126                &dummy
127            }
128        };
129        self.run(
130            ctx,
131            &self.flash_attn_v2,
132            &[
133                (0, q.buffer()),
134                (1, k.buffer()),
135                (2, v.buffer()),
136                (3, out.buffer()),
137                (4, &params),
138                (5, mask_buf),
139            ],
140            (1, dims.seq_q.div_ceil(128), dims.batch * dims.heads),
141        );
142    }
143
144    /// Single-head attention backward pass for `[seq_len, head_dim]` tensors.
145    /// Writes (does not accumulate into) `grad_q`, `grad_k` and `grad_v`.
146    #[allow(clippy::too_many_arguments)]
147    pub fn flash_attention_backward(
148        &self,
149        ctx: &WgpuContext,
150        q: &UnifiedTensor,
151        k: &UnifiedTensor,
152        v: &UnifiedTensor,
153        out: &UnifiedTensor,
154        grad_out: &UnifiedTensor,
155        grad_q: &UnifiedTensor,
156        grad_k: &UnifiedTensor,
157        grad_v: &UnifiedTensor,
158        causal: bool,
159    ) {
160        let views = [q, k, v, out, grad_out, grad_q, grad_k, grad_v].map(single_head_view);
161        let [q4, k4, v4, o4, go4, gq4, gk4, gv4] = &views;
162        self.flash_attention_v2_backward(ctx, q4, k4, v4, o4, go4, gq4, gk4, gv4, causal);
163    }
164
165    /// Backward pass of [`flash_attention_v2`](Self::flash_attention_v2)
166    /// without mask or dropout. See
167    /// [`flash_attention_v2_backward_with`](Self::flash_attention_v2_backward_with).
168    #[allow(clippy::too_many_arguments)]
169    pub fn flash_attention_v2_backward(
170        &self,
171        ctx: &WgpuContext,
172        q: &UnifiedTensor,
173        k: &UnifiedTensor,
174        v: &UnifiedTensor,
175        out: &UnifiedTensor,
176        grad_out: &UnifiedTensor,
177        grad_q: &UnifiedTensor,
178        grad_k: &UnifiedTensor,
179        grad_v: &UnifiedTensor,
180        causal: bool,
181    ) {
182        let opts = AttentionOptions {
183            causal,
184            ..AttentionOptions::default()
185        };
186        self.flash_attention_v2_backward_with(
187            ctx, q, k, v, out, grad_out, grad_q, grad_k, grad_v, &opts, None,
188        );
189    }
190
191    /// Backward pass of [`flash_attention_v2`](Self::flash_attention_v2).
192    ///
193    /// `out` is the forward output for the same inputs; `opts` and `mask`
194    /// must match the forward call (dropout decisions are regenerated from
195    /// `opts.seed`). Writes `grad_q`, `grad_k` and
196    /// `grad_v` (same shapes as `q`, `k`, `v`); no gradient is produced for
197    /// the mask.
198    #[allow(clippy::too_many_arguments)]
199    pub fn flash_attention_v2_backward_with(
200        &self,
201        ctx: &WgpuContext,
202        q: &UnifiedTensor,
203        k: &UnifiedTensor,
204        v: &UnifiedTensor,
205        out: &UnifiedTensor,
206        grad_out: &UnifiedTensor,
207        grad_q: &UnifiedTensor,
208        grad_k: &UnifiedTensor,
209        grad_v: &UnifiedTensor,
210        opts: &AttentionOptions,
211        mask: Option<&UnifiedTensor>,
212    ) {
213        let dims = AttentionDims::new(q, k);
214        let params = dims.params_bwd(ctx, opts, mask);
215        let dummy;
216        let mask_buf = match mask {
217            Some(m) => m.buffer(),
218            None => {
219                dummy = dummy_storage(ctx);
220                &dummy
221            }
222        };
223        // Per (batch, head, query) row: [Delta, LSE].
224        let stats = ctx.device.create_buffer(&wgpu::BufferDescriptor {
225            label: Some("Flash Bwd Stats"),
226            size: (dims.batch * dims.heads * dims.seq_q).max(1) as u64 * 8,
227            usage: wgpu::BufferUsages::STORAGE,
228            mapped_at_creation: false,
229        });
230        let (b, h, hkv) = (dims.batch, dims.heads, dims.kv_heads);
231        self.run(
232            ctx,
233            &self.flash_attn_v2_bw_dq,
234            &[
235                (0, q.buffer()),
236                (1, k.buffer()),
237                (2, v.buffer()),
238                (3, out.buffer()),
239                (4, grad_out.buffer()),
240                (5, grad_q.buffer()),
241                (8, &params),
242                (9, mask_buf),
243                (10, &stats),
244            ],
245            (1, dims.seq_q.div_ceil(128), b * h),
246        );
247        self.run(
248            ctx,
249            &self.flash_attn_v2_bw_dkdv,
250            &[
251                (0, q.buffer()),
252                (1, k.buffer()),
253                (2, v.buffer()),
254                (4, grad_out.buffer()),
255                (6, grad_k.buffer()),
256                (7, grad_v.buffer()),
257                (8, &params),
258                (9, mask_buf),
259                (10, &stats),
260            ],
261            (1, dims.seq_k.div_ceil(128), b * hkv),
262        );
263    }
264}
265
266/// `[seq, dim]` tensor viewed as `[1, seq, 1, dim]`.
267fn single_head_view(t: &UnifiedTensor) -> UnifiedTensor {
268    assert_eq!(t.shape.len(), 2, "expected a [seq_len, head_dim] tensor");
269    t.reshaped(&[1, t.shape[0], 1, t.shape[1]])
270}
271
272/// Shapes of a `[batch, seq, heads, head_dim]` attention call.
273struct AttentionDims {
274    batch: u32,
275    seq_q: u32,
276    seq_k: u32,
277    heads: u32,
278    kv_heads: u32,
279    head_dim: u32,
280}
281
282impl AttentionDims {
283    fn new(q: &UnifiedTensor, k: &UnifiedTensor) -> Self {
284        assert_eq!(q.shape.len(), 4, "q must be [batch, seq, heads, head_dim]");
285        assert_eq!(
286            k.shape.len(),
287            4,
288            "k must be [batch, seq, kv_heads, head_dim]"
289        );
290        let dims = Self {
291            batch: q.shape[0] as u32,
292            seq_q: q.shape[1] as u32,
293            heads: q.shape[2] as u32,
294            head_dim: q.shape[3] as u32,
295            seq_k: k.shape[1] as u32,
296            kv_heads: k.shape[2] as u32,
297        };
298        assert!(
299            dims.head_dim <= 128,
300            "flash attention supports head_dim <= 128, got {}",
301            dims.head_dim
302        );
303        assert!(
304            dims.kv_heads > 0 && dims.heads.is_multiple_of(dims.kv_heads),
305            "num_heads ({}) must be a multiple of num_kv_heads ({})",
306            dims.heads,
307            dims.kv_heads
308        );
309        dims
310    }
311
312    fn scale(&self) -> f32 {
313        1.0 / (self.head_dim as f32).sqrt()
314    }
315
316    /// Uniform for the forward kernel (`FlashV2Params`).
317    fn params(
318        &self,
319        ctx: &WgpuContext,
320        opts: &AttentionOptions,
321        mask: Option<&UnifiedTensor>,
322    ) -> Buffer {
323        let (has_mask, s) = mask_strides(mask);
324        let params: [u32; 20] = [
325            self.batch,
326            self.heads,
327            self.seq_q,
328            self.seq_k,
329            self.head_dim,
330            opts.causal as u32,
331            self.scale().to_bits(),
332            self.kv_heads,
333            opts.dropout_prob.to_bits(),
334            opts.seed as u32,
335            has_mask,
336            opts.causal_offset,
337            s[0],
338            s[1],
339            s[2],
340            s[3],
341            opts.window,
342            0,
343            0,
344            0,
345        ];
346        uniform(ctx, bytemuck::cast_slice(&params))
347    }
348
349    /// Uniform for the backward kernels (`FlashBwdParams`).
350    fn params_bwd(
351        &self,
352        ctx: &WgpuContext,
353        opts: &AttentionOptions,
354        mask: Option<&UnifiedTensor>,
355    ) -> Buffer {
356        let (has_mask, s) = mask_strides(mask);
357        let params: [u32; 20] = [
358            self.batch,
359            self.heads,
360            self.seq_q,
361            self.seq_k,
362            self.head_dim,
363            self.kv_heads,
364            self.scale().to_bits(),
365            opts.dropout_prob.to_bits(),
366            opts.seed as u32,
367            opts.causal as u32,
368            has_mask,
369            opts.causal_offset,
370            s[0],
371            s[1],
372            s[2],
373            s[3],
374            opts.window,
375            0,
376            0,
377            0,
378        ];
379        uniform(ctx, bytemuck::cast_slice(&params))
380    }
381}
382
383/// Options for [`ComputeEngine::flash_attention_v2_with`] and
384/// [`ComputeEngine::flash_attention_v2_backward_with`].
385///
386/// Query `i` sits at absolute position `i + causal_offset`; key `j` is
387/// visible when `j <= i + causal_offset` (if `causal`) and
388/// `|i + causal_offset - j| < window` (if `window > 0`).
389#[derive(Debug, Clone, Copy, Default)]
390pub struct AttentionOptions {
391    /// Hide keys after each query's absolute position.
392    pub causal: bool,
393    /// Absolute position of the first query, e.g. the number of tokens
394    /// already in a KV cache when decoding.
395    pub causal_offset: u32,
396    /// Sliding-window size in tokens; 0 means unlimited.
397    pub window: u32,
398    /// Attention dropout probability; 0.0 disables dropout.
399    pub dropout_prob: f32,
400    /// Seed for the dropout pattern (only the low 32 bits are used).
401    pub seed: u64,
402}
403
404/// Broadcast strides `[b, h, q, k]` for an additive attention mask of shape
405/// `[batch, heads, seq_q, seq_k]`; size-1 dimensions get stride 0.
406fn mask_strides(mask: Option<&UnifiedTensor>) -> (u32, [u32; 4]) {
407    let Some(m) = mask else {
408        return (0, [0; 4]);
409    };
410    let shape = &m.shape;
411    if shape.len() != 4 {
412        return (1, [0; 4]);
413    }
414    let dim_b = shape[0] as u32;
415    let dim_h = shape[1] as u32;
416    let dim_q = shape[2] as u32;
417    let dim_k = shape[3] as u32;
418    let s = |dim: u32, stride: u32| if dim > 1 { stride } else { 0 };
419    (
420        1,
421        [
422            s(dim_b, dim_k * dim_q * dim_h),
423            s(dim_h, dim_k * dim_q),
424            s(dim_q, dim_k),
425            s(dim_k, 1),
426        ],
427    )
428}
ferrotensor_backend/engine/attention.rs

ferrotensor_backend/engine/
attention.rs
