Struct MultiHeadSelfAttention 

pub struct MultiHeadSelfAttention<B: Backend> { /* private fields */ }

Scaled dot-product multi-head self-attention with optional chunked computation.

When chunk_size > 0 the query sequence is processed in windows of chunk_size rows, keeping the forward-pass peak attention memory at O(B · H · chunk_size · N) instead of O(B · H · N²), and ensuring each individual WGPU GPU dispatch remains small enough to avoid OS watchdog (TDR) timeouts.

⚠ Training memory — chunking reduces dispatch size but NOT total tape

Burn’s forward pass builds an autodiff tape for every transformer layer before loss.backward() runs. At the forward→backward boundary all depth layers’ chunk tensors are simultaneously in GPU memory:

peak = depth × 2 × ceil(N/chunk) × B × H × chunk × N × 4 bytes
     = 12 × 2 × 39 × B × 12 × 64 × 2448 × 4   (ViT-B defaults)
     ≈ 6.56 GB × B

Chunking (small chunk_size) keeps individual GPU dispatch sizes small (preventing OS watchdog / TDR timeouts), but the cumulative tape size is the same as full attention. The only way to reduce training memory is gradient checkpointing (recompute attention during backward instead of storing it) — not yet implemented in this codebase.

Safe configurations (24 GB GPU, ViT-B):

• batch_size = 2 → all-layers peak ≈ 13 GB ✓
• batch_size = 4 → all-layers peak ≈ 26 GB ✗ OOM

The crate::training::learner::train function guards against unsafe configurations using --vram-gb to derive the correct limit.

Forward memory comparison (N = 2 448, H = 12, B = 8, fp32)

mode	peak fwd attn tensor	size
full (chunk=0)	(8, 12, 2448, 2448)	~18 GB
chunk=256	(8, 12, 256, 2448)	~1.9 GB
chunk=128	(8, 12, 128, 2448)	~960 MB
chunk=64	(8, 12, 64, 2448)	~480 MB

Implementations

impl<B: Backend> MultiHeadSelfAttention<B>

pub fn new( d_model: usize, num_heads: usize, dropout: f64, chunk_size: usize, device: &B::Device, ) -> Self

Construct MHSA.

• chunk_size – query chunk window; 0 disables chunking.

pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3>

Self-attention: (B, N, D) → (B, N, D).

When chunk_size > 0 the computation is split into ceil(N / chunk_size) passes, each allocating an attention matrix of shape (B, H, chunk_size, N) rather than (B, H, N, N).

Trait Implementations

impl<B> AutodiffModule<B> for MultiHeadSelfAttention<B>

where B: AutodiffBackend + Backend, <B as AutodiffBackend>::InnerBackend: Backend,

type InnerModule = MultiHeadSelfAttention<<B as AutodiffBackend>::InnerBackend>

Inner module without auto-differentiation.

fn valid(&self) -> Self::InnerModule

Get the same module, but on the inner backend without auto-differentiation.

impl<B: Backend> Clone for MultiHeadSelfAttention<B>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<B: Debug + Backend> Debug for MultiHeadSelfAttention<B>

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

Formats the value using the given formatter.

impl<B: Backend> Display for MultiHeadSelfAttention<B>

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

Formats the value using the given formatter.

impl<B: Backend> Module<B> for MultiHeadSelfAttention<B>

type Record = MultiHeadSelfAttentionRecord<B>

Type to save and load the module.

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

fn visit<Visitor: ModuleVisitor<B>>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.

fn map<Mapper: ModuleMapper<B>>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.

fn collect_devices(&self, devices: Devices<B>) -> Devices<B>

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn to_device(self, device: &B::Device) -> Self

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &B::Device) -> Self

Fork the module and all of its sub-modules to the given device.

fn devices(&self) -> Vec<<B as Backend>::Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

fn save_file<FR, PB>( self, file_path: PB, recorder: &FR, ) -> Result<(), RecorderError>

where FR: FileRecorder<B>, PB: Into<PathBuf>,

Save the module to a file using the provided file recorder.

fn load_file<FR, PB>( self, file_path: PB, recorder: &FR, device: &<B as Backend>::Device, ) -> Result<Self, RecorderError>

where FR: FileRecorder<B>, PB: Into<PathBuf>,

Load the module from a file using the provided file recorder.

fn quantize_weights<C>(self, quantizer: &mut Quantizer<C>) -> Self

where C: Calibration,

Quantize the weights of the module.

impl<B: Backend> ModuleDisplay for MultiHeadSelfAttention<B>

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl<B: Backend> ModuleDisplayDefault for MultiHeadSelfAttention<B>

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.