ferrotensor_graph/flash_attention.rs
1//! Flash Attention - Memory-Efficient Attention Mechanism
2//!
3//! Implements FlashAttention-2 algorithm for O(N) memory instead of O(N²).
4//! Uses tiled matrix multiplication with online softmax normalization.
5
6use ferrotensor_backend::{AttentionOptions, ComputeEngine, UnifiedTensor, WgpuContext};
7use std::sync::Arc;
8
9/// Flash Attention layer with tiled computation and online softmax.
10///
11/// Memory efficient: O(N) instead of O(N²) for standard attention.
12/// Speed: 2-4x faster than naive implementation.
13///
14/// # Example
15/// ```no_run
16/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
17/// # use ferrotensor_graph::flash_attention::{FlashAttention};
18/// let ctx = pollster::block_on(WgpuContext::new());
19/// let engine = ComputeEngine::new(&ctx);
20/// // [batch, seq_len, hidden]
21/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
22/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
23/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
24/// let flash_attn = FlashAttention::new(768, 12, true);
25/// let output = flash_attn.forward(&q, &k, &v, &ctx, &engine);
26/// ```
27pub struct FlashAttention {
28 dim: usize,
29 num_heads: usize,
30 head_dim: usize,
31 causal: bool,
32 #[allow(dead_code)]
33 scale: f32,
34
35 // Optional: Use optimized multi-head kernel
36 #[allow(dead_code)]
37 use_multi_head_kernel: bool,
38}
39
40/// KV-Cache for efficient autoregressive generation
41pub struct KVCache {
42 /// Cached keys for each layer
43 pub k_cache: Vec<Arc<UnifiedTensor>>,
44 /// Cached values for each layer
45 pub v_cache: Vec<Arc<UnifiedTensor>>,
46 /// Maximum sequence length
47 pub max_seq_len: usize,
48 /// Current position in cache
49 pub current_pos: usize,
50}
51
52impl FlashAttention {
53 /// Create a new FlashAttention layer.
54 ///
55 /// # Arguments
56 /// * `dim` - Total dimension (will be split across heads)
57 /// * `num_heads` - Number of attention heads
58 /// * `causal` - Whether to apply causal masking (for autoregressive models)
59 pub fn new(dim: usize, num_heads: usize, causal: bool) -> Self {
60 let head_dim = dim / num_heads;
61 let scale = 1.0 / (head_dim as f32).sqrt();
62 Self {
63 dim,
64 num_heads,
65 head_dim,
66 causal,
67 scale,
68 use_multi_head_kernel: num_heads > 1,
69 }
70 }
71
72 /// `[seq, dim]`, `[batch, seq, dim]` or `[batch, seq, heads, head_dim]`
73 /// viewed as `[batch, seq, num_heads, head_dim]`.
74 fn heads_view(&self, t: &UnifiedTensor) -> UnifiedTensor {
75 let (batch, seq) = match t.shape.len() {
76 2 => (1, t.shape[0]),
77 3 | 4 => (t.shape[0], t.shape[1]),
78 n => panic!("FlashAttention expects a 2-D, 3-D or 4-D tensor, got {n}-D"),
79 };
80 assert_eq!(
81 t.shape.iter().product::<usize>(),
82 batch * seq * self.dim,
83 "tensor {:?} does not have dim {} per token",
84 t.shape,
85 self.dim
86 );
87 t.reshaped(&[batch, seq, self.num_heads, self.head_dim])
88 }
89
90 /// Forward pass using Flash Attention algorithm.
91 ///
92 /// # Arguments
93 /// * `q` - Query tensor `[batch, seq_q, dim]`, `[seq_q, dim]` or
94 /// `[batch, seq_q, num_heads, head_dim]`
95 /// * `k` - Key tensor with the same layout and `seq_k` tokens
96 /// * `v` - Value tensor with the same layout as `k`
97 ///
98 /// # Returns
99 /// Output tensor with the same shape as `q`.
100 pub fn forward(
101 &self,
102 q: &UnifiedTensor,
103 k: &UnifiedTensor,
104 v: &UnifiedTensor,
105 ctx: &WgpuContext,
106 engine: &ComputeEngine,
107 ) -> UnifiedTensor {
108 let output = UnifiedTensor::empty(ctx, &q.shape);
109 engine.flash_attention_v2(
110 ctx,
111 &self.heads_view(q),
112 &self.heads_view(k),
113 &self.heads_view(v),
114 &self.heads_view(&output),
115 self.causal,
116 0.0,
117 0,
118 None,
119 );
120 output
121 }
122
123 /// Same as [`forward`](Self::forward), which already splits heads.
124 ///
125 /// Kept for compatibility; it will be removed in 0.3.
126 #[deprecated(since = "0.2.1", note = "use `forward`, which already splits heads")]
127 pub fn forward_multi_head(
128 &self,
129 q: &UnifiedTensor,
130 k: &UnifiedTensor,
131 v: &UnifiedTensor,
132 ctx: &WgpuContext,
133 engine: &ComputeEngine,
134 ) -> UnifiedTensor {
135 self.forward(q, k, v, ctx, engine)
136 }
137
138 /// Self-attention with `hidden` used directly as Q, K and V.
139 ///
140 /// There are no Q, K, V or output projections: project `hidden` yourself
141 /// and call [`forward`](Self::forward). Kept for compatibility; it will be
142 /// removed in 0.3.
143 #[deprecated(
144 since = "0.2.1",
145 note = "applies no projections; project the input and call `forward`"
146 )]
147 pub fn multi_head_forward(
148 &self,
149 hidden: &UnifiedTensor,
150 ctx: &WgpuContext,
151 engine: &ComputeEngine,
152 ) -> UnifiedTensor {
153 self.forward(hidden, hidden, hidden, ctx, engine)
154 }
155
156 /// Memory for one head's attention scores: the full `seq x seq` matrix
157 /// that standard attention stores versus the `seq x head_dim` accumulator
158 /// Flash Attention keeps instead (f32, per head, scores only).
159 pub fn memory_savings(&self, seq_len: usize) -> String {
160 let standard_memory = seq_len * seq_len * 4;
161 let flash_memory = seq_len * self.head_dim * 4;
162 let savings = standard_memory as f32 / flash_memory as f32;
163 format!(
164 "{:.1}x less memory ({} vs {} per head for seq={})",
165 savings,
166 format_bytes(flash_memory),
167 format_bytes(standard_memory),
168 seq_len
169 )
170 }
171
172 /// Backward pass for Flash Attention.
173 ///
174 /// Computes gradients dQ, dK, dV using recomputation to avoid
175 /// storing the full attention matrix.
176 ///
177 /// # Arguments
178 /// * `q` - Query tensor from forward pass
179 /// * `k` - Key tensor from forward pass
180 /// * `v` - Value tensor from forward pass
181 /// * `output` - Output from forward pass
182 /// * `grad_output` - Gradient of the output
183 ///
184 /// # Returns
185 /// FlashAttentionGrads containing dQ, dK, dV
186 #[allow(clippy::too_many_arguments)]
187 pub fn backward(
188 &self,
189 q: &UnifiedTensor,
190 k: &UnifiedTensor,
191 v: &UnifiedTensor,
192 output: &UnifiedTensor,
193 grad_output: &UnifiedTensor,
194 ctx: &WgpuContext,
195 engine: &ComputeEngine,
196 ) -> FlashAttentionGrads {
197 let grad_q = UnifiedTensor::empty(ctx, &q.shape);
198 let grad_k = UnifiedTensor::empty(ctx, &k.shape);
199 let grad_v = UnifiedTensor::empty(ctx, &v.shape);
200
201 engine.flash_attention_v2_backward(
202 ctx,
203 &self.heads_view(q),
204 &self.heads_view(k),
205 &self.heads_view(v),
206 &self.heads_view(output),
207 &self.heads_view(grad_output),
208 &self.heads_view(&grad_q),
209 &self.heads_view(&grad_k),
210 &self.heads_view(&grad_v),
211 self.causal,
212 );
213
214 FlashAttentionGrads {
215 grad_q,
216 grad_k,
217 grad_v,
218 }
219 }
220}
221
222/// Gradients from Flash Attention backward pass.
223#[derive(Debug)]
224pub struct FlashAttentionGrads {
225 /// Gradient with respect to queries
226 pub grad_q: UnifiedTensor,
227 /// Gradient with respect to keys
228 pub grad_k: UnifiedTensor,
229 /// Gradient with respect to values
230 pub grad_v: UnifiedTensor,
231}
232
233impl std::fmt::Debug for FlashAttention {
234 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
235 f.debug_struct("FlashAttention")
236 .field("dim", &self.dim)
237 .field("num_heads", &self.num_heads)
238 .field("head_dim", &self.head_dim)
239 .field("causal", &self.causal)
240 .finish()
241 }
242}
243
244/// Format a byte count with a binary unit (B, KiB, MiB, GiB).
245fn format_bytes(bytes: usize) -> String {
246 const UNITS: [&str; 4] = ["B", "KiB", "MiB", "GiB"];
247 let mut value = bytes as f64;
248 let mut unit = 0;
249 while value >= 1024.0 && unit < UNITS.len() - 1 {
250 value /= 1024.0;
251 unit += 1;
252 }
253 if unit == 0 {
254 format!("{bytes} B")
255 } else {
256 format!("{value:.1} {}", UNITS[unit])
257 }
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263
264 #[test]
265 fn test_memory_savings() {
266 let attn = FlashAttention::new(512, 8, true);
267 let savings = attn.memory_savings(2048);
268 assert!(savings.contains("x less memory"));
269 // head_dim 64: 512*64*4 = 128 KiB vs 512*512*4 = 1 MiB.
270 let attn = FlashAttention::new(512, 8, true);
271 assert_eq!(
272 attn.memory_savings(512),
273 "8.0x less memory (128.0 KiB vs 1.0 MiB per head for seq=512)"
274 );
275 }
276
277 #[test]
278 fn test_format_bytes() {
279 assert_eq!(format_bytes(0), "0 B");
280 assert_eq!(format_bytes(1023), "1023 B");
281 assert_eq!(format_bytes(1536), "1.5 KiB");
282 assert_eq!(format_bytes(3 * 1024 * 1024 * 1024), "3.0 GiB");
283 }
284
285 #[test]
286 fn test_creation() {
287 let attn = FlashAttention::new(768, 12, true);
288 assert_eq!(attn.dim, 768);
289 assert_eq!(attn.num_heads, 12);
290 assert_eq!(attn.head_dim, 64);
291 assert!(attn.causal);
292 }
293}
294
295// ============================================================================
296// GROUPED QUERY ATTENTION (GQA) - Used by Llama 3, Mistral, etc.
297// ============================================================================
298
299/// Grouped Query Attention (GQA) - Efficient attention with fewer KV heads
300///
301/// GQA uses fewer key-value heads than query heads, reducing memory and
302/// computation while maintaining quality. This is the attention mechanism
303/// used by Llama 3, Mistral, and many modern LLMs.
304///
305/// # Memory Savings
306/// - Standard MHA: `num_heads` K and V projections
307/// - GQA: `num_kv_heads` K and V projections (typically 2-8x fewer)
308///
309/// # Example
310/// ```no_run
311/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
312/// # use ferrotensor_graph::flash_attention::{GroupedQueryAttention};
313/// let ctx = pollster::block_on(WgpuContext::new());
314/// let engine = ComputeEngine::new(&ctx);
315/// // [batch, seq_len, hidden]
316/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 4096]);
317/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 1024]);
318/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 1024]);
319/// // Llama 3 8B: 32 query heads share 8 KV heads (no dropout, no mask)
320/// let gqa = GroupedQueryAttention::new(4096, 32, 8, true, 0.0);
321/// let output = gqa.forward(&q, &k, &v, None, &ctx, &engine);
322/// ```
323pub struct GroupedQueryAttention {
324 /// Total hidden dimension
325 pub dim: usize,
326 /// Number of query heads
327 pub num_heads: usize,
328 /// Number of key-value heads (fewer than query heads)
329 pub num_kv_heads: usize,
330 /// Dimension per head
331 pub head_dim: usize,
332 /// Number of query heads per KV head
333 pub num_groups: usize,
334 /// Whether to use causal masking
335 pub causal: bool,
336 /// Scaling factor for attention scores
337 #[allow(dead_code)]
338 scale: f32,
339 /// Dropout probability (0.0 = no dropout)
340 pub dropout_prob: f32,
341}
342
343impl GroupedQueryAttention {
344 /// Create a new GQA layer.
345 ///
346 /// # Arguments
347 /// * `dim` - Total hidden dimension
348 /// * `num_heads` - Number of query heads
349 /// * `num_kv_heads` - Number of key-value heads (must divide num_heads evenly)
350 /// * `causal` - Whether to apply causal masking
351 /// * `dropout_prob` - Dropout probability
352 ///
353 /// # Panics
354 /// Panics if `num_heads` is not divisible by `num_kv_heads`
355 pub fn new(
356 dim: usize,
357 num_heads: usize,
358 num_kv_heads: usize,
359 causal: bool,
360 dropout_prob: f32,
361 ) -> Self {
362 assert!(
363 num_heads.is_multiple_of(num_kv_heads),
364 "num_heads ({}) must be divisible by num_kv_heads ({})",
365 num_heads,
366 num_kv_heads
367 );
368
369 let head_dim = dim / num_heads;
370 let num_groups = num_heads / num_kv_heads;
371 let scale = 1.0 / (head_dim as f32).sqrt();
372
373 Self {
374 dim,
375 num_heads,
376 num_kv_heads,
377 head_dim,
378 num_groups,
379 causal,
380 scale,
381 dropout_prob,
382 }
383 }
384
385 /// Forward pass with grouped query attention.
386 ///
387 /// # Arguments
388 /// * `q` - Query tensor [batch, seq_q, num_heads, head_dim]
389 /// * `k` - Key tensor [batch, seq_k, num_kv_heads, head_dim]
390 /// * `v` - Value tensor [batch, seq_k, num_kv_heads, head_dim]
391 ///
392 /// # Returns
393 /// Output tensor [batch, seq_q, num_heads, head_dim]
394 pub fn forward(
395 &self,
396 q: &UnifiedTensor,
397 k: &UnifiedTensor,
398 v: &UnifiedTensor,
399 mask: Option<&UnifiedTensor>,
400 ctx: &WgpuContext,
401 engine: &ComputeEngine,
402 ) -> UnifiedTensor {
403 self.forward_at(q, k, v, mask, 0, ctx, engine)
404 }
405
406 /// Like [`forward`](Self::forward), with the first query at absolute
407 /// position `pos`: query `i` attends to keys `0..=pos + i` when causal.
408 /// Used to decode against a KV cache, where `k`/`v` hold the cache.
409 #[allow(clippy::too_many_arguments)]
410 pub fn forward_at(
411 &self,
412 q: &UnifiedTensor,
413 k: &UnifiedTensor,
414 v: &UnifiedTensor,
415 mask: Option<&UnifiedTensor>,
416 pos: usize,
417 ctx: &WgpuContext,
418 engine: &ComputeEngine,
419 ) -> UnifiedTensor {
420 // Validation: Expect 4D tensors [batch, seq, heads, dim]
421 if q.shape.len() != 4 || k.shape.len() != 4 || v.shape.len() != 4 {
422 panic!("GroupedQueryAttention expects 4D tensors (batch, seq, heads, dim)");
423 }
424
425 let batch_size = q.shape[0];
426 let seq_q = q.shape[1];
427
428 let output = UnifiedTensor::empty(ctx, &[batch_size, seq_q, self.num_heads, self.head_dim]);
429
430 // Use time-based seed for stochastic dropout
431 let seed = std::time::SystemTime::now()
432 .duration_since(std::time::UNIX_EPOCH)
433 .unwrap_or_default()
434 .as_nanos() as u64;
435
436 let opts = AttentionOptions {
437 causal: self.causal,
438 causal_offset: pos as u32,
439 dropout_prob: self.dropout_prob,
440 seed,
441 ..AttentionOptions::default()
442 };
443 engine.flash_attention_v2_with(ctx, q, k, v, &output, &opts, mask);
444
445 output
446 }
447
448 /// Get memory savings compared to standard MHA
449 pub fn memory_savings(&self) -> String {
450 let mha_kv_params = 2 * self.num_heads * self.head_dim;
451 let gqa_kv_params = 2 * self.num_kv_heads * self.head_dim;
452 let savings = mha_kv_params as f32 / gqa_kv_params as f32;
453
454 format!(
455 "{:.1}x KV memory savings ({} vs {} KV heads)",
456 savings, self.num_kv_heads, self.num_heads
457 )
458 }
459}
460
461// ============================================================================
462// MULTI-QUERY ATTENTION (MQA) - Single KV head for maximum efficiency
463// ============================================================================
464
465/// Multi-Query Attention (MQA) - Extreme efficiency with single KV head
466///
467/// MQA uses a single key-value head shared across all query heads.
468/// This provides maximum memory efficiency at some quality cost.
469///
470/// # Use Cases
471/// - Inference-optimized models
472/// - Memory-constrained deployments
473/// - High-throughput serving
474///
475/// # Example
476/// ```no_run
477/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
478/// # use ferrotensor_graph::flash_attention::{MultiQueryAttention};
479/// let ctx = pollster::block_on(WgpuContext::new());
480/// let engine = ComputeEngine::new(&ctx);
481/// // [batch, seq_len, hidden]
482/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
483/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 64]);
484/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 64]);
485/// // 12 query heads share one KV head
486/// let mqa = MultiQueryAttention::new(768, 12, true, 0.0);
487/// let output = mqa.forward(&q, &k, &v, None, &ctx, &engine);
488/// ```
489pub struct MultiQueryAttention {
490 /// Total hidden dimension
491 pub dim: usize,
492 /// Number of query heads
493 pub num_heads: usize,
494 /// Dimension per head
495 pub head_dim: usize,
496 /// Whether to use causal masking
497 pub causal: bool,
498 /// Internal GQA with num_kv_heads = 1
499 gqa: GroupedQueryAttention,
500}
501
502impl MultiQueryAttention {
503 /// Create a new MQA layer.
504 ///
505 /// # Arguments
506 /// * `dim` - Total hidden dimension
507 /// * `num_heads` - Number of query heads
508 /// * `causal` - Whether to apply causal masking
509 pub fn new(dim: usize, num_heads: usize, causal: bool, dropout_prob: f32) -> Self {
510 let head_dim = dim / num_heads;
511 let gqa = GroupedQueryAttention::new(dim, num_heads, 1, causal, dropout_prob);
512
513 Self {
514 dim,
515 num_heads,
516 head_dim,
517 causal,
518 gqa,
519 }
520 }
521
522 /// Forward pass with multi-query attention.
523 ///
524 /// # Arguments
525 /// * `q` - Query tensor [batch, seq_q, num_heads, head_dim]
526 /// * `k` - Key tensor [batch, seq_k, 1, head_dim]
527 /// * `v` - Value tensor [batch, seq_k, 1, head_dim]
528 ///
529 /// # Returns
530 /// Output tensor [batch, seq_q, num_heads, head_dim]
531 pub fn forward(
532 &self,
533 q: &UnifiedTensor,
534 k: &UnifiedTensor,
535 v: &UnifiedTensor,
536 mask: Option<&UnifiedTensor>,
537 ctx: &WgpuContext,
538 engine: &ComputeEngine,
539 ) -> UnifiedTensor {
540 // Delegate to GQA with num_kv_heads = 1
541 self.gqa.forward(q, k, v, mask, ctx, engine)
542 }
543
544 /// Get memory savings compared to MHA
545 pub fn memory_savings(&self) -> String {
546 format!(
547 "{}x KV memory savings (1 vs {} KV heads)",
548 self.num_heads, self.num_heads
549 )
550 }
551}
552
553// ============================================================================
554// SLIDING WINDOW ATTENTION - For long sequences
555// ============================================================================
556
557/// Sliding Window Attention for efficient long-context processing
558///
559/// Limits attention to a local window around each position, enabling
560/// linear memory scaling for very long sequences.
561///
562/// # Use Cases
563/// - Long document processing (100K+ tokens)
564/// - Streaming input processing
565/// - Memory-efficient inference
566///
567/// # Example
568/// ```no_run
569/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
570/// # use ferrotensor_graph::flash_attention::{SlidingWindowAttention};
571/// let ctx = pollster::block_on(WgpuContext::new());
572/// let engine = ComputeEngine::new(&ctx);
573/// // [batch, seq_len, hidden]
574/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
575/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
576/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
577/// let swa = SlidingWindowAttention::new(768, 12, 4096, true);
578/// let output = swa.forward(&q, &k, &v, &ctx, &engine);
579/// ```
580pub struct SlidingWindowAttention {
581 /// Total hidden dimension
582 pub dim: usize,
583 /// Number of attention heads
584 pub num_heads: usize,
585 /// Dimension per head
586 pub head_dim: usize,
587 /// Window size for local attention
588 pub window_size: usize,
589 /// Whether to use causal masking
590 pub causal: bool,
591 #[allow(dead_code)]
592 scale: f32,
593}
594
595impl SlidingWindowAttention {
596 /// Create a new sliding window attention layer.
597 ///
598 /// # Arguments
599 /// * `dim` - Total hidden dimension
600 /// * `num_heads` - Number of attention heads
601 /// * `window_size` - Maximum distance for attention (tokens only attend within this window)
602 /// * `causal` - Whether to apply causal masking
603 pub fn new(dim: usize, num_heads: usize, window_size: usize, causal: bool) -> Self {
604 let head_dim = dim / num_heads;
605 let scale = 1.0 / (head_dim as f32).sqrt();
606
607 Self {
608 dim,
609 num_heads,
610 head_dim,
611 window_size,
612 causal,
613 scale,
614 }
615 }
616
617 /// Forward pass with sliding window attention.
618 ///
619 /// Each position can only attend to positions within `window_size` distance.
620 ///
621 /// # Arguments
622 /// * `q` - Query tensor [batch, seq_q, dim]
623 /// * `k` - Key tensor [batch, seq_k, dim]
624 /// * `v` - Value tensor [batch, seq_k, dim]
625 ///
626 /// # Returns
627 /// Output tensor [batch, seq_q, dim]
628 pub fn forward(
629 &self,
630 q: &UnifiedTensor,
631 k: &UnifiedTensor,
632 v: &UnifiedTensor,
633 ctx: &WgpuContext,
634 engine: &ComputeEngine,
635 ) -> UnifiedTensor {
636 let view = |t: &UnifiedTensor| {
637 t.reshaped(&[t.shape[0], t.shape[1], self.num_heads, self.head_dim])
638 };
639 let output = UnifiedTensor::empty(ctx, &q.shape);
640 let opts = AttentionOptions {
641 causal: self.causal,
642 window: self.window_size as u32,
643 ..AttentionOptions::default()
644 };
645 engine.flash_attention_v2_with(
646 ctx,
647 &view(q),
648 &view(k),
649 &view(v),
650 &view(&output),
651 &opts,
652 None,
653 );
654 output
655 }
656
657 /// Get memory savings compared to full attention
658 pub fn memory_savings(&self, seq_len: usize) -> String {
659 let full_attention = seq_len * seq_len;
660 let window_attention = seq_len * self.window_size.min(seq_len);
661 let savings = full_attention as f32 / window_attention as f32;
662
663 format!(
664 "{:.1}x attention savings (window={}, seq={})",
665 savings, self.window_size, seq_len
666 )
667 }
668}
669
670// ============================================================================
671// ATTENTION MASK - Custom masking patterns
672// ============================================================================
673
674/// Attention mask for custom attention patterns
675#[derive(Debug, Clone)]
676pub enum AttentionMask {
677 /// No mask (full attention)
678 None,
679 /// Causal mask (lower triangular)
680 Causal,
681 /// Padding mask to ignore padded tokens
682 Padding(Vec<usize>), // lengths of non-padded sequences
683 /// Sliding window mask
684 SlidingWindow(usize),
685 /// Custom mask tensor
686 Custom(Arc<UnifiedTensor>),
687}
688
689impl AttentionMask {
690 /// Create a causal mask
691 pub fn causal() -> Self {
692 Self::Causal
693 }
694
695 /// Create a padding mask from sequence lengths
696 pub fn from_lengths(lengths: Vec<usize>) -> Self {
697 Self::Padding(lengths)
698 }
699
700 /// Create a sliding window mask
701 pub fn sliding_window(window_size: usize) -> Self {
702 Self::SlidingWindow(window_size)
703 }
704}
705
706impl std::fmt::Debug for GroupedQueryAttention {
707 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
708 f.debug_struct("GroupedQueryAttention")
709 .field("dim", &self.dim)
710 .field("num_heads", &self.num_heads)
711 .field("num_kv_heads", &self.num_kv_heads)
712 .field("num_groups", &self.num_groups)
713 .field("head_dim", &self.head_dim)
714 .field("causal", &self.causal)
715 .field("dropout", &self.dropout_prob)
716 .finish()
717 }
718}
719
720impl std::fmt::Debug for MultiQueryAttention {
721 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
722 f.debug_struct("MultiQueryAttention")
723 .field("dim", &self.dim)
724 .field("num_heads", &self.num_heads)
725 .field("head_dim", &self.head_dim)
726 .field("causal", &self.causal)
727 .finish()
728 }
729}
730
731impl std::fmt::Debug for SlidingWindowAttention {
732 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
733 f.debug_struct("SlidingWindowAttention")
734 .field("dim", &self.dim)
735 .field("num_heads", &self.num_heads)
736 .field("head_dim", &self.head_dim)
737 .field("window_size", &self.window_size)
738 .field("causal", &self.causal)
739 .finish()
740 }
741}