Expand description
Sparse Attention Masks
This module implements prompt-aware attention head masking to skip computation for heads that don’t contribute to the output.
§Key Insight
Not all attention heads are equally important for every prompt. Portrait prompts activate face-focused heads while landscape prompts activate background/composition heads. By predicting which heads matter, we can skip 50-70% of attention computation.
§Architecture
┌─────────────────────────────────────────────────────────────────┐
│ Sparse Attention │
├─────────────────────────────────────────────────────────────────┤
│ │
│ Standard Attention (32 heads × 64 layers = 2048 computations) │
│ ════════════════════════════════════════════════════════════ │
│ [████████████████████████████████] 100% compute │
│ │
│ Sparse Attention (prompt-aware masking) │
│ ════════════════════════════════════════════════════════════ │
│ "Portrait of a woman" │
│ [████████░░░░░░░░████░░░░░░░░░░░░] 35% compute │
│ ↑ face ↑ skip ↑ style │
│ │
│ "Mountain landscape at sunset" │
│ [░░░░░░░░████████████████████░░░░] 45% compute │
│ ↑ skip ↑ background/lighting ↑ skip │
└─────────────────────────────────────────────────────────────────┘Modules§
- prelude
- Prelude for common imports
Structs§
- Attention
Mask - A mask indicating which attention heads to compute
- Category
Mapping - Mapping from head indices to categories
- Head
Analysis - Complete analysis of all attention heads
- Head
Importance - Importance score for a single attention head
- Importance
Analyzer - Analyzer for computing head importance
- Importance
Stats - Statistics about head importance distribution
- Kernel
Config - Configuration for sparse attention kernel
- Kernel
Stats - Kernel statistics
- Mask
Builder - Builder for attention masks
- Mask
Predictor - Predicts attention masks from prompt information
- Prediction
- Prediction result with confidence
- Predictor
Config - Configuration for mask prediction
- Sparse
Kernel - Manager for sparse attention kernel execution.
Enums§
- Head
Category - Categories of attention heads based on their learned function
- Mask
Pattern - Pattern for mask generation
- Prompt
Category - Categories of prompts that influence head activation
- Sparse
Error - Errors that can occur during sparse attention operations
Constants§
- DEFAULT_
SPARSITY - Default sparsity target (fraction of heads to skip)
- MAX_
QUALITY_ LOSS - Maximum quality degradation allowed
- MIN_
ACTIVE_ HEADS - Minimum heads to keep active per layer
Type Aliases§
- Result
- Result type for sparse attention operations