Struct BounceTile 

pub struct BounceTile<N: Numeric> {
    pub cmma: CmmaTile<N>,
    pub smem: Shared<[N]>,
    pub fragment: WhiteboxFragment<N>,
}

Fields

cmma: CmmaTile<N>

smem: Shared<[N]>

fragment: WhiteboxFragment<N>

Implementations

impl<N: Numeric> BounceTile<N>

pub fn new(cmma: CmmaTile<N>, cfg: BounceConfig) -> BounceTile<N>

pub fn __expand_new( scope: &Scope, cmma: <CmmaTile<N> as CubeType>::ExpandType, cfg: BounceConfig, ) -> <BounceTile<N> as CubeType>::ExpandType

impl<E: Float> BounceTile<E>

pub fn cmma_to_fragment(&mut self)

Synchronizes the fragment view from the cmma fragment via smem. Call before any rowwise/elementwise op so the fragment reflects the current cmma state.

pub fn fragment_to_cmma(&mut self)

Synchronizes the cmma fragment from the fragment view via smem. Call after rowwise/elementwise edits to make the cmma side current for the next mma.

pub fn row_max(&self, acc: &mut RowWise<E>, base: &RowWise<E>)

pub fn row_sum(&self, acc: &mut RowWise<E>)

pub fn exp_diff(&mut self, rowwise: &RowWise<E>)

pub fn rowwise_scale(&mut self, scale: &RowWise<E>)

pub fn scale_and_mask<M: Mask>(&mut self, scale: E, mask: &M)

pub fn fill_zero(&mut self)

Zeros the cmma fragment. The fragment view is not the live storage at fill_zero call sites (always invoked before any cmma_to_fragment), so only cmma needs clearing.

pub fn write_fragment_to<Lhs: Float, Sc: TileScope>( &self, softmaxed: &mut Tile<Lhs, Sc>, )

Writes the (already-softmaxed) fragment view of this bounce tile into softmaxed. The source fragment is plane-fragmented; for a Bounce destination this routes through the destination’s smem into its cmma fragment.

pub fn __expand_cmma_to_fragment( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, )

Synchronizes the fragment view from the cmma fragment via smem. Call before any rowwise/elementwise op so the fragment reflects the current cmma state.

pub fn __expand_fragment_to_cmma( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, )

Synchronizes the cmma fragment from the fragment view via smem. Call after rowwise/elementwise edits to make the cmma side current for the next mma.

pub fn __expand_row_max( scope: &Scope, this: &<Self as CubeType>::ExpandType, acc: &mut <RowWise<E> as CubeType>::ExpandType, base: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_row_sum( scope: &Scope, this: &<Self as CubeType>::ExpandType, acc: &mut <RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_exp_diff( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, rowwise: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_rowwise_scale( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, scale: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_scale_and_mask<M: Mask>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, scale: <E as CubeType>::ExpandType, mask: &<M as CubeType>::ExpandType, )

pub fn __expand_fill_zero( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, )

Zeros the cmma fragment. The fragment view is not the live storage at fill_zero call sites (always invoked before any cmma_to_fragment), so only cmma needs clearing.

pub fn __expand_write_fragment_to<Lhs: Float, Sc: TileScope>( scope: &Scope, this: &<Self as CubeType>::ExpandType, softmaxed: &mut <Tile<Lhs, Sc> as CubeType>::ExpandType, )

Writes the (already-softmaxed) fragment view of this bounce tile into softmaxed. The source fragment is plane-fragmented; for a Bounce destination this routes through the destination’s smem into its cmma fragment.

impl<N: Numeric> BounceTile<N>

pub fn copy_from<SE: Numeric, SS: Size, Sc: TileScope>( &mut self, source: &Tile<SE, Sc>, ident: StageIdent, )

Copies into the bounce tile’s cmma fragment from source. Bounce always loads through its CMMA representation (the WhiteboxFragment view is synced lazily on demand by softmax/scale ops); supported sources mirror CmmaTile::copy_from.

pub fn init_zero(&mut self)

Zero-init the bounce tile (clears its cmma fragment).

pub fn __expand_copy_from<SE: Numeric, SS: Size, Sc: TileScope>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, source: &<Tile<SE, Sc> as CubeType>::ExpandType, ident: StageIdent, )

Copies into the bounce tile’s cmma fragment from source. Bounce always loads through its CMMA representation (the WhiteboxFragment view is synced lazily on demand by softmax/scale ops); supported sources mirror CmmaTile::copy_from.

pub fn __expand_init_zero( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, )

Zero-init the bounce tile (clears its cmma fragment).

impl<Acc: Float> BounceTile<Acc>

pub fn softmax<Lhs: Float, M: Mask>( &mut self, mask: &M, softmaxed: &mut Tile<Lhs, Plane>, state: &mut (RowWise<Acc>, RowWise<Acc>), head_dim_factor: Acc, ) -> RowWise<Acc>

Online softmax for the Bounce variant. cmma → fragment once at entry so all subsequent rowwise ops read/write the fragment view; the post- exp values are still in the fragment at the end (we skip fragment_to_cmma on score because its cmma is cleared next iteration), and we stream straight into softmaxed via write_fragment_to.

pub fn __expand_softmax<Lhs: Float, M: Mask>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, mask: &<M as CubeType>::ExpandType, softmaxed: &mut <Tile<Lhs, Plane> as CubeType>::ExpandType, state: &mut (<RowWise<Acc> as CubeType>::ExpandType, <RowWise<Acc> as CubeType>::ExpandType), head_dim_factor: <Acc as CubeType>::ExpandType, ) -> <RowWise<Acc> as CubeType>::ExpandType

Online softmax for the Bounce variant. cmma → fragment once at entry so all subsequent rowwise ops read/write the fragment view; the post- exp values are still in the fragment at the end (we skip fragment_to_cmma on score because its cmma is cleared next iteration), and we stream straight into softmaxed via write_fragment_to.

Trait Implementations

impl<N: Numeric> CubeType for BounceTile<N>

type ExpandType = BounceTileExpand<N>