Struct SlidingWindowAttention 

pub struct SlidingWindowAttention {
    pub dim: usize,
    pub num_heads: usize,
    pub head_dim: usize,
    pub window_size: usize,
    pub causal: bool,
    /* private fields */
}

Sliding Window Attention for efficient long-context processing

Limits attention to a local window around each position, enabling linear memory scaling for very long sequences.

Use Cases

• Long document processing (100K+ tokens)
• Streaming input processing
• Memory-efficient inference

Example

let ctx = pollster::block_on(WgpuContext::new());
let engine = ComputeEngine::new(&ctx);
// [batch, seq_len, hidden]
let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
let k = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
let v = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
let swa = SlidingWindowAttention::new(768, 12, 4096, true);
let output = swa.forward(&q, &k, &v, &ctx, &engine);

Fields

dim: usize

Total hidden dimension

num_heads: usize

Number of attention heads

head_dim: usize

Dimension per head

window_size: usize

Window size for local attention

causal: bool

Whether to use causal masking

Implementations

impl SlidingWindowAttention

pub fn new( dim: usize, num_heads: usize, window_size: usize, causal: bool, ) -> Self

Create a new sliding window attention layer.

Arguments

• dim - Total hidden dimension
• num_heads - Number of attention heads
• window_size - Maximum distance for attention (tokens only attend within this window)
• causal - Whether to apply causal masking

pub fn forward( &self, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, ctx: &WgpuContext, engine: &ComputeEngine, ) -> UnifiedTensor

Forward pass with sliding window attention.

Each position can only attend to positions within window_size distance.

Arguments

• q - Query tensor [batch, seq_q, dim]
• k - Key tensor [batch, seq_k, dim]
• v - Value tensor [batch, seq_k, dim]

Returns

Output tensor [batch, seq_q, dim]

pub fn memory_savings(&self, seq_len: usize) -> String

Get memory savings compared to full attention

Trait Implementations

impl Debug for SlidingWindowAttention

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.