Struct CausalSoftmaxInplace 

pub struct CausalSoftmaxInplace {
    pub offset: usize,
}

The same causal softmax, written in place.

Why this exists beside [CausalSoftmax]

The score tensor is (1, 9, 1142, 1142) — 47 MB. apply_op1_no_bwd allocates a second one and writes the result there, so every layer allocates 47 MB, writes 47 MB, and drops 47 MB. Over 30 layers that is 1.4 GB of allocation churn and an extra 1.4 GB of writes, and it was the largest part of the 18 % of prefill that per-op timing could not account for (examples/text_inline_prof).

The scores come straight out of q.matmul(k.t()) and nothing else holds a reference, so the buffer is ours to overwrite. Softmax is row-local — each row’s max, sum and normalisation touch only that row — so writing over the input as we go is safe and produces identical values.

Fields

offset: usize

Where this query block starts in the full sequence.

Trait Implementations

impl InplaceOp1 for CausalSoftmaxInplace

fn name(&self) -> &'static str

fn cpu_fwd(&self, storage: &mut CpuStorage, layout: &Layout) -> CandleResult<()>

The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides, offsets etc so the associated layout should be used to access it.

fn cuda_fwd( &self, _storage: &mut CudaStorage, _layout: &Layout, ) -> Result<(), Error>

The forward pass, as run on a gpu device. Note that the storage can use arbitrary strides, offsets etc so the associated layout should be used to access it.

fn metal_fwd( &self, _storage: &mut MetalStorage, _layout: &Layout, ) -> Result<(), Error>

The forward pass, as run on a metal gpu device. Note that the storage can use arbitrary strides, offsets etc so the associated layout should be used to access it.