Struct CrossAttention 

pub struct CrossAttention { /* private fields */ }

Cross-attention module with optional flash attention support.

When flash attention is enabled (via feature flag and configuration), uses memory-efficient O(N) block-wise attention computation instead of the standard O(N^2) attention matrix.

Implementations

impl CrossAttention

pub fn new( vs: VarBuilder<'_>, query_dim: usize, context_dim: Option<usize>, heads: usize, dim_head: usize, ) -> Result<Self>

Create a new cross-attention module with standard attention.

pub fn new_with_flash( vs: VarBuilder<'_>, query_dim: usize, context_dim: Option<usize>, heads: usize, dim_head: usize, use_flash_attention: bool, flash_block_size: usize, ) -> Result<Self>

Create a new cross-attention module with optional flash attention.

Arguments

• vs - Variable builder for weight initialization
• query_dim - Query input dimension
• context_dim - Context dimension (None for self-attention)
• heads - Number of attention heads
• dim_head - Dimension per head
• use_flash_attention - Whether to use flash attention
• flash_block_size - Block size for flash attention tiling

pub fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result<Tensor>

Scaled-dot-product attention (standard or flash based on configuration).

Automatically dispatches to flash attention when enabled and the feature is available, otherwise uses standard O(N^2) attention.

pub fn is_flash_attention_enabled(&self) -> bool

Check if flash attention is enabled for this module.

Returns true only if flash attention was requested during construction AND the flash_attention feature is enabled.

Trait Implementations

impl Debug for CrossAttention

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

Formats the value using the given formatter.